repro-bundle-v0/src/problems.py

"""Synthetic distributed problems used to numerically audit FAB's theory.

Bilevel problem (Section 2, Eq. 1): a decentralized ridge-regression
hyperparameter-tuning task in the style of the paper's own "distributed
hyperparameter tuning" experiment (Sec 4, Appendix B). Each agent i holds
local training data (A_i, b_i) and validation data (C_i, e_i):

    g_i(x, y) = (1/2 m_i) ||A_i y - b_i||^2 + (mu/2) ||y - x||^2      (lower level, mu-strongly convex in y)
    f_i(x, y) = (1/2 p_i) ||C_i y - e_i||^2 + kappa_i (1 - cos(1^T x))  (upper level, smooth but NONconvex in x)

x is a learned "prior/anchor" hyperparameter that the lower-level ridge
regression is regularized towards; y is the regression weight vector.
Because g_i is jointly quadratic, y*(x) is affine in x, which lets us
compute the true hypergradient of F*(x) := (1/n) sum_i f_i(x, y*(x)) in
closed form (via the implicit function theorem) as an exact ground truth --
no need for a separate high-precision centralized solve.

Single-level problem (Eq. 9, used for Theorem 3.6 / Push-Pull only): the same
kind of smooth nonconvex coupling term directly on x, no bilevel structure.
"""
from __future__ import annotations

import numpy as np


class BilevelHyperTuning:
    def __init__(self, n: int, d: int, m: int, p: int, mu: float = 1.0,
                 heterogeneity: float = 0.5, seed: int = 0):
        rng = np.random.default_rng(seed)
        self.n, self.d, self.m, self.p, self.mu = n, d, m, p, mu

        y_true = rng.normal(size=d) / np.sqrt(d)
        self.A = np.empty((n, m, d))
        self.b = np.empty((n, m))
        self.C = np.empty((n, p, d))
        self.e = np.empty((n, p))
        self.kappa = rng.uniform(0.5, 1.5, size=n)

        for i in range(n):
            # heterogeneity in (0, 1]: smaller => agents see more different local targets
            y_i = y_true + heterogeneity * rng.normal(size=d) / np.sqrt(d)
            self.A[i] = rng.normal(size=(m, d)) / np.sqrt(m)
            self.b[i] = self.A[i] @ y_i + 0.01 * rng.normal(size=m)
            self.C[i] = rng.normal(size=(p, d)) / np.sqrt(p)
            self.e[i] = self.C[i] @ y_true + 0.01 * rng.normal(size=p)

        # M = (1/n) sum_i (1/m) A_i^T A_i + mu I ; y*(x) = M^{-1}(v + mu x)
        M = mu * np.eye(d)
        v = np.zeros(d)
        for i in range(n):
            M += (self.A[i].T @ self.A[i]) / (n * m)
            v += (self.A[i].T @ self.b[i]) / (n * m)
        self.M = M
        self.v = v
        self.Minv = np.linalg.inv(M)

    def y_star(self, x: np.ndarray) -> np.ndarray:
        return self.Minv @ (self.v + self.mu * x)

    def f_i(self, i: int, x, y):
        r = self.C[i] @ y - self.e[i]
        return 0.5 * (r @ r) / self.p + self.kappa[i] * (1 - np.cos(x.sum()))

    def g_i(self, i: int, x, y):
        r = self.A[i] @ y - self.b[i]
        return 0.5 * (r @ r) / self.m + 0.5 * self.mu * np.sum((y - x) ** 2)

    def grad_y_f_i(self, i, x, y):
        return (self.C[i].T @ (self.C[i] @ y - self.e[i])) / self.p

    def grad_x_f_i(self, i, x, y):
        return np.full_like(x, self.kappa[i] * np.sin(x.sum()))

    def grad_y_g_i(self, i, x, y):
        return (self.A[i].T @ (self.A[i] @ y - self.b[i])) / self.m + self.mu * (y - x)

    def grad_x_g_i(self, i, x, y):
        return self.mu * (x - y)

    def true_hypergradient(self, x: np.ndarray) -> np.ndarray:
        """Exact grad F*(x) via the implicit function theorem: dy*/dx = mu * M^{-1}."""
        y = self.y_star(x)
        grad_y_f_avg = np.mean([self.grad_y_f_i(i, x, y) for i in range(self.n)], axis=0)
        grad_x_f_avg = np.mean([self.grad_x_f_i(i, x, y) for i in range(self.n)], axis=0)
        dy_dx = self.mu * self.Minv  # d x d
        return dy_dx.T @ grad_y_f_avg + grad_x_f_avg

    def F_star(self, x: np.ndarray) -> float:
        y = self.y_star(x)
        return float(np.mean([self.f_i(i, x, y) for i in range(self.n)]))


class SingleLevelNonconvex:
    """min_x (1/n) sum_i f_i(x), f_i(x) = kappa_i (1 - cos(1^T x)) + 0.5*c_i ||x - r_i||^2.

    Smooth, bounded-below, genuinely nonconvex (periodic term), with per-agent
    heterogeneity (r_i, c_i, kappa_i) -- used only for Theorem 3.6 / Prop 3.5's
    single-level Push-Pull special case (Eq. 9), no bilevel structure at all.
    """

    def __init__(self, n: int, d: int, heterogeneity: float = 0.5, seed: int = 1):
        rng = np.random.default_rng(seed)
        self.n, self.d = n, d
        self.kappa = rng.uniform(0.5, 1.5, size=n)
        self.c = rng.uniform(0.5, 1.5, size=n)
        self.r = heterogeneity * rng.normal(size=(n, d)) / np.sqrt(d)

    def f_i(self, i, x):
        return self.kappa[i] * (1 - np.cos(x.sum())) + 0.5 * self.c[i] * np.sum((x - self.r[i]) ** 2)

    def grad_i(self, i, x):
        return np.full_like(x, self.kappa[i] * np.sin(x.sum())) + self.c[i] * (x - self.r[i])

    def F(self, x):
        return float(np.mean([self.f_i(i, x) for i in range(self.n)]))

    def true_gradient(self, x):
        return np.mean([self.grad_i(i, x) for i in range(self.n)], axis=0)