# Oja Gradient-Subspace Optimizer

- ID: 1571
- Canonical URL: https://synthcore.org/idea/1571/oja-gradient-subspace-optimizer
- API JSON: https://synthcore.org/api/idea/1571.json
- API Markdown: https://synthcore.org/api/idea/1571.md
- Verification status: unverified
- Source: [arXiv:2607.29034](https://arxiv.org/abs/2607.29034)
- Category: dynamics
- Solves: stability, speedup, scalability
- ML areas: optimizer, training-dynamics, mlp, transformer
- Math tags: dynamical-systems, linear-algebra, spectral-theory, control-theory, numerical-analysis
- Ratings: usefulness 6/10; difficulty 5/10; novelty 6/10

## Idea description

Track the dominant rank-r subspace of the gradient covariance online, then use that basis to construct a low-rank adaptive update or a controlled preconditioner. Unlike offline PCA refreshes, the Oja flow continuously follows changing training geometry while preserving orthonormality, potentially reducing the cost of second-order or Shampoo-like methods.

## Mathematical statement

Let g_t in R^d be the flattened minibatch gradient and let A_t in R^(d x d) be an exponentially averaged gradient covariance, A_t = beta A_(t-1) + (1-beta) g_t g_t^T. Maintain U_t in St(r,d) = {U: U^T U = I_r} using the Oja flow epsilon dU/dt = (I_d - U U^T) A_t U, where epsilon > 0 is the response-time parameter. For a fixed symmetric A with eigenvalues lambda_1 >= ... >= lambda_d and gap Delta = lambda_r - lambda_(r+1) > 0, the dominant-subspace error is locally expected to decay on a time scale epsilon / Delta. For a time-varying eigenspace, tracking requires its drift rate to be sufficiently smaller than this response rate. The neural update uses P_t = U_t U_t^T and either retains only P_t g_t or applies a bounded gain to the tracked subspace: g_tilde = [P_t + alpha(I_d - P_t)] g_t, with 0 <= alpha <= 1.

## Key formulas

- $$\varepsilon\frac{dU(t)}{dt}=(I_d-U(t)U(t)^\top)A(t)U(t),\qquad U(t)^\top U(t)=I_r.$$
- $$A_t=\beta A_{t-1}+(1-\beta)g_tg_t^\top,\qquad U_{t+1}=\operatorname{qf}\!\left(U_t+\frac{h}{\varepsilon}(I_d-U_tU_t^\top)A_tU_t\right),$$
- $$\widetilde g_t=\left[U_tU_t^\top+\alpha(I_d-U_tU_t^\top)\right]g_t,\qquad \theta_{t+1}=\theta_t-\eta_t\widetilde g_t,$$
- $$\Delta=\lambda_r(A)-\lambda_{r+1}(A)>0,\qquad \rho\approx\exp\left(-\frac{\Delta h}{\varepsilon}\right).$$

## Implementation notes

1. Integration point: start with a small MLP or transformer block and maintain one Oja basis per weight matrix, or one basis for the concatenated gradients of a selected layer. Flatten the minibatch gradient to g_t, update the covariance estimate, update U_t, and use the projected or softly projected gradient in the ordinary optimizer. Do not materialize A_t for large layers; estimate A_t U_t directly as beta times the previous A_U plus (1-beta) times g_t times (g_t^T U_t), followed by the projected Oja increment. 2. Pseudocode: compute g; set A_U = beta A_U + (1-beta) g(g^T U); set Z = U + (h/epsilon)(A_U - U(U^T A_U)); set U = qr(Z); set P_g = U(U^T g); set g_tilde = P_g + alpha(g - P_g); apply Adam or SGD to theta using g_tilde. 3. Computed from the paper: the Oja vector field, Stiefel projection, and eigengap-dependent response mechanism. Estimated empirically: lambda_r, lambda_(r+1), subspace drift, and the useful value of epsilon. 4. First cheap experiment: train a width-512 MLP on CIFAR-10 with SGD, comparing SGD, low-rank projected SGD, and soft projection for r in {8, 32, 128}. Freeze covariance statistics after a stationary phase and measure principal-angle error against a periodically computed exact PCA basis. For a fixed covariance with measured gap Delta, vary h/epsilon and fit the error decay; the predicted slope is -Delta/epsilon in continuous time, or log(rho)/h approximately equal to -Delta/epsilon. Under a controlled rotating covariance whose top subspace rotates at angular speed omega, the tracking error should increase sharply when omega becomes comparable to Delta/epsilon; test whether the observed transition is within 20 percent of this boundary. Reject settings where the discrete increment loses orthogonality or causes exploding training loss.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.