The Numerical Experiment on Homogeneous Spaces

GeometricOptimizers generalizes Adam to homogeneous spaces; Optimization on Homogeneous Spaces derives how. This page is the numerical evidence, reproducing [1].

The numerical experiment

The paper trains a vision transformer on MNIST and Fashion-MNIST: 16 transformer blocks, 7 attention heads, $49 \times 16$ inputs ($28 \times 28$ images cut into 16 patches of $7 \times 7$), batch size 2048, 500 epochs, Float32, no hyperparameter tuning. The projection matrices $W^Q_i, W^K_i, W^V_i$ of the multi-head attention layers are the weights that are put on $St(7, 49)$; the feedforward and classification weights stay unconstrained. Four configurations are compared:

Constraining the projections to the Stiefel manifold is what makes the network trainable at all. The unconstrained baseline suffers from vanishing gradients: with 16 blocks and none of the usual remedies — layer normalization, dropout, regularization, pre-training — the gradient that reaches the early blocks dies, the network collapses onto a trivial prediction $e_i$ and stays there. The plateau is exactly $\sqrt{\tfrac{9}{10} \cdot 2} \approx 1.34$, the $L_2$ distance between a one-hot guess and a one-hot target that disagree on 9 of 10 digits. An orthonormal $Y$ neither amplifies nor damps what passes through a block, which is what removes the problem.

Two conclusions follow. First, hard geometric constraints can replace the heuristics a transformer normally needs, and they do so without adding a hyperparameter — unlike a soft orthogonality penalty $\tilde{L} = L + \mu\Sigma$, which adds $\mu$ and guarantees nothing. Second, among the Stiefel optimizers, Adam gives speed-ups over gradient descent and momentum of the same order as it does in the vector space case, which is the point of generalizing it rather than settling for a first-order method.

Reproducing the experiment

The scripts in scripts/geometric_optimizers/ run that experiment against GeometricOptimizers itself — mnist.jl on the CPU, mnist_cuda.jl on an NVIDIA GPU, mnist_metal.jl on Apple silicon. Running the Experiments says what each of them does and how to start it. Everything below is one run of mnist_cuda.jl on an RTX 4090: the same four configurations, 500 epochs of 29 batches each at a batch size of 2048, Float32, the Cayley retraction — none of the scripts passes retraction, so they all take the default — 6 h 53 min for all four.

The three series the figures plot are checked in under docs/src/data/ — 540 rows in total, distilled from the run by scripts/geometric_optimizers/distill_mnist_results.jl. The figures are therefore rebuilt with the documentation and need neither a GPU nor a rerun.

The training loss

The curves reproduce the MNIST panel of the paper's training-loss figure, and closely: the run ends at a loss of 0.234 for Adam on the Stiefel manifold, 0.693 for momentum, 0.726 for gradient descent and 1.342 for unconstrained Adam. The ordering Adam $<$ momentum $<$ gradient descent $<$ unconstrained is the paper's, and the plateau of the unconstrained baseline is the trivial prediction $\sqrt{1.8} \approx 1.342$ derived above — it is reached within three epochs and held for the remaining 497.

The test accuracy

The paper reports the training loss only, so the accuracies are this repository's addition rather than a reproduction. They say what the loss curves imply: the unconstrained network sits at 0.098, the accuracy of always guessing the same digit, for all 500 epochs, while the three Stiefel runs separate in the same order as their losses. Adam is also the only one of the three that has converged — it reaches 0.86 by epoch 100 and then moves within $\pm 0.01$, whereas gradient descent and momentum are still climbing at epoch 500.

Drift off the manifold

A retraction maps onto the manifold by construction, so in infinite-precision real arithmetic $Y^TY = \mathbb{I}$ would hold after every one of the 14500 steps and this figure would be a flat line at zero. What it actually plots is the rounding error of a finite format — the experiment runs in Float32, as the paper's does — accumulated over those steps.

Two things are worth noting. The growth is linear in the step count rather than a random walk: between epoch 25 and epoch 500 the three curves grow by factors of 19.6, 19.6 and 16.7 while the step count grows by 20, which is an exponent of 0.99, 0.99 and 0.94 against the slope-one guide. So the departure is a systematic accumulation and not noise, and it is set by the number of steps rather than by anything in the geometry. A random walk would give an exponent near 0.5 and, at epoch 500, a drift an order of magnitude smaller. And Adam departs about twenty times faster than gradient descent or momentum, which is what its normalized update $\hat{m}/(\sqrt{\hat{v}} + \varepsilon)$ predicts: its steps are far larger than a raw gradient's, and the error a retraction leaves behind grows with the size of the step it is given.

In absolute terms none of this matters at this length. After 500 epochs Adam is at $8.6\cdot10^{-3}$, which is still orthonormal to two digits, and the other two are two orders of magnitude below that. It is worth knowing rather than worth fixing: linearity means the drift is predictable, so a run several times longer, or one that needs more than two digits, would want the section rebuilt at intervals — a cost that scales with the number of steps and not with the size of the network. Moving to Float64 would push the whole figure down by roughly the ratio of the two epsilons and change nothing about its shape.

References