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

TL;DR

This paper tackles the data inefficiency and training instability of in-context reinforcement learning (ICRL) by integrating N-Gram Induction Heads (NGH) into transformer-based agents. Building on Algorithm Distillation, it adds an NGH layer that directly encodes n-gram patterns into attention, with n-gram matching extended to pixel-based observations via vector quantization. The results show substantial data efficiency gains (up to 27x less data in some tasks), faster hyperparameter optimization, and successful application to both discrete and image-based environments, while preserving or improving baseline performance. This method offers a practical, more scalable path to robust ICRL across diverse observation spaces and task distributions.

Abstract

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.

Figures (10)

• Figure 1: Performance comparison for different number of training goals between our method and Algorithm Distillation (AD), an in-context reinforcement learning method laskin2022context. Our method demonstrates similar performance with less training goals (128 vs. 512) and in general outperforms the baseline. See \ref{['sec:res']} for results.
• Figure 2: (Left) The Key-to-Door environment. The key and the door are shown for illustrative purposes only; the agent does not see their location during training. (Right) An observation from the Miniworld environment.
• Figure 3: Results on Dark Room. We search through hyperparameters in random order and report expected maximum performance dodge2019show. We also constrain the number of optimization steps by 10K and use equal batch size to ensure both methods use the same amount of data. The top row shows experiments with different number of learning histories, with the total number of training goals fixed. It is seen that our method needs much less hyperparameter assignments (20 for 1K histories) to find the optimal model, while the baseline performance increases only asymptotically (full plots are shown in \ref{['apndx:plots']}). The number of traning tasks for this experiment is 60. The bottom row presents experiments with varied number of goals and fixed number of learning histories. Our method makes it possible to find the optimal hyperparameters with only 15 hyperparameter assignments, while the baseline fails to work in such low data conditions. However, none of the methods can learn to generalize from only 10 goals. The number of learning histories for this task is 1K.
• Figure 4: Results on Key-to-Door. We demonstrate the ability of our method to generalize when the task diversity is limited. We fix the total number of goals with 100, significantly shrinking the number of learning histories. Keep in mind that for the baseline method to converge to a model with the same performance, it needs 2048 goals and 2048 learning histories laskin2022context. We show that our method needs 27x less data comparing to baseline (see \ref{['apndx:data_calc']} for justification). The baseline method can no longer converge with that few data and its performance plateaus with the increasing number of hyperparameter assignments, while N-Gram model shows near-optimal performance.
• Figure 5: Results on Miniword environments. We show that our method is applicable not only for environments with discrete observations, but also for the image-based ones. The settings of the Miniworld environments are similar to Dark Room and Key-to-Door. The main outcome of these experiments is that we can successfully implement n-gram matching in for images and get similar results to the discrete environments. The details of the setup are described in \ref{['res:pixel']}. (Left) We fix the number of tasks in Miniworld-Dark to 30 and the number of learning histories to 50. The N-Gram layer significantly enhances performance, resulting in nearly-optimal model, while the baseline quickly saturates around suboptimal return. The evaluation is done on 50 goals. (Right) For the experiment in Miniworld-Key-to-Door we fix 300 training tasks and 50 learning histories. The results follow the pattern observed in discrete environments, where the N-Gram layer refines the performance of the baseline model. We evaluate the models on 100 unseen goals.