In plain words: Treats each transformer layer as one optimization step on token embeddings, where attention and the MLP are gradient steps on two different energies. Adding a classic momentum trick to those same steps beat a standard GPT-style model on two text datasets.
Abstract
We propose a variational framework that interprets transformer layers as iterations of an optimization algorithm acting on token embeddings. In this view, self-attention implements a gradient step of an interaction energy, while MLP layers correspond to gradient updates of a potential energy. Standard GPT-style transformers emerge as vanilla gradient descent on the resulting composite objective, implemented via Lie--Trotter splitting between these two energy functionals. This perspective enables principled architectural design using classical optimization ideas. As a proof of concept, we introduce a Nesterov-style accelerated transformer that preserves the same attention and MLP oracles. The resulting architecture consistently outperforms a nanoGPT baseline on TinyStories and OpenWebText, demonstrating that optimization-theoretic insights can translate into practical gains.
Aleksandr Zimin, Yury Polyanskiy, Philippe Rigollet
arXiv:2601.23236 · cs.LG, cs.AI, math.OC, stat.ML · submitted Jan 30, 2026 · updated Mar 4, 2026
abstract · pdf · html