YuriiFormer: A Suite of Nesterov-Accelerated Transformers

TL;DR

This work reframes transformer blocks as discrete steps in optimizing a composite energy $\mathcal{E}(X)+\mathcal{F}(X)$ over token configurations, where attention updates follow the gradient of an interaction energy and MLPs follow a potential energy gradient. By applying a Nesterov accelerated gradient template within a Lie--Trotter splitting, the authors introduce YuriiFormer, a family of momentum-enhanced transformers that preserves the same attention and MLP oracles. Empirical results on TinyStories and OpenWebText show consistent improvements in validation loss and downstream few-shot accuracy over a vanilla GPT-style baseline of equal size and budget, with Lie--Trotter variants generally outperforming Euler discretizations. The study demonstrates that optimization-theoretic insights can meaningfully guide architectural design in transformers, offering a modular framework to borrow ideas from numerical optimization and apply them to deep learning models with practical performance benefits.

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.

Figures (5)

• Figure 1: YuriiFormer pre-layerNorm and Euler discretization (left) and Lie--Trotter splitting (right), with velocity LayerNorm $\mathrm{LN}_v$ applied after each velocity update.
• Figure 2: TinyStories (small): loss vs. training step (steps $\ge 1000$). Training loss is Gaussian-smoothed ($\sigma=50$ steps) for readability.
• Figure 3: OpenWebText: validation loss vs. training step for 12-layer and 24-layer models, excluding the first 3k warm-up steps and lightly Gaussian-smoothed ($\sigma=200$ steps) for readability. The ordering is consistent across training: GD+Euler is highest, Nesterov+Lie--Trotter is lowest, and GD+Lie--Trotter (nanoGPT) and Nesterov+Euler lie in between. See Figure \ref{['fig:nesterov-vs-polyak-openwebtext']} for a direct Nesterov vs. Polyak comparison.
• Figure 4: Polyak variant of YuriiFormer with Lie--Trotter splitting (no lookahead), with velocity LayerNorm $\mathrm{LN}_v$ applied after each velocity update.
• Figure 5: OpenWebText: Nesterov vs. Polyak variants YuriiFormer validation loss for small and medium models (steps $\ge 10$k). Curves are nearly identical across sizes, with a small but consistent advantage for Nesterov.