mlsanityData to 7 Oct 2026

Toward generalizable learning of all (linear) first-order methods via memory augmented Transformers

arXiv 2410.07263 ↗First posted 8 Oct 2024cs.LG math.OCcited by 0

On the map: Mixture-of-experts → Expert routing strategies

Abstract (arXiv)

We show that memory-augmented Transformers can implement the entire class of linear first-order methods (LFOMs), a class that contains gradient descent (GD) and more advanced methods such as conjugate gradient descent (CGD), momentum methods and all other variants that linearly combine past gradients. Building on prior work that studies how Transformers simulate GD, we provide theoretical and empirical evidence that memory-augmented Transformers can learn more advanced algorithms. We then take a first step toward turning the learned algorithms into actually usable methods by developing a mixture-of-experts (MoE) approach for test-time adaptation to out-of-distribution (OOD) samples. Lastly, we show that LFOMs can themselves be treated as learnable algorithms, whose parameters can be learned from data to attain strong performance.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts. Builds on 23 of its 31 references that are in the map. First posted on arXiv 8 Oct 2024.

Builds on

23 of 31 references.

References · 23

  1. 2023GPT-4 Technical Reportcited by 18,132
  2. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  3. 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
  4. 2024Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameterscited by 1,928
  5. 2021What Makes Good In-Context Examples for GPT-3?cited by 1,184
  6. 2021Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivitycited by 1,122
  7. 2021An Explanation of In-context Learning as Implicit Bayesian Inferencecited by 915
  8. 2022Transformers learn in-context by gradient descentcited by 714
  9. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  10. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  11. 2021Linear Transformers Are Secretly Fast Weight Programmerscited by 471
  12. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
11 more
  1. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
  2. 2022Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizerscited by 204
  3. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  4. 2022Self-Adaptive In-Context Learning: An Information Compression Perspective for In-Context Example Selection and Orderingcited by 158
  5. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 97
  6. 2023A Theory of Emergent In-Context Learning as Implicit Structure Inductioncited by 86
  7. 2020Memformer: A Memory-Augmented Transformer for Sequence Modelingcited by 63
  8. 2024Linear Transformers are Versatile In-Context Learnerscited by 28
  9. 2024How Well Can Transformers Emulate In-context Newton's Method?cited by 23
  10. 2024Superiority of Multi-Head Attention in In-Context Linear Regressioncited by 17
  11. 2024How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regressioncited by 13

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression
  2. 2025Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought
  3. 2025Mixture-of-Transformers Learn Faster: A Theoretical Study on Classification Problems
  4. 2021EvoMoE: An Evolutional Mixture-of-Experts Training Framework via Dense-To-Sparse Gate
  5. 2024Natural GaLore: Accelerating GaLore for memory-efficient LLM Training and Fine-tuning
  6. 2024S2FT: Efficient, Scalable and Generalizable LLM Fine-tuning by Structured Sparsity
  7. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization
  8. 2019Learning in Gated Neural Networks