EGOP-Based Reparameterization

• The paper demonstrates that reparameterizing with EGOP yields improved convergence rates and invariance in adaptive optimizers.
• The method involves estimating the EGOP matrix, performing eigendecomposition, and transforming coordinates to align with the gradient's spectral structure.
• Empirical tests confirm faster convergence and enhanced robustness against arbitrary rotations in supervised learning and convex optimization tasks.

The Expected Gradient Outer Product (EGOP)–based reparameterization is an orthogonal, data-driven change of basis for parameter space in optimization, designed to improve the behavior of coordinate-wise adaptive gradient methods such as Adam and Adagrad. Motivated by the non-equivariance of these optimizers under orthogonal transformations, EGOP-based reparameterization constructs a basis aligned to the spectral structure of the expected outer product of the gradients: $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$. Diagonal-adaptive methods applied in this basis exhibit improved convergence and invariance properties, especially when $G$ is low-rank or exhibits strong spectral decay. The method is supported by both theoretical guarantees and a range of empirical findings in supervised learning and convex optimization contexts (DePavia et al., 27 Oct 2025, DePavia et al., 3 Feb 2025).

1. Expected Gradient Outer Product and Choice of Basis

The EGOP matrix, $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$, captures the typical geometry of the gradient field under a chosen parameter distribution $\rho$ (DePavia et al., 3 Feb 2025). The key insight is that the eigensystem of $G$ provides an optimal basis, in the sense that the expected directional variation of the gradient is maximally concentrated along the leading eigenvectors. Since methods like Adam employ coordinate-wise statistics, the efficiency and implicit bias of these methods are intimately coupled to the chosen basis. When rotated into the EGOP eigenbasis, the EGOP becomes diagonal, mitigating the adverse effects of arbitrary rotations on adaptive optimizers.

2. EGOP-Based Reparameterization Procedure

The standard workflow comprises:

1. EGOP Estimation: Draw $M$ independent samples $\theta^{(i)}\sim\rho$; compute gradients $g_i = \nabla f(\theta^{(i)})$; estimate $\widehat G = \tfrac{1}{M}\sum_{i=1}^M g_i g_i^\top$.
2. Eigen-Decomposition: Compute the eigendecomposition $\widehat G = U\Lambda U^\top$, with $G$0 and diagonal $G$1.
3. Coordinate Change: Set new variables $G$2; the reparameterized objective is $G$3.
4. Optimization: Run Adam, Adagrad, or similar optimizer directly in $G$4-space, with all standard updates carried over. For Adam: compute $G$5, maintain moments, and update $G$6; reconstruct $G$7 if needed.

This process is "one-shot," typically performed once prior to training, though online extensions recompute the EGOP and basis every $G$8 steps (DePavia et al., 27 Oct 2025, DePavia et al., 3 Feb 2025).

3. Theoretical Properties: Equivariance and Convergence Acceleration

The fundamental theoretical advantage is rotational equivariance: for any orthogonal input transformation $G$9, Adam or Adagrad run in the EGOP basis produces iterates that align with the input rotation (Theorem 4.3 of (DePavia et al., 27 Oct 2025)). This endows the optimizer with invariance to arbitrary basis changes of the data or parameterization, eliminating the instabilities and implicit bias reversals otherwise observed.

Convergence improvements stem from the alignment with the spectral structure of $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$0. The Adagrad (or Adam, under analogous assumptions) convergence rate in the reparameterized coordinates replaces the worst-case sum of per-coordinate smoothness constants by a sum dominated by the leading eigendirections of $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$1. If the EGOP spectrum decays rapidly (low stable rank $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$2), a speed-up by up to $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$3 is theoretically supported, with the criterion for acceleration being $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$4 and dense leading eigenvectors (large $$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$5) (DePavia et al., 3 Feb 2025).

4. Empirical Evidence

Empirical tests confirm these theoretical claims:

• Synthetic 2D Gaussian Classification: In (DePavia et al., 27 Oct 2025), Adam's "richness bias" is destroyed by arbitrary small input rotations, often converging to simple linear solutions. With EGOP reparameterization, Adam consistently recovers nonlinear decision boundaries invariant to rotation, matching the Bayes-optimal up to isometry.
• Multilayer Linear Networks: (DePavia et al., 3 Feb 2025) reports faster and lower-loss convergence for Adagrad/Adam with one-shot EGOP reparameterization across synthetic architectures with enforced spectral decay, with no observable overfitting.
• Real-World Data (UCI digits, Fashion-MNIST): Block-wise EGOP reparameterization in two-layer ReLU networks yields improved training loss and validation accuracy for Adam and Adagrad, while SGD (rotation-invariant) is unaffected.
• Convex Problems and Ablations: Reparameterized Adam/Adagrad converge more rapidly for log-sum-exp, logistic regression, and least-squares tasks with decaying spectrum, while ablations show that merely increasing regularization (ε in Adam) does not restore rotational robustness (DePavia et al., 27 Oct 2025).

Empirical metrics include distance to Bayes-optimal boundary, test error, and boundary curvature.

5. Algorithmic and Implementation Aspects

Key steps in the EGOP-based reparameterization algorithm are:

Step	Description	Computational Cost
EGOP Estimation	$$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$6 independent gradients, empirical mean	$$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$7
Eigen-Decomposition	$$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$8 eigendecomposition (full or top-$$G = \mathbb{E}_{\theta\sim\rho}[\nabla f(\theta)\nabla f(\theta)^\top]$$9)	$\rho$0 (or less)
Basis Transformation	Forward-backward pass in rotated coordinates	$\rho$1 per step

Practical heuristics include using a top-$\rho$2 truncated basis, layer-wise ("block") EGOP for neural networks, or randomized SVD for large $\rho$3. The additional overhead is limited to an upfront cost, with negligible per-iteration penalty.

Suggested hyperparameter choices are $\rho$4 and $\rho$5 usually isotropic Gaussian or uniform on small balls. For numerical stability, Adam's $\rho$6 often needs to be set higher (e.g., $\rho$7 in (DePavia et al., 27 Oct 2025)).

6. Applicability, Limitations, and Extensions

EGOP-based reparameterization is effective when:

• The EGOP spectrum exhibits strong decay ($\rho$8).
• The leading eigenvectors are dense ($\rho$9).
• Adaptive methods outperform SGD in the native coordinates.

Primary limitations are the upfront computational cost and the practicality of full-basis updates for very large models. Layer-wise or top-$G$0 eigenbasis truncations can mitigate this overhead.

The approach extends naturally to any coordinate-wise adaptive method, including AdaMax, Shampoo/SOAP, and others. A plausible implication is that similar schemes may improve the behavior of these optimizers on large neural architectures (e.g., transformers), but detailed analysis remains an open question (DePavia et al., 3 Feb 2025).

Open questions include optimizing incremental-basis updates, scaling to large models, and further characterization of implicit bias under EGOP reparameterization.