Breaking through the learning plateaus of in-context learning in Transformer

TL;DR

This paper investigates the causes of learning plateaus in Transformer models' in-context learning and proposes strategies to overcome these barriers, enhancing the models' learning efficiency and capabilities.

Contribution

It introduces a novel conceptual framework separating model components and develops strategies to accelerate in-context learning in Transformers.

Findings

Learning plateaus are linked to impaired weights component functionality.

Proposed strategies effectively break learning plateaus in synthetic tasks.

Strategies also improve in-context learning in NLP applications.

Abstract

In-context learning, i.e., learning from context examples, is an impressive ability of Transformer. Training Transformers to possess this in-context learning skill is computationally intensive due to the occurrence of learning plateaus, which are periods within the training process where there is minimal or no enhancement in the model's in-context learning capability. To study the mechanism behind the learning plateaus, we conceptually seperate a component within the model's internal representation that is exclusively affected by the model's weights. We call this the "weights component", and the remainder is identified as the "context component". By conducting meticulous and controlled experiments on synthetic tasks, we note that the persistence of learning plateaus correlates with compromised functionality of the weights component. Recognizing the impaired performance of the weights component as a fundamental behavior drives learning plateaus, we have developed three strategies to expedite the learning of Transformers. The effectiveness of these strategies is further confirmed in natural language processing tasks. In conclusion, our research demonstrates the feasibility of cultivating a powerful in-context learning ability within AI systems in an eco-friendly manner.