repro-bundle-v0/src/objective.py

"""Synthetic nonconvex stochastic objective used to verify the Lion convergence
theorems (Jiang & Zhang, arXiv:2508.12327).

f(x) = (1/N) sum_i phi(a_i^T x - b_i),   phi(z) = z^2 / (1 + z^2)

This is the standard smooth-but-nonconvex "correntropy" / robust-regression
loss: L-smooth (bounded Hessian), bounded below (0 <= f), and its gradient is
uniformly bounded (needed for the paper's Assumption 9, used by the
communication-efficient variants). Both properties are checked in
src/sanity_checks.py.

phi'(z)  = 2z / (1+z^2)^2
grad f(x) = (1/N) sum_i phi'(a_i^T x - b_i) * a_i
"""
import numpy as np


def phi(z):
    return z * z / (1.0 + z * z)


def phi_prime(z):
    denom = (1.0 + z * z) ** 2
    return 2.0 * z / denom


def make_dataset(N, d, seed, x_star=None, feature_scale=1.0, label_noise=0.0):
    """i.i.d. Gaussian design matrix + labels generated from a "true" x_star,
    with optional label noise. Returns (A, b, x_star)."""
    rng = np.random.default_rng(seed)
    A = rng.normal(scale=feature_scale, size=(N, d)) / np.sqrt(d)
    if x_star is None:
        x_star = rng.normal(size=d)
    b = A @ x_star
    if label_noise > 0:
        b = b + rng.normal(scale=label_noise, size=N)
    return A, b, x_star


def full_grad(x, A, b):
    r = A @ x - b
    coeff = phi_prime(r)
    return (A * coeff[:, None]).mean(axis=0)


def full_value(x, A, b):
    r = A @ x - b
    return phi(r).mean()


def stochastic_grad(x, A, b, batch_size, rng):
    """Minibatch stochastic gradient: uniform sampling without replacement."""
    N = A.shape[0]
    idx = rng.integers(0, N, size=batch_size)  # with replacement -> i.i.d. noise, matches Assumption 3
    Ab = A[idx]
    bb = b[idx]
    r = Ab @ x - bb
    coeff = phi_prime(r)
    return (Ab * coeff[:, None]).mean(axis=0)


def make_heterogeneous_nodes(n, N_per_node, d, seed, heterogeneity=0.5, feature_scale=1.0):
    """n nodes with DIFFERENT local minimizers x_star_j = x_star_global +
    heterogeneity * N(0, I) -- satisfies the paper's heterogeneous setting
    (f_j allowed to differ significantly across nodes) while keeping the
    global average f = (1/n) sum f_j well-defined and smooth.

    Returns stacked arrays A (n, N_per_node, d) and b (n, N_per_node) so the
    per-node stochastic-gradient computation in src/lion.py can be fully
    vectorized across nodes (no Python-level loop over n)."""
    rng = np.random.default_rng(seed)
    x_star_global = rng.normal(size=d)
    A = np.empty((n, N_per_node, d))
    b = np.empty((n, N_per_node))
    for j in range(n):
        x_star_j = x_star_global + heterogeneity * rng.normal(size=d)
        A_j, b_j, _ = make_dataset(N_per_node, d, seed=seed * 100003 + j,
                                    x_star=x_star_j, feature_scale=feature_scale)
        A[j], b[j] = A_j, b_j
    return A, b


def full_grad_nodes(x, A, b):
    """Vectorized per-node full gradient. A: (n,N,d), b: (n,N). Returns (n,d)."""
    r = np.einsum('njd,d->nj', A, x) - b
    coeff = phi_prime(r)
    return np.einsum('nj,njd->nd', coeff, A) / A.shape[1]


def stochastic_grad_nodes(x, A, b, idx):
    """Vectorized per-node minibatch stochastic gradient.
    A: (n,N,d), b: (n,N), idx: (n,batch) per-node sample indices (int).
    Returns (n,d)."""
    n = A.shape[0]
    rows = np.arange(n)[:, None]
    Ab = A[rows, idx]          # (n, batch, d)
    bb = b[rows, idx]          # (n, batch)
    r = np.einsum('nbd,d->nb', Ab, x) - bb
    coeff = phi_prime(r)
    return np.einsum('nb,nbd->nd', coeff, Ab) / idx.shape[1]


def max_grad_coeff_bound():
    """max_z |phi'(z)| -- loose analytical bound on the per-SAMPLE gradient
    contribution phi'(r_i), not on a gradient (mean-of-N, feature-scaled)
    component -- see empirical_grad_bound for the bound that actually
    matters for Assumption 9 / the S_G(.) compression radius."""
    # phi'(z) = 2z/(1+z^2)^2, maximized at z = 1/sqrt(3): value = 3*sqrt(3)/8
    return 3.0 * np.sqrt(3.0) / 8.0


def empirical_grad_bound(A, b, x_radius, n_samples=150, seed=0, margin=1.5):
    """sup_x ||grad f_j(x)||_inf, estimated by sampling x over a ball of the
    radius actually traversed during optimization (x starts at 0 and moves
    toward x_star, so x_radius ~ typical ||x_star||), with a safety margin.
    A, b are the stacked per-node arrays (n,N,d) / (n,N); a single (N,d)/(N,)
    pair also works (treated as one "node")."""
    rng = np.random.default_rng(seed)
    stacked = A.ndim == 3
    d = A.shape[-1]
    worst = 0.0
    for _ in range(n_samples):
        x = rng.normal(size=d) * rng.uniform(0, x_radius)
        if stacked:
            g = full_grad_nodes(x, A, b)
        else:
            g = full_grad(x, A, b)
        worst = max(worst, np.abs(g).max())
    return margin * worst