Macierze projekcji query i key w mechanizmie attention transformerów są obecnie trenowane przy użyciu optymalizatorów euklidesowych bez żadnych ograniczeń na ich geometrię. Zespół badawczy zaproponował inną drogę - ograniczenie tych macierzy do rozmaitości Stiefel i optymalizowanie ich za pomocą Riemannian Adam, który operuje jednym skalarem drugiego momentu na ramkę, ogranicza krok poprzez region zaufania i zastosowuje polarną retrakcję.
Mathematyczne dowody pokazują, że ta aktualizacja jest spadkiem najbardziej stromym w osadzonej metryce, niezależnym od skali gradientu, dobrze uwarunkowanym i dokładnie O(d)-ekwiwariantnym. Kluczem do sukcesu jest zjawisko odkryte przez badaczy: weight decay ma identycznie zerowy gradient Riemannowski na rozmaitości Stiefel, ponieważ W = WIr leży w przestrzeni normalnej. To oznacza, że nauczona geometria attention przetrwa cykle upadku, które decay zwykle powoduje w pozostałej części modelu.
Wyniki eksperymentów są imponujące. Na zadaniu modularnej arytmetyki grokking samo jedno uruchomienie osiągnęło 97,0% dokładności walidacji w epoce 20000 w porównaniu z 61,1% dla linii bazowej. Na łatach CIFAR-10 metoda zyskała średnio 8,98 punktu procentowego nad AdamW z wynikami na poziomie t=60,6 dla 12 sparowanych startów, a różnica rosła wraz z dodatkową ilością danych zamiast się zmniejszać. Mechanika bazowa prosta - aktualizacja Riemannowska ze stałym krokiem ma stopień jeden w stosunku do gradientu, co powoduje, że przesuwa się 24-40 razy mniej na krok niż identycznie ukształtowana macierz AdamW, co oznacza że ramki ledwie opuszczają swoje inicjalizacje.