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

arXiv: https://arxiv.org/abs/2411.01958 · HTML: https://mlsanity.com/paper/2411.01958/

## 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.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 25 of 25)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Many-Shot In-Context Learning](https://mlsanity.com/paper/2404.11018/) (2024)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Show Your Work: Improved Reporting of Experimental Results](https://mlsanity.com/paper/1909.03004/) (2019)
- [Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens](https://mlsanity.com/paper/2401.17377/) (2024)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [CORL: Research-oriented Deep Offline Reinforcement Learning Library](https://mlsanity.com/paper/2210.07105/) (2022)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [In-Context Language Learning: Architectures and Algorithms](https://mlsanity.com/paper/2401.12973/) (2024)
- [The Transient Nature of Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2311.08360/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents](https://mlsanity.com/paper/2310.09971/) (2023)
- [Emergence of In-Context Reinforcement Learning from Noise Distillation](https://mlsanity.com/paper/2312.12275/) (2023)
- [Showing Your Offline Reinforcement Learning Work: Online Evaluation Budget Matters](https://mlsanity.com/paper/2110.04156/) (2021)
- [In-Context Reinforcement Learning for Variable Action Spaces](https://mlsanity.com/paper/2312.13327/) (2023)
- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [N-Grammer: Augmenting Transformers with latent n-grams](https://mlsanity.com/paper/2207.06366/) (2022)
- [XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning](https://mlsanity.com/paper/2406.08973/) (2024)

## Built on by (most cited first, 7 of 7)

- [Vintix: Action Model via In-Context Reinforcement Learning](https://mlsanity.com/paper/2501.19400/) (2025)
- [Yes, Q-learning Helps Offline In-Context RL](https://mlsanity.com/paper/2502.17666/) (2025)
- [Zero-Shot Adaptation of Behavioral Foundation Models to Unseen Dynamics](https://mlsanity.com/paper/2505.13150/) (2025)
- [Towards Provable Emergence of In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.18389/) (2025)
- [Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.25582/) (2025)
- [Induction Signatures Are Not Enough: A Matched-Compute Study of Load-Bearing Structure in In-Context Learning](https://mlsanity.com/paper/2509.22947/) (2025)
- [Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner](https://mlsanity.com/paper/2604.05112/) (2026)

## 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.