1"""Projection-Regularized Gradient optimizer (small-model exact-matrix MVP).
 2
 3This implementation follows the supplied formula literally. It is intended for
 4small models; for large models, replace _metric with a diagonal or Woodbury
 5implementation rather than materializing p x p matrices.
 6"""
 7from collections import deque
 8import torch
 9
10
11class ProjectionRegularizedSGD(torch.optim.Optimizer):
12    """SGD using a regularized recent-gradient subspace metric.
13
14    Args:
15        params: model parameters
16        lr: step size
17        history_size: number of recent flattened gradients
18        beta: covariance EMA coefficient (the MVP uses the current window)
19        ridge: lambda in S + lambda I
20        rho: ridge in the m-by-m projected Gram system
21        alpha: residual identity strength, in [0, 1]
22        max_exact_parameters: safety limit for the exact p-by-p implementation
23    """
24    def __init__(self, params, lr=1e-2, history_size=4, beta=0.0,
25                 ridge=1e-2, rho=0.1, alpha=0.1, max_exact_parameters=4096):
26        if lr < 0 or history_size < 1 or not (0 <= beta < 1):
27            raise ValueError("invalid lr/history_size/beta")
28        if ridge <= 0 or rho <= 0 or not (0 <= alpha <= 1):
29            raise ValueError("require ridge,rho > 0 and alpha in [0,1]")
30        super().__init__(params, dict(lr=lr))
31        self.history_size, self.beta = history_size, beta
32        self.ridge, self.rho, self.alpha = ridge, rho, alpha
33        self.max_exact_parameters = max_exact_parameters
34        self._history = deque(maxlen=history_size)
35        self._S = None
36
37    @torch.no_grad()
38    def _flat_grad(self):
39        chunks = []
40        for group in self.param_groups:
41            for p in group["params"]:
42                if p.grad is None:
43                    chunks.append(torch.zeros_like(p).reshape(-1))
44                else:
45                    chunks.append(p.grad.reshape(-1))
46        return torch.cat(chunks)
47
48    @torch.no_grad()
49    def _metric(self, g):
50        G = torch.stack(list(self._history), dim=1)
51        p = g.numel()
52        if p > self.max_exact_parameters:
53            raise RuntimeError("exact projection is limited to small models")
54        eye = torch.eye(p, device=g.device, dtype=g.dtype)
55        cov = (G @ G.T) / G.shape[1]
56        if self._S is None:
57            S = cov + self.ridge * eye
58        else:
59            S = self.beta * self._S + (1.0 - self.beta) * cov + self.ridge * eye
60        self._S = S.detach()
61        # eigh gives a symmetric, numerically stable inverse square root.
62        eigval, eigvec = torch.linalg.eigh(S)
63        invsqrt = (eigvec * eigval.clamp_min(1e-12).rsqrt()) @ eigvec.T
64        if self.alpha == 1.0:
65            return invsqrt @ invsqrt
66        SinvG = torch.linalg.solve(S, G)
67        K = G.T @ SinvG + self.rho * torch.eye(
68            G.shape[1], device=g.device, dtype=g.dtype)
69        P = invsqrt @ G @ torch.linalg.solve(K, G.T) @ invsqrt
70        return invsqrt @ ((1.0 - self.alpha) * P + self.alpha * eye) @ invsqrt
71
72    @torch.no_grad()
73    def step(self, closure=None):
74        loss = None
75        if closure is not None:
76            with torch.enable_grad():
77                loss = closure()
78        g = self._flat_grad()
79        self._history.append(g.detach().clone())
80        d = self._metric(g) @ g
81        pos = 0
82        for group in self.param_groups:
83            for p in group["params"]:
84                n = p.numel()
85                if p.grad is not None:
86                    p.add_(d[pos:pos+n].view_as(p), alpha=-group["lr"])
87                pos += n
88        return loss