模型最佳化
New Research Finds Adam Breaks Matrix-Factorization Symmetry, Allowing Identical Functions to Converge to Different Attention Representations
The study traces the difference between Adam and low-rank recovery to the coordinate-wise preconditioner’s lack of gauge equivariance. In published experiments, two functionally equivalent initializations trained with Adam produced attention invariants that differed by 56%, though the findings have not yet been validated in large-model pretraining.
A new preprint finds that Adam’s coordinate-wise scaling does more than change convergence speed: the coordinate basis used to represent the parameters may directly affect the function representation ultimately learned. For a factorized matrix `W = UVᵀ`, jointly rotating the two factors to `(UQ, VQ)` leaves `W` unchanged, so the loss function cannot detect this gauge symmetry. Adam’s second-moment estimates, however, accumulate separately for each coordinate, meaning that update directions are no longer equivalent after rotation.
The authors divide optimizers into two groups. Gradient descent, momentum, Adam with a single shared scale, Muon, and Shampoo preserve gauge equivariance; coordinate-wise methods including Adam, RMSProp, Lion, signum, and Adafactor do not. The paper further proves that any memoryless update rule satisfying this symmetry must be a left preconditioner determined by the Gram matrix. Using a parameter that continuously interpolates between coordinate-wise and shared scaling, the researchers experimentally recover the low-rank preference step by step, pinpointing preconditioner anisotropy as the source of the difference.
In underdetermined matrix sensing, every method reduces training error into the interpolation regime, but they differ in the solutions they select: symmetry-preserving updates recover solutions closer to the planted low-rank ground truth. Across two hyperspectral datasets, at equal training loss and the lowest sampling density, gradient descent achieved held-out errors 43% to 44% lower than Adam. The Transformer tests began from two rotated but functionally identical attention initializations. Adam caused their trajectories to diverge after the first step, ultimately producing a 56% relative Frobenius-distance difference between each head’s invariant `W_QᵀW_K`, while the difference under symmetry-preserving optimizers remained near floating-point error.
The public repository includes nine update rules, raw JSONL logs, 25 theorem-identity tests, and a FlowAdam prototype. Its `precond_power` setting interpolates between standard Adam and a shared RMS scale, but fully preserving the symmetry also requires global-norm clipping. The current implementation accepts only one parameter group and requires a closure. For engineering practice, the findings suggest that reparameterization may not be harmless when using Adam. However, the current evidence comes primarily from matrix recovery, small-scale attention experiments, and hyperspectral data, and is not yet sufficient to conclude that replacing Adam with symmetry-preserving optimizers would improve the overall quality of large language models.