Momentum Streams for Optimizer-Inspired Transformers

• URL: http://arxiv.org/abs/2605.24425v1
• Date: Sat, 23 May 2026 06:40:27 GMT
• Title: Momentum Streams for Optimizer-Inspired Transformers
• Abstract summary: We build a family of Transformers inspired by momentum-based designs.<n>In our main pretraining experiment, the triple-momentum TMMFormer achieves the lowest validation loss.<n>We show that TMMFormer and other momentum-based designs reach flatter minima than the vanilla Transformer.
• Score: 8.060910957442397
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: The residual update of a pre-norm Transformer layer admits an interpretation as one step of a first-order optimizer acting on a surrogate token energy, wherein the attention and MLP sublayers function as gradient oracles. Based on this observation, we build a family of optimizer-inspired Transformers (triple-momentum, Adam/AdamW, Muon, SOAP) and compare them under matched compute. In our main pretraining experiment, the triple-momentum TMMFormer achieves the lowest validation loss, outperforming the vanilla Transformer and prior architectural variants. A controlled ablation and supporting theory show that momentum, not preconditioning, is the main source of the gain. We further show that TMMFormer and other momentum-based designs reach flatter minima than the vanilla Transformer, which leads to less forgetting and better generalization.

Related papers

• A Constrained Optimization Perspective of Unrolled Transformers [77.12297732942095]
  We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms.<n>We observe constrained transformers achieve stronger to perturbations robustness and maintain higher out-of-distribution generalization.
  arXiv  Detail & Related papers  (2026-01-24T02:12:39Z)
• EMAformer: Enhancing Transformer through Embedding Armor for Time Series Forecasting [21.876566019196677]
  EMAformer is a model that enhances the Transformer with an auxiliary embedding suite.<n>It achieves state-of-the-art performance on 12 real-world benchmarks.<n>It reduces forecasting errors by an average of 2.73% in MSE and 5.15% in MAE.
  arXiv  Detail & Related papers  (2025-11-11T16:12:44Z)
• Tricks and Plug-ins for Gradient Boosting with Transformers [17.43386196818751]
  BoostTransformer augments transformers with boosting principles through subgrid token selection and importance-weighted sampling.<n>Our method incorporates a least square boosting objective directly into the transformer pipeline, enabling more efficient training and improved performance.
  arXiv  Detail & Related papers  (2025-08-04T21:54:16Z)
• Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization [88.5582111768376]
  We study the optimization of a Transformer composed of a self-attention layer with softmax followed by a fully connected layer under gradient descent on a certain data distribution model. Our results establish a sharp condition that can distinguish between the small test error phase and the large test error regime, based on the signal-to-noise ratio in the data model.
  arXiv  Detail & Related papers  (2024-09-28T13:24:11Z)
• MoEUT: Mixture-of-Experts Universal Transformers [75.96744719516813]
  Universal Transformers (UTs) have advantages over standard Transformers in learning compositional generalizations. Layer-sharing drastically reduces the parameter count compared to the non-shared model with the same dimensionality. No previous work has succeeded in proposing a shared-layer Transformer design that is competitive in parameter count-dominated tasks such as language modeling.
  arXiv  Detail & Related papers  (2024-05-25T03:24:32Z)
• Setting the Record Straight on Transformer Oversmoothing [35.125957267464756]
  As model depth increases, Transformers oversmooth, i.e., inputs become more and more similar. We show that smoothing behavior depends on the eigenspectrum of the value and projection weights. Our analysis reveals a simple way to parameterize the weights of the Transformer update equations to influence smoothing behavior.
  arXiv  Detail & Related papers  (2024-01-09T01:19:03Z)
• Emergent Agentic Transformer from Chain of Hindsight Experience [96.56164427726203]
  We show that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches. This is the first time that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches.
  arXiv  Detail & Related papers  (2023-05-26T00:43:02Z)
• Transformers learn in-context by gradient descent [58.24152335931036]
  Training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass.
  arXiv  Detail & Related papers  (2022-12-15T09:21:21Z)
• Transformers from an Optimization Perspective [24.78739299952529]
  We study the problem of finding an energy function underlying the Transformer model. By finding such a function, we can reinterpret Transformers as the unfolding of an interpretable optimization process. This work contributes to our intuition and understanding of Transformers, while potentially laying the ground-work for new model designs.
  arXiv  Detail & Related papers  (2022-05-27T10:45:15Z)
• Pay Attention to MLPs [84.54729425918164]
  We show that gMLP can perform as well as Transformers in key language and applications. Our comparisons show that self-attention is not critical for Vision Transformers, as gMLP can achieve the same accuracy. In general, our experiments show that gMLP can scale as well as Transformers over increased data and compute.
  arXiv  Detail & Related papers  (2021-05-17T17:55:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.