How do simple rotations affect the implicit bias of Adam?

TL;DR

The paper investigates how simple data rotations affect the implicit bias of Adam-like optimizers. It shows that the inherent richness bias of adaptive methods can be destroyed or reversed by small rotations, yielding linear decision boundaries that generalize worse than those learned by SGD. To address this, it introduces EGOP-reparameterization, an orthogonal transform of the optimization objective based on the eigenbasis of the expected gradient outer product, proving rotation equivariance for any first-order method and demonstrating empirical restoration of nonlinear, Bayes-like boundaries. The findings highlight the critical role of data geometry in optimizer-induced generalization and offer a practical, broadly applicable remedy that extends to related preconditioning methods like Shampoo and SOAP.

Abstract

Adaptive gradient methods such as Adam and Adagrad are widely used in machine learning, yet their effect on the generalization of learned models -- relative to methods like gradient descent -- remains poorly understood. Prior work on binary classification suggests that Adam exhibits a ``richness bias,'' which can help it learn nonlinear decision boundaries closer to the Bayes-optimal decision boundary relative to gradient descent. However, the coordinate-wise preconditioning scheme employed by Adam renders the overall method sensitive to orthogonal transformations of feature space. We show that this sensitivity can manifest as a reversal of Adam's competitive advantage: even small rotations of the underlying data distribution can make Adam forfeit its richness bias and converge to a linear decision boundary that is farther from the Bayes-optimal decision boundary than the one learned by gradient descent. To alleviate this issue, we show that a recently proposed reparameterization method -- which applies an orthogonal transformation to the optimization objective -- endows any first-order method with equivariance to data rotations, and we empirically demonstrate its ability to restore Adam's bias towards rich decision boundaries.

Figures (5)

• Figure 1: Illustration of main results. Small data rotations can drastically change generalization. (a) Base algorithms $\texttt{Adam}$ and $\mathtt{SignGD}$ produce nonlinear decision boundaries, similar to the nonlinear Bayes-optimal rule, compared with GD's linear decision boundary. (b) After rotating the data by a small angle ($\gamma = \pi/32$ radians), the base $\texttt{Adam}$ and $\mathtt{SignGD}$ algorithms produce linear boundaries with poor generalization. (c) Training EGOP-reparameterized $\texttt{Adam}$ and $\mathtt{SignGD}$ restores $\texttt{Adam}$'s implicit bias towards nonlinear decision boundaries.
• Figure 2: Illustration of Theorem \ref{['thm:informal-invariant']}. Decision boundaries produced by EGOP-reparameterized $\texttt{Adam}$ and $\mathtt{SignGD}$ are equivariant. As the data distribution rotates by varying angles $\gamma$, the decision boundaries learned by these algorithms also rotate by $\gamma$ radians. For all data rotations, EGOP-reparameterized $\texttt{Adam}$ and $\mathtt{SignGD}$ learn nonlinear decision boundaries that are closer to the Bayes-optimal decision boundary than those produced by gradient descent. Experiment details in Section \ref{['sec:empirical_results']}.
• Figure 3: (a) Illustration of Theorem \ref{['thm:general-rotation-linear-boundary']}. When $\gamma = \pi/4$, $\texttt{Adam}$ and $\mathtt{SignGD}$ learn a linear decision boundary identical to that learned by gradient descent. Compare with \ref{['fig:OG_gamma=0']}, which shows that when $\gamma=0$, the decision boundaries produced by $\texttt{Adam}$ and $\mathtt{SignGD}$ are nonlinear and are closer to the Bayes' optimal decision boundary than that learned by GD. (b) For each epoch $t$ pictured, we visualize the weights produced by training with $\texttt{Adam}$. For every $k\in \{1,\dots,m\}$, we scatter $w_k/\lVert w_k\rVert$ where $w_k \stackrel{\mathrm{ def}}{=} (W_t)_{k,:}$ denotes the $k$th row of weights. As predicted by Theorem \ref{['thm:general-rotation-linear-boundary']}, these values rapidly converge to $a_k \cdot [1, 1]^\mathsf{T}/ \sqrt{2}$. Details in Section \ref{['sec:empirical_results']}.
• Figure 4: Relaxing simplifying assumptions: for $\gamma=\pi/32$, decision boundaries produced by a 2-layer ReLU network including bias terms and trainable outer-layer weights, defined in \ref{['eq:full-network-class']}. $\texttt{Adam}$ and $\mathtt{SignGD}$ still yield the decision boundary predicted by Thm. \ref{['thm:general-rotation-linear-boundary']}.
• Figure 5: For $\texttt{Adam}$ and $\mathtt{SignGD}$ without reparameterization, increasing $\epsilon$ alone does not lead to the nonlinear, good generalization, equivariant boundaries produced by EGOP-repameterization. Here we visualize results using $\texttt{Adam}$ and $\mathtt{SignGD}$, with $\epsilon=3.0$, and without reparameterization. Compare these results with those in Figures \ref{['fig:OG_gamma=0']}, \ref{['fig:OG_gamma=pi_over_32']}, and \ref{['fig:OG_gamma=pi_over_4']}, where $\texttt{Adam}$ and $\mathtt{SignGD}$ are used without reparameterization but with a lower value of $\epsilon$, namely $\epsilon=10^{-8}$.

Theorems & Definitions (27)

• Theorem 1: Informal; see \ref{['thm:general-rotation-linear-boundary']}
• Definition 1: EGOP
• Theorem 2: Informal; see \ref{['thm:EGOP-invariant-decision-boundaries-2d-gamma']}
• Theorem 3
• Lemma 1
• Definition 2
• Lemma 2
• Theorem 4
• Lemma 3
• proof : Proof of Lemma \ref{['lemma:grads-under-rotation']}.