Positional Information Matters for Invariant In-Context Learning: A Case Study of Simple Function Classes
This addresses the problem of improving ICL robustness for AI models, particularly in handling distribution shifts, though it is incremental as it builds on existing architectures like transformers and DeepSet.
The study investigated in-context learning (ICL) limitations in transformers by analyzing linear regression tasks under distribution shifts, finding that preserving permutation invariance (termed ICL invariance) is crucial for out-of-distribution performance, with transformers using identical positional encodings achieving state-of-the-art results across various shifts.
In-context learning (ICL) refers to the ability of a model to condition on a few in-context demonstrations (input-output examples of the underlying task) to generate the answer for a new query input, without updating parameters. Despite the impressive ICL ability of LLMs, it has also been found that ICL in LLMs is sensitive to input demonstrations and limited to short context lengths. To understand the limitations and principles for successful ICL, we conduct an investigation with ICL linear regression of transformers. We characterize several Out-of-Distribution (OOD) cases for ICL inspired by realistic LLM ICL failures and compare transformers with DeepSet, a simple yet powerful architecture for ICL. Surprisingly, DeepSet outperforms transformers across a variety of distribution shifts, implying that preserving permutation invariance symmetry to input demonstrations is crucial for OOD ICL. The phenomenon specifies a fundamental requirement by ICL, which we termed as ICL invariance. Nevertheless, the positional encodings in LLMs will break ICL invariance. To this end, we further evaluate transformers with identical positional encodings and find preserving ICL invariance in transformers achieves state-of-the-art performance across various ICL distribution shifts