---
title: Rotational Adam Optimizer
url: https://www.emergentmind.com/topics/rotational-adam-optimizer
type: topic
---

# Rotational Adam Optimizer

Rotational Adam Optimizer is an extension of adaptive gradient methods that seeks to restore or enforce rotation-equivariance—insensitivity to orthogonal changes of basis—in Adam-style preconditioning. Standard Adam operates with strictly per-coordinate second-moment adaptation, which breaks rotation equivariance and can result in significant algorithmic artifacts, degraded training speed, and lost generalization benefits when the parameterization or data undergoes orthogonal transformation. Recent work has developed both theoretically principled and pragmatically scalable forms of Rotational Adam, including matrix preconditioner diagonalization, expected gradient outer product reparameterization, symmetry-aware adaptation, Riemannian manifold generalizations, and vector-grouped updates.

## 1. Failure Modes of Standard Adam under Rotations

Adam maintains running averages of the first and second moments of stochastic gradients for each parameter entry, updating as
$$
m_t = \beta_1 m_{t-1} + (1-\beta_1) \nabla f(\theta_{t-1}) \\
v_t = \beta_2 v_{t-1} + (1-\beta_2) [\nabla f(\theta_{t-1})]^2,
$$
with elementwise operations. The parameter update is
$$
\theta_t = \theta_{t-1} - \alpha \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}.
$$
This formulation is sensitive to the coordinate basis: under an arbitrary rotation $R$, the moments transform (first moment: $m_t \mapsto R m_t$; second moment: $v_t \mapsto$ not generally $R v_t$, due to the elementwise square), but the elementwise division does not commute with $R$. Standard SGD, by contrast, is fully rotation-equivariant. Empirically, Adam’s performance—both convergence speed and implicit bias—degrades under random, layerwise, or global basis rotations in transformer and vision architectures, with the degree of slowdown correlated with the scale and type of rotation. For instance, global rotations can increase GPT-2 training time by ~16% and decrease ViT/S ImageNet-1K convergence rate by up to 96%, while ResNet-50 is robust to such transformations [2410.19964].

## 2. Mathematical Foundations of Rotation Equivariance

An optimizer $A$ is rotation-equivariant if for any orthogonal $R$ and all $t$,
$$
R w_{t+1} = A(\{R w_i\}_{i\leq t}, f^{(R)}, t),
$$
where $f^{(R)}(w) = f(R^T w)$. Adam’s coordinatewise adaptation (division by $\sqrt{v_t}+\epsilon$) violates this property because, in general, $R [m_t/(\sqrt{v_t}+\epsilon)] \neq [R m_t]/(\sqrt{v_t'}+\epsilon)$; the update trajectory depends on the choice of parameter basis [2410.19964, 2205.13599].

Several approaches have been developed to achieve rotation-equivariance by restructuring Adam’s adaptation, including:

1. **Full-matrix second moment adaptation**:
   $$ M_t = \beta_2 M_{t-1} + (1-\beta_2) g_t g_t^T $$
   $$ \theta_t = \theta_{t-1} - \alpha M_t^{-1/2} \hat m_t $$
   (This update is formally rotation-invariant.)

2. **Dynamic basis diagonalization**: Diagonalize $M_t$ as $Q_t D_t Q_t^T$, then update in the rotated basis, mapping back to original coordinates. This generalizes SVD-based rotation strategies to full preconditioning [2410.19964, 2502.07488].
   
3. **Block-wise or vector-grouped moments**: Aggregate second moments over structured subsets (e.g., rows, channels, global vectors), yielding equivariance and eliminating axis artifacts [2205.13599].

## 3. Rotation-Equivariant Adam Variants

A variety of practical rotational strategies for Adam optimization have emerged:

### 3.1 EGOP-Reparameterized (Covariance-Aligned) Adam

Using the Expected Gradient Outer Product (EGOP) matrix,
$$
P = \mathbb{E}_{\theta \sim \rho}[\nabla f(\theta)\nabla f(\theta)^T]
$$
with eigendecomposition $P = Q\Lambda Q^T$, a fixed orthonormal rotation $T=Q^T$ is applied to parameters, mapping $\varphi = T\theta$. The Adam update is performed in $\varphi$-space and mapped back. This approach is rotation-equivalent and exploits the dominant principal gradient directions present in tasks with spectral decay, yielding improved step efficiency and insensitivity to parameterization [2502.01594, 2510.23804].

### 3.2 Adaptive Preconditioner Diagonalization (AdaDiag++/Rotational Adam)

The empirical gradient covariance $M_t$ is diagonalized via periodic SVD (on matrix-shaped parameters),
$$
G_t = P_t\Sigma_t Q_t^T, \qquad T_t = Q_t \otimes P_t
$$
and parameter updates are performed in the basis diagonalizing $M_t$, with diagonal preconditioning, and mapped back:
$$
\theta_{t+1} = \theta_t - \alpha T_t^T \frac{\hat{\tilde m}_t}{\sqrt{\hat{\tilde v}_t} + \epsilon}
$$
This yields large improvements in convergence metrics for large-scale language and vision models, with significant step and epoch reductions (LLaMA pretraining: 2x fewer steps vs Adam; ResNet/ImageNet-1K: 1.2–1.5x fewer epochs) [2502.07488].

### 3.3 VectorAdam

VectorAdam groups vector-valued parameters (e.g., 3D coordinates, neural features) and computes a scalar second moment for each block,
$$
m_t^{(i)} = \beta_1 m_{t-1}^{(i)} + (1-\beta_1) g_t^{(i)}, \\
v_t^{(i)} = \beta_2 v_{t-1}^{(i)} + (1-\beta_2) \|g_t^{(i)}\|_2^2, \\
\theta_{t+1}^{(i)} = \theta_t^{(i)} - \alpha \frac{\hat m_t^{(i)}}{\sqrt{\hat v_t^{(i)}}+\epsilon}
$$
removing axis-aligned update artifacts in geometric learning and adversarial point-cloud optimization [2205.13599].

### 3.4 Adaptive and Symmetry-Aware Rotational Policies (ARO)

ARO (Adaptively Rotated Optimization) introduces an adaptive rotation for each parameter matrix, selecting $R_t$ via a Procrustes-QR policy to maximize a dual-norm loss decrease proxy, then performing the Adam-style or alternative normed descent in this rotated basis. ARO unifies and extends prior “eigen-rotation,” “SOAP,” and “Muon” schemes, and delivers up to 1.3–1.35x step-speedup over AdamW at LLM scales with negligible added wall-clock cost [2602.09006].

## 4. Algorithms and Implementation

### 4.1 EGOP-Reparameterized Rotational Adam

- Estimate EGOP $P$ by sampling gradients under distribution $\rho$.
- Compute $P = Q\Lambda Q^T$; set $T=Q^T$.
- Transform initial parameters: $\varphi_0 = T\theta_0$.
- For $t=0..T-1$:
  - $g_t = \nabla f(\theta_t)$; $\hat g_t = T g_t$.
  - Update Adam moments in $\varphi$-space: $m_t$, $v_t$.
  - Bias-correct, step update: $\varphi_{t+1} = \varphi_t - \eta m_t/\left(\sqrt{v_t}+\epsilon\right)$.
  - Map back: $\theta_{t+1} = Q \varphi_{t+1}$.
- Use block-wise or low-rank $Q$ for tractability at scale [2502.01594, 2510.23804].

### 4.2 Preconditioner Diagonalization (AdaDiag++)

- Aggregate $M_t$ (second-moment estimator) as a moving average of gradient outer products.
- Periodically diagonalize $M_t$ via SVD (reshaping gradient), yielding rotation $T_t$.
- Apply Adam updates in rotated space, then inverse-rotate updates back.
- AdafacDiag integrates with Adafactor for memory efficiency; only row/column second-moment statistics are needed [2502.07488].

### 4.3 Riemannian/Stiefel Manifold Rotational Adam

- For optimization on orthogonality-constrained spaces (e.g., Stiefel manifold),
- Project gradients, accumulate moments in the intrinsic tangent space,
- Retract by computing exponential map update: $X_{t+1} = \exp(\Omega_t)X_t$.
- Preserves manifold constraints, and carries Adam’s adaptivity over non-Euclidean geometries [2305.16901].

### 4.4 Adaptively Rotated Optimization (ARO)

- Maintain running momentum $M_t$ (matrix).
- Compute lookahead step via base optimizer (e.g., Adam).
- Select $R_t$ by maximizing a dual-norm proxy; implement via Cholesky–QR on Gram matrices.
- Transform gradients, moments, and apply update in $R_t$-rotated basis.
- Retain per-iteration overhead $O(mn + m^3)$ for $m \times n$ matrices, compatible with LLM-training workloads.
- Supports hierarchical/global/shared rotation schemes [2602.09006].

## 5. Theoretical and Empirical Properties

Rotational Adam methodologies inherit the convergence properties of standard Adam in the Euclidean or Riemannian sense, as all bias-correction, adaptation, and isometric mappings (rotations, eigen-bases) preserve the necessary convexity and bounded-gradient assumptions in online convex optimization [2502.01594, 2502.07488, 2305.16901]. When the EGOP spectrum decays, or when curvature structure is block-diagonalizable, rotational approaches deliver substantial convergence acceleration and restore invariance. Edge cases where the curvature is isotropic (flat EGOP) do not benefit from rotation.

Empirical results include:
- 2x step reduction in LLaMA pretraining compared to Adam [2502.07488].
- Restoration of “richness bias” in small-rotation ReLU nets (decision boundaries remain nonlinear and Bayes-optimal) [2510.23804].
- Elimination of axis-aligned artifacts in geometry and adversarial optimization [2205.13599].
- Stable, memory-efficient integration with Adafactor yields similar performance at order-of-magnitude lower storage [2502.07488].
- ARO delivers 1.3–1.35x step-speedup over AdamW and 1.1–1.15x over orthogonalization methods across a range of LLM scales, with controlled benchmarking [2602.09006].

## 6. Practical Considerations and Memory/Complexity

Rotational Adam variants incur additional costs:
- EGOP or full-matrix approaches: $O(d^2)$ per iteration for $d$-dimensional parameters (impractical for very large $d$).
- Block-wise or low-rank schemes decrease per-step cost to $O(Lk^2)$, $L$ blocks of size $k$.
- Periodic SVD/amortized QR rotation cost is small relative to forward/backward passes in large models, especially with windowed basis updates [2502.07488, 2602.09006].
- VectorAdam reduces memory footprint for vector-structured blocks by a factor of $n$ (for $r\times n$ parameter blocks), and slightly decreases computation by eliminating per-coordinate variances [2205.13599].
- AdafacDiag maintains sublinear memory comparable to Adafactor, suitable for large-scale deployment.

## 7. Connections, Limitations, and Future Work

Rotational Adam aligns with a broader framework of symmetry-aware optimization, where natural group actions (e.g., rotations) leave the objective function invariant. Generalizations extend to:
- Dynamic or per-layer rotation sharing; cross-module rotational symmetry exploitation.
- Hybrid permutation-equivariant and gauge-invariant extensions.
- Riemannian gradient and moment transport for non-Euclidean manifolds.
- Adaptive rotation selection via data-driven or curvature-informed policies (e.g., via dual-norm maximization) [2602.09006].

Current limitations include:
- Intractability of full-matrix preconditioning at extreme ($d > 10^5$) scale unless blocked/low-rank.
- Diminishing returns as the eigenvalue spectrum of the second-moment matrix flattens.
- The need for expert tuning of block granularity, SVD update periods, and memory-efficient representation in large transformer settings.

By exposing and counteracting Adam’s rotation-pathologies, Rotational Adam optimizers provide a unified, theoretically-grounded, and empirically-validated path to efficient and robust adaptive optimization in modern large-scale learning [2410.19964, 2502.01594, 2510.23804, 2502.07488, 2205.13599, 2602.09006].

Source: https://www.emergentmind.com/topics/rotational-adam-optimizer