mlsanityData to 14 Oct 2026

N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs

arXiv 2411.01958 ↗First posted 4 Nov 2024cs.LGcited by 7

On the map: Deep reinforcement learning methods → Deep RL algorithms and representations

Abstract (arXiv)

In-context learning allows models like transformers to adapt to new tasks from a few examples without updating their weights, a desirable trait for reinforcement learning (RL). However, existing in-context RL methods, such as Algorithm Distillation (AD), demand large, carefully curated datasets and can be unstable and costly to train due to the transient nature of in-context learning abilities. In this work, we integrated the n-gram induction heads into transformers for in-context RL. By incorporating these n-gram attention patterns, we considerably reduced the amount of data required for generalization and eased the training process by making models less sensitive to hyperparameters. Our approach matches, and in some cases surpasses, the performance of AD in both grid-world and pixel-based environments, suggesting that n-gram induction heads could improve the efficiency of in-context RL.

Key facts. Cited by 7 papers in the map (data through 14 Oct 2026). Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods. Builds on 25 of its 35 references that are in the map. First posted on arXiv 4 Nov 2024.

Builds on

25 of 35 references.

Closest ancestors · 1

Same direction, newest first

  1. 2024XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning13 Jun

Other references · 24

  1. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2017Neural Discrete Representation Learningcited by 4,712
  4. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,465
  5. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,546
  6. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,204
  7. 2022In-context Learning and Induction Headscited by 964
  8. 2021Transformers Can Do Bayesian Inferencecited by 360
  9. 2024Many-Shot In-Context Learningcited by 206
  10. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 190
  11. 2019Show Your Work: Improved Reporting of Experimental Resultscited by 169
  12. 2024Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokenscited by 153
12 more
  1. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 139
  2. 2022CORL: Research-oriented Deep Offline Reinforcement Learning Librarycited by 119
  3. 2024The Evolution of Statistical Induction Heads: In-Context Learning Markov Chainscited by 118
  4. 2024In-Context Language Learning: Architectures and Algorithmscited by 89
  5. 2023The Transient Nature of Emergent In-Context Learning in Transformerscited by 76
  6. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 55
  7. 2023AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agentscited by 41
  8. 2023Emergence of In-Context Reinforcement Learning from Noise Distillationcited by 27
  9. 2021Showing Your Offline Reinforcement Learning Work: Online Evaluation Budget Matterscited by 24
  10. 2023In-Context Reinforcement Learning for Variable Action Spacescited by 24
  11. 2024Retrieval-Augmented Decision Transformer: External Memory for In-context RLcited by 21
  12. 2022N-Grammer: Augmenting Transformers with latent n-gramscited by 14

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 6; 2026: 1.

7 papers.

  1. 2025Vintix: Action Model via In-Context Reinforcement Learningcited by 10
  2. 2025Yes, Q-learning Helps Offline In-Context RLcited by 10
  3. 2025Zero-Shot Adaptation of Behavioral Foundation Models to Unseen Dynamicscited by 8
  4. 2025Towards Provable Emergence of In-Context Reinforcement Learningcited by 3
  5. 2025Safe In-Context Reinforcement Learningcited by 3
  6. 2025Induction Signatures Are Not Enough: A Matched-Compute Study of Load-Bearing Structure in In-Context Learningcited by 2
  7. 2026Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learnercited by 1

Similar, not cited

Close in wording, no citation either way

  1. 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation
  2. 2025Scalable In-Context Q-Learning
  3. 2025Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning
  4. 2026Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement
  5. 2025Improving In-Context Learning with Reasoning Distillation
  6. 2026Do We Need Adam? Surprisingly Strong and Sparse Reinforcement Learning with SGD in LLMs
  7. 2024Random Policy Enables In-Context Reinforcement Learning within Trust Horizons
  8. 2024In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought