Does Adam Favor a Basis?

1. The model in one screen

The task is associative recall. An episode writes \(k\) key-value pairs and then asks for the value of each key again in a shuffled order. The model has exactly one trainable object, a matrix \(W_K \in \mathbb{R}^{d_k \times V}\) whose columns, after normalization to unit length, are the key vectors \(k_a\). Everything else is fixed. Values are one-hot vectors \(v_b = e_b\), and the memory is a matrix \(S \in \mathbb{R}^{V_b \times d_k}\) that starts at zero every episode.

Two operations touch \(S\). A read at key \(a\) computes scores \(r = S k_a\) and predicts \(\mathrm{softmax}(g\, r)\) over value classes. A write when the value \(v_b\) arrives corrects the memory by the prediction error,

This is the delta rule, and it is one gradient step on \(\tfrac12\|S k_a - v_b\|^2\) with respect to \(S\). Because \(\|k_a\| = 1\), the new read at that key is \((1-\beta)\,S_{\text{old}}k_a + \beta v_b\), so at \(\beta = 1\) the current pair is stored exactly. The catch is that every other key \(c\) also moves, by \(\beta\,\delta\,(k_a^{\top}k_c)\). Orthogonal keys ignore each other. Overlapping keys interfere.

Training runs across many episodes and updates only \(W_K\), differentiating through all the writes. So the outer loop learns where to put the keys, and the inner loop (the delta rule) does the remembering. The quantity that summarizes the key layout is the Gram matrix \(G = K^{\top}K\), the table of all pairwise overlaps \(k_a^{\top}k_c\). A Gram matrix is just that table, with ones on the diagonal because the keys are unit length.

Things to try. Put two keys close together and watch the later write damage the earlier one. Put two keys at a negative overlap (pointing roughly opposite) and notice the interference changes sign, suppressing the wrong class instead of boosting it. That sign asymmetry is why a condition number alone does not capture retrieval quality. With three keys in two dimensions, no arrangement stores all three exactly. The memory has only two input directions.

2. The model has no favorite axes

Here is the fact the whole question rests on. Take any orthogonal matrix \(Q \in \mathbb{R}^{d_k \times d_k}\). An orthogonal matrix is a rotation, possibly combined with a reflection, so \(Q^{\top}Q = I\) and it preserves lengths and angles. Apply it to every key and compensate in the memory,

Then every read is unchanged, since \(S' k_a' = S Q^{\top} Q k_a = S k_a\), and every write stays consistent, since \(\delta' k_a'^{\top} = \delta\, k_a^{\top} Q^{\top}\) is exactly how \(S'\) should move. The Gram matrix is unchanged too, because \(K'^{\top}K' = K^{\top} Q^{\top} Q K = G\). Two key layouts that differ by a rotation produce identical predictions and identical loss on every episode.

In plain terms, the model cannot tell which way its coordinate axes point. Only the pairwise geometry of the keys matters. So if a trained model's keys end up lined up with the axes of \(\mathbb{R}^{d_k}\), the model's loss cannot be what asked for that. Something else did. The hypothesis is that the something is Adam.

3. What a step looks like, GD vs Adam

Gradient descent moves the flattened parameter vector \(w\) by \(w \leftarrow w - \eta\, \nabla L\). Rotate the whole problem by \(Q\) and the gradient rotates with it, since \(\nabla_{w'} L(Qw) = Q\, \nabla_w L(w)\) when \(L\) itself is rotation invariant, as ours is. So the GD step rotates too. This property has a name, rotational equivariance. Rotate first then step, or step first then rotate, you land in the same place. Two GD runs started from \(W_0\) and \(QW_0\) on the same data stay exact rotated copies of each other forever.

Adam keeps two running averages per coordinate. With \(g_t = \nabla L(w_t)\),

where \(g_t^2\), the square root, and the division all act elementwise, and the hats denote the standard bias corrections \(\hat m_t = m_t / (1-\beta_1^t)\), \(\hat v_t = v_t/(1-\beta_2^t)\). Squaring coordinate by coordinate does not commute with a rotation, so the Adam step does not rotate when the gradient rotates.

The cleanest way to see it is the very first step. At \(t = 1\) the bias corrections give \(\hat m_1 = g_1\) and \(\hat v_1 = g_1^2\), so with \(\epsilon \approx 0\)

Adam's first step is sign descent. Every coordinate moves by exactly \(\eta\), regardless of how big its gradient was. Later steps are a smoothed version of this, each coordinate rescaled by its own recent gradient magnitude.

4. Why coordinatewise steps can prefer axes

Look again at the left canvas in Demo 3. The GD step always ends on the circle, and the sign step always ends on a corner of the square. That is not decoration. Gradient descent is steepest descent when you measure step length with the ordinary Euclidean norm, whose unit ball is a sphere. Sign descent is steepest descent when you measure step length with the max norm \(\|x\|_\infty = \max_j |x_j|\), whose unit ball is a cube. A sphere looks the same from every direction. A cube does not. Its corners and faces are tied to the coordinate axes.

The symmetries a cube does keep are permutations of the axes and sign flips of the axes. Adam has exactly those symmetries and no others. Relabel your parameters, or negate some of them, and Adam does the same thing in the relabeled coordinates. Rotate them by \(37^\circ\) and it does something different. So Adam cannot prefer a particular axis, but it can prefer the axes as a set, which is why the right measurement is a quantity that is invariant under permutations and sign flips but not under general rotations.

That is a reason the optimizer could break the symmetry. It is not yet a mechanism for why learned features would end up aligned with axes rather than merely pushed around. The honest status in the literature is that the alignment is observed in transformers trained with Adam and absent with SGD, and the per-coordinate normalization is the accepted explanation, but the step-by-step story of how the alignment forms is not settled. A rough intuition that is at least consistent with the update rule is this. A feature that lives in a single coordinate has its whole gradient in one slot, and Adam gives that slot a full-sized step. A feature spread across many coordinates shares its gradient among them, and each slot still gets a step of roughly size \(\eta\), so the feature as a whole moves faster in the Euclidean sense. Both effects make the dynamics depend on how features sit relative to the axes, which GD never does. Which one wins, and whether either produces alignment in a given model, is an empirical question.

5. The controlled test, running live

The clean way to test this is a paired-run design. Draw one initialization \(W_0\) and one random orthogonal \(Q\). Train from \(W_0\) and from \(QW_0\) on the same episode stream with zero initial optimizer moments. Do this once with GD and once with Adam. Then track three things.

1. Orientation divergence \(\|W_t' - QW_t\|_F\). Frobenius norm means square root of the sum of squared entries. For GD this must stay at floating-point noise. For Adam it grows. The GD pair is the control that proves the harness is correct.
2. Gram divergence \(\|G_t' - G_t\|_F\). This is the one that can matter for recall. If the two Adam runs disagree only in orientation, their Gram matrices stay equal and their loss stays equal. If their Gram matrices split, Adam's basis sensitivity changed something the model can feel.
3. Axis concentration \[ A(K) = \frac{1}{V}\sum_{a=1}^{V}\sum_{j=1}^{d_k} (k_a)_j^{4}. \] For one unit vector, \(\sum_j k_j^4\) is \(1\) when the vector sits on an axis and \(1/d_k\) when it is spread evenly across all axes. For a direction drawn uniformly at random it averages \(3/(d_k+2)\), so a freshly initialized Gaussian \(W_K\) starts near that value. A sharper baseline than this constant is available. Take the learned keys, rotate them by many random \(Q\), and compute \(A\) each time. That baseline has the same intrinsic geometry as the learned keys and differs only in orientation. If the learned \(A\) sits well outside the spread of that baseline, the orientation is special.

The simulation below implements exactly the model from Section 1, with analytic gradients through every write, and runs all four trainings at once in your browser. It is tiny on purpose. The point is to make the metrics concrete, not to settle the question.

Run it once at the defaults. You should see the GD divergence flat at around \(10^{-15}\) (that is double-precision noise, the proof that the harness respects the symmetry) and the Adam divergence climbing within a few dozen steps to order \(10^{-1}\). That part is a theorem, and the plot just confirms the code. The interesting plots are the Gram divergence column of the table and the \(A(K)\) comparison. At this scale the Adam pair's Gram matrices usually do split (so the two Adam runs learned different key geometries, not just different orientations), while the axis-concentration excess over baseline is small and swings between seeds. The seed sweep is there so you do not fool yourself with one lucky run.

Also watch the mean column norm. GD leaves the raw norm of each column growing slowly, since for a scale-invariant parameter the gradient is orthogonal to the parameter. Adam's steps are not orthogonal to anything in particular. The forward pass is blind to these norms, but the effective learning rate on the normalized key is not, which is a confound you must control when comparing optimizers.

6. Reading results honestly

Suppose someone runs this experiment at real scale, say on a transformer, and gets numbers. Three separate statements govern what they can and cannot mean, and each rules out a tempting but wrong conclusion.

Rotation alone cannot change recall. Section 2 proved it. So if you find that Adam-trained keys are more axis-aligned than GD-trained keys, that fact by itself says nothing about which model remembers better. You have found a fingerprint, not a performance effect.

Orientation can matter for recall only through intrinsic geometry. Adam's basis preference is a statement about dynamics. The dynamics decide where in Gram-matrix space training ends up. If Adam's trajectory lands at a different \(G\) than GD's would (different signed overlaps, different active-set spectra), then recall differs, and the orientation preference was the cause through the path training took. The Gram divergence between the two Adam runs is the simplest witness that the path depends on orientation at all.

Separate speed from representation. Adam usually trains faster in wall-clock steps. If you compare at a fixed step count, every difference you see mixes "Adam got further" with "Adam went somewhere different". Make two comparisons, one at matched training budget and one at matched held-out loss. A difference in \(A(K)\) that survives the matched-loss comparison is a representation effect. One that disappears is probably a speed effect.

7. Blind spots and the wider map

Things adjacent to this question that are worth knowing if you want to test it properly.

Where "privileged basis" comes from

Anthropic's 2023 interpretability note on privileged bases in the residual stream observed outlier dimensions in a 200M-parameter model, ruled out several architectural causes, and arrived by elimination at Adam's per-coordinate normalization as the likely culprit without proving it. Caples and Neuhaus then trained matched 12M-parameter models with Adam and with SGD plus momentum and found excess kurtosis above 100 under Adam and near zero under SGD. The experiment on this page is the same logic ported to a setting where the symmetry is exact and provable, which is a real advantage. In a transformer, "rotational equivariance of the architecture" is only approximately true (embeddings, unembeddings and the data all have their own structure), and commenters on the Caples and Neuhaus post pressed on exactly that point.

Adam's sensitivity to rotation is now a research topic in its own right

Maes, Zhang and coauthors (2024) measured how Adam's training loss on GPT-2-scale models degrades under random rotations of parameter space, and showed that structured rotations can also help. Their conclusion is that the standard rotation-invariant convergence theory for Adam cannot explain why Adam works, because Adam's advantage lives in the basis. DePavia and coauthors (2025) show the implicit bias of Adam on linear classification can be brittle under small data rotations, and that an orthogonal reparameterization restores rotation equivariance. Both papers give ready-made ways to measure how much an optimizer depends on the basis.

GD is not the only control

Any optimizer that is equivariant under \(W \mapsto QW\) works as a control, and some are much stronger learners than plain GD. Muon, which orthogonalizes the momentum matrix before applying it, is equivariant under orthogonal maps on either side of a matrix parameter, because the polar factor of \(QM\) is \(Q\) times the polar factor of \(M\). Shampoo and SOAP are also natural comparison points, since they adapt in a learned basis rather than the coordinate basis. A three-arm experiment (GD, Adam, Muon) separates "adaptive step sizes" from "coordinatewise adaptive step sizes", which a two-arm one cannot.

Adam is only sometimes sign descent

The first-step identity in Section 3 is exact. After that, Adam equals sign descent only when gradients are consistent across steps relative to their noise. Balles and Hennig (2018) decompose Adam into a sign part and a variance-adaptation part. Kunstner and coauthors (2023) argue that sign descent, not noise handling, explains most of Adam's advantage over SGD on transformers. A cleaner test of the hypothesis adds plain sign descent as a fourth arm. It has the cube geometry with none of the moving-average machinery.

Confounds specific to this model

• The forward pass is invariant to each column's raw norm, but the optimizers are not. Equalize or track norms, or the comparison is partly about effective learning rate.
• \(\epsilon\) in Adam interpolates between sign descent (\(\epsilon \to 0\)) and plain momentum SGD (\(\epsilon\) large). Sweep it. If the axis effect fades as \(\epsilon\) grows, that is strong evidence it comes from the normalization.
• Small \(d_k\) makes \(A(K)\) a coarse instrument. The gap between isotropic (\(3/(d_k+2)\)) and fully aligned (\(1\)) is only \(0.4\) at \(d_k = 3\), and chance fluctuations of a few keys fill much of it. A convincing test needs \(d_k\) of 16 or more, with the random-rotation baseline and its spread reported every time.
• The axis concentration of the keys is a different object from the kurtosis of activations in the transformer result. In this model, the analog of an activation is the memory state \(S\) or the read \(Sk_a\). Measuring axis structure there too would make the bridge to the transformer result explicit.