repro_bundle/code/lsr1.py

"""L-SR1: Learned Symmetric-Rank-One Preconditioning.

Independent PyTorch reimplementation of the analytic (learned-optimizer) core of

    Lifshitz, Zuler, Fouks, Raviv. "L-SR1: Learned Symmetric-Rank-One
    Preconditioning." ICML 2026 (OpenReview w1fkDwiZgN, arXiv:2508.12270).

No official code exists ("Code will be made publicly available upon acceptance"),
so everything here is reconstructed from the paper text, Algorithm 1, Eq. 8/9,
and Tables 4-6.

The optimizer has three *element-wise* MLP modules (Claim 1) that operate on the
per-coordinate optimizer state, a limited-memory SR1 preconditioner built from
generated vectors, and a learned PSD projection enforced through a secant-residual
penalty in the meta-loss (Claim 2).
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Callable

import torch
import torch.nn as nn


# ---------------------------------------------------------------------------
# Learned modules (Tables 4 & 5) -- all applied element-wise over coordinates.
# ---------------------------------------------------------------------------
class BasicBlock(nn.Module):
    """Table 5: Linear->BN->PReLU->Dropout->Linear->BN->Dropout, residual add."""

    def __init__(self, dim: int, dropout: float = 0.0):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim)
        self.bn1 = nn.BatchNorm1d(dim)
        self.act = nn.PReLU()
        self.drop1 = nn.Dropout(dropout)
        self.fc2 = nn.Linear(dim, dim)
        self.bn2 = nn.BatchNorm1d(dim)
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x):
        h = self.drop1(self.act(self.bn1(self.fc1(x))))
        h = self.drop2(self.bn2(self.fc2(h)))
        return x + h


class MLP(nn.Module):
    """Table 4: fc1->BN->PReLU->Dropout->BasicBlock x2->fc2. Element-wise.

    Input is a tensor whose last dim is ``d_in``; every leading dim is treated as
    an independent sample (so the same MLP acts on every problem coordinate,
    giving the dimension-invariance the paper relies on).
    """

    def __init__(self, d_in: int, d_out: int, d_hidden: int = 128,
                 n_blocks: int = 2, dropout: float = 0.0):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.bn1 = nn.BatchNorm1d(d_hidden)
        self.act = nn.PReLU()
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList(
            [BasicBlock(d_hidden, dropout) for _ in range(n_blocks)])
        self.fc2 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        lead = x.shape[:-1]
        x = x.reshape(-1, x.shape[-1])
        h = self.drop(self.act(self.bn1(self.fc1(x))))
        for blk in self.blocks:
            h = blk(h)
        h = self.fc2(h)
        return h.reshape(*lead, h.shape[-1])


@dataclass
class LSR1Config:
    d_hidden: int = 128
    n_encoder_inputs: int = 5      # x, p, d_prev, g, q
    latent: int = 128              # encoder output width M (>5)
    buffer_L: int = 8              # limited-memory buffer size
    K: int = 16                    # unrolled optimization steps (meta-training)
    gamma1: float = 0.4            # alpha = gamma1 * exp(gamma2 * alpha_tilde)
    gamma2: float = 0.001
    lambda_sec: float = 100.0      # secant-residual weight in meta-loss
    dropout: float = 0.0
    step_includes_identity: bool = True   # B0=I term in the *step* direction:
    # the SR1 preconditioner is B~=I+Sum vv^T (matching the secant penalty), so
    # even the untrained optimizer does gradient descent and the learned vectors
    # add curvature. Verified essential: with it, the learned step reaches cosine
    # ~1.0 with the Newton direction on N=2 quadratics; without it, the optimizer
    # must learn the whole preconditioner from scratch and fails to align.
    relative_secant: bool = False         # divide secant residual by ||q||^2
    trust_region: float | None = None     # cap on per-problem step L2 norm.
    # NOTE: not described in the paper. The learned SR1 direction d=(Sum vv^T)g is
    # linear in the gradient, so on stiff problems (Rosenbrock) an untrained/early
    # optimizer produces runaway steps (x->g->step feedback) and meta-training
    # NaNs. A trust region is standard practice for learned optimizers (cf.
    # gradient clipping in Andrychowicz+16). It is set large enough that it does
    # not bind near convergence, so quadratic/Newton-alignment results are
    # unaffected; it only tames the early nonconvex-training transient.


class LSR1Optimizer(nn.Module):
    """The three learned modules + the limited-memory SR1 rollout (Algorithm 1)."""

    def __init__(self, cfg: LSR1Config):
        super().__init__()
        self.cfg = cfg
        self.encoder = MLP(cfg.n_encoder_inputs, cfg.latent, cfg.d_hidden,
                           dropout=cfg.dropout)
        self.vector_gen = MLP(cfg.latent, 1, cfg.d_hidden, dropout=cfg.dropout)
        self.lr_gen = MLP(cfg.latent, 1, cfg.d_hidden, dropout=cfg.dropout)
        # Small-init the generator heads so the *untrained* preconditioner is
        # near-zero: this keeps the initial K-step rollout numerically stable
        # (large random v would make d = Sum v(v.g) overshoot and diverge) while
        # leaving non-zero gradients so the modules can learn to grow. Standard
        # practice for meta-trained (learned) optimizers.
        with torch.no_grad():
            self.vector_gen.fc2.weight.mul_(0.01)
            self.vector_gen.fc2.bias.zero_()
            self.lr_gen.fc2.weight.zero_()      # alpha_tilde=0 -> alpha=gamma1
            self.lr_gen.fc2.bias.zero_()

    # -- one learned step -------------------------------------------------
    def step_direction(self, x, p, d_prev, g, q, buffer):
        """Return (v_k, alpha_k, d_k, latent) for the current state.

        All tensors are [B, N] (batch of problems x coordinates); the buffer is a
        list of past v vectors, each [B, N]. ``d_k`` is the SR1 descent direction
        Sum_{v in B_L} v (v . g)  (Algorithm 1), optionally + g (B0=I term).
        """
        feats = torch.stack([x, p, d_prev, g, q], dim=-1)          # [B, N, 5]
        latent = self.encoder(feats)                               # [B, N, M]
        v = self.vector_gen(latent).squeeze(-1)                    # [B, N]
        alpha_tilde = self.lr_gen(latent).squeeze(-1)              # [B, N]
        alpha = self.cfg.gamma1 * torch.exp(self.cfg.gamma2 * alpha_tilde)

        vecs = buffer + [v]
        vecs = vecs[-self.cfg.buffer_L:]
        d = torch.zeros_like(g)
        if self.cfg.step_includes_identity:
            d = d + g
        for vv in vecs:
            coeff = (vv * g).sum(dim=-1, keepdim=True)             # [B, 1]
            d = d + coeff * vv
        return v, alpha, d, vecs

    @staticmethod
    def secant_residual(q, p, buffer, relative=False):
        """||p - B~ q||^2 with B~ = I + Sum v v^T  (Eq. 8), per problem.

        ``relative`` divides by ||q||^2 so the penalty is scale-invariant across
        problem families (needed only for the mixed benchmark, where gradient
        magnitudes span orders of magnitude; the quadratic claim-2 study uses the
        absolute form exactly as written in the paper).
        """
        Bq = q.clone()
        for vv in buffer:
            Bq = Bq + (vv * q).sum(dim=-1, keepdim=True) * vv
        res = ((p - Bq) ** 2).sum(dim=-1)                         # [B]
        if relative:
            res = res / ((q ** 2).sum(dim=-1) + 1e-8)
        return res


def rollout(opt: LSR1Optimizer, problem, x0, K=None, create_graph=True,
            record=False):
    """Unroll the learned optimizer for K steps on a batch of problems.

    ``problem`` exposes ``value(x) -> [B]`` and (via autograd) its gradient.
    Returns a dict with the accumulated meta-loss terms and, if ``record``, the
    trajectory / per-step diagnostics (used at eval time).
    """
    cfg = opt.cfg
    K = cfg.K if K is None else K
    x = x0.clone().requires_grad_(True)
    B, N = x.shape
    zeros = torch.zeros_like(x)
    p = zeros; d_prev = zeros; q = zeros
    g = _grad(problem.value(x).sum(), x, create_graph)
    g_prev = g
    buffer: list[torch.Tensor] = []

    f_terms, sec_terms = [], []
    traj = [x.detach().clone()] if record else None
    dirs = [] if record else None

    for k in range(K):
        v, alpha, d, buffer = opt.step_direction(x, p, d_prev, g, q, buffer)
        step = alpha * d
        if cfg.trust_region is not None:
            # per-coordinate (infinity-norm) trust region -> dimension-invariant
            step = step.clamp(-cfg.trust_region, cfg.trust_region)
        x_next = x - step
        g_next = _grad(problem.value(x_next).sum(), x_next, create_graph)
        p_next = x_next - x
        q_next = g_next - g
        sec = opt.secant_residual(q_next, p_next, buffer,
                                  relative=cfg.relative_secant)

        f_terms.append(problem.value(x_next))
        sec_terms.append(sec)
        if record:
            traj.append(x_next.detach().clone())
            dirs.append((-alpha * d).detach().clone())

        x, g_prev, g = x_next, g, g_next
        p, q, d_prev = p_next, q_next, d

    f_stack = torch.stack(f_terms)            # [K, B]
    sec_stack = torch.stack(sec_terms)        # [K, B]
    meta = (f_stack + cfg.lambda_sec * sec_stack).mean()
    out = {
        "meta_loss": meta,
        "f_final": f_stack[-1].detach(),
        "f_traj": f_stack.detach(),
        "sec_traj": sec_stack.detach(),
    }
    if record:
        out["traj"] = torch.stack(traj)       # [K+1, B, N]
        out["dirs"] = torch.stack(dirs)       # [K, B, N]
    return out


def _grad(scalar, x, create_graph):
    (g,) = torch.autograd.grad(scalar, x, create_graph=create_graph)
    return g


# ---------------------------------------------------------------------------
# Analytic test problems (Sec. 5.1, App. C.2)
# ---------------------------------------------------------------------------
class Quadratic:
    """f(x) = 1/2 x^T H x + b^T x, H PSD (cond capped, unit Frobenius), |b|=1."""

    def __init__(self, H, b):
        self.H = H                        # [B, N, N]
        self.b = b                        # [B, N]
        self.name = "quadratic"

    def value(self, x):                   # x: [B, N] -> [B]
        Hx = torch.einsum("bij,bj->bi", self.H, x)
        return 0.5 * (x * Hx).sum(-1) + (self.b * x).sum(-1)

    def newton_dir(self, x):
        g = torch.einsum("bij,bj->bi", self.H, x) + self.b
        return -torch.linalg.solve(self.H, g)

    def optimum(self):
        return -torch.linalg.solve(self.H, self.b)


def make_quadratics(B, N, cond=1000.0, unit_frobenius=True, diagonal=False,
                    generator=None, device="cpu", dtype=torch.float32):
    g = generator
    if diagonal:
        # eigenvalues geometrically spaced in [1/cond, 1] -> condition number cond
        if cond == 1.0:
            eig = torch.ones(B, N, device=device, dtype=dtype)
        else:
            t = torch.linspace(0, 1, N, device=device, dtype=dtype)
            eig = (1.0 / cond) ** (1 - t)          # from 1/cond up to 1
            eig = eig.unsqueeze(0).expand(B, N).contiguous()
        H = torch.diag_embed(eig)
    else:
        A = torch.randn(B, N, N, device=device, dtype=dtype, generator=g)
        H = A.transpose(-1, -2) @ A / N
        # cap condition number at `cond`
        evals, evecs = torch.linalg.eigh(H)
        lo = evals[..., -1:] / cond
        evals = torch.clamp(evals, min=lo)
        H = evecs @ torch.diag_embed(evals) @ evecs.transpose(-1, -2)
    if unit_frobenius:
        fro = torch.linalg.matrix_norm(H, ord="fro").unsqueeze(-1).unsqueeze(-1)
        H = H / fro
    b = torch.randn(B, N, device=device, dtype=dtype, generator=g)
    b = b / b.norm(dim=-1, keepdim=True)
    return Quadratic(H, b)


class Rosenbrock:
    name = "rosenbrock"

    def value(self, x):
        xi, xn = x[..., :-1], x[..., 1:]
        return (100.0 * (xn - xi ** 2) ** 2 + (1 - xi) ** 2).sum(-1)


class Rastrigin:
    name = "rastrigin"

    def value(self, x):
        N = x.shape[-1]
        return 10 * N + (x ** 2 - 10 * torch.cos(2 * math.pi * x)).sum(-1)


def init_points(B, N, problem_name, scale=1.0, generator=None, device="cpu",
                dtype=torch.float32):
    x = torch.randn(B, N, device=device, dtype=dtype, generator=generator)
    if problem_name == "rosenbrock":
        x = x * 0.5                       # start near the valley but off-optimum
    elif problem_name == "rastrigin":
        x = x * 2.0                       # multimodal basin
    else:
        x = x * scale
    return x