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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Zisman, Ilya, Nikulin, Alexander, Sinii, Viacheslav, Tarasov, Denis, Lyubaykin, Nikita, Polubarov, Andrei, Kiselev, Igor, Kurenkov, Vladislav
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915141900042240
author Zisman, Ilya
Nikulin, Alexander
Sinii, Viacheslav
Tarasov, Denis
Lyubaykin, Nikita
Polubarov, Andrei
Kiselev, Igor
Kurenkov, Vladislav
author_facet Zisman, Ilya
Nikulin, Alexander
Sinii, Viacheslav
Tarasov, Denis
Lyubaykin, Nikita
Polubarov, Andrei
Kiselev, Igor
Kurenkov, Vladislav
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2411_01958
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs
Zisman, Ilya
Nikulin, Alexander
Sinii, Viacheslav
Tarasov, Denis
Lyubaykin, Nikita
Polubarov, Andrei
Kiselev, Igor
Kurenkov, Vladislav
Machine Learning
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.
title N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs
topic Machine Learning
url https://arxiv.org/abs/2411.01958