Rethinking Invariance in In-context Learning

TL;DR

This work tackles the sensitivity of in-context learning (ICL) to the order of context examples by formalizing three key principles: permutation invariance, information non-leakage, and context interdependence. It introduces InvICL, a two-stage design that first encodes context independently to preserve leakage constraints and then applies a leave-one-out pre-encoding to entangle context representations, all while maintaining a diagonal attention structure on context rows to ensure invariance. A parallel two-pass implementation enables practical, single-forward-pass inference with $O(n^2)$ complexity, matching the efficiency of baseline methods. Empirically, InvICL outperforms both invariant and non-invariant baselines on synthetic and real-world datasets, including strong length-generalization and OOD performance, underscoring the value of principled invariant ICL for robust generalization and practical deployment.

Abstract

In-Context Learning (ICL) has emerged as a pivotal capability of auto-regressive large language models, yet it is hindered by a notable sensitivity to the ordering of context examples regardless of their mutual independence. To address this issue, recent studies have introduced several variant algorithms of ICL that achieve permutation invariance. However, many of these do not exhibit comparable performance with the standard auto-regressive ICL algorithm. In this work, we identify two crucial elements in the design of an invariant ICL algorithm: information non-leakage and context interdependence, which are not simultaneously achieved by any of the existing methods. These investigations lead us to the proposed Invariant ICL (InvICL), a methodology designed to achieve invariance in ICL while ensuring the two properties. Empirically, our findings reveal that InvICL surpasses previous models, both invariant and non-invariant, in most benchmark datasets, showcasing superior generalization capabilities across varying input lengths. Code is available at https://github.com/PKU-ML/InvICL.

Figures (10)

• Figure 1: Performance of existing ICL algorithms under the settings of zhang2024batch, including auto-regressive (AR) ICL, Prefix ICL raffel2020exploring, BatchICL zhang2024batch and PCW ratner2022parallel. Task prompts are removed for fair comparison.
• Figure 2: The attention masks of four types of ICL, corresponding to different types of ICL methods.
• Figure 3: ICL performance of different models that are trained with (a) 50k epochs and (b) 200k epochs. "Least Squares" is the optimal algorithm for the linear regression task.
• Figure 4: The length generalization behavior of InvICL and AR ICL on HR$\rightarrow$LR setting. The models are meta-trained by sequences with 8 context examples.
• Figure 5: The symmetric positional encoding applied in our work. ${p_i}$ refers to the learned absolute positional embeddings that are added to the token embeddings at position $i$. Figure (a) shows the positional encoding under the standard ICL input sequence. As for the duplicated input of InvICL, we apply the same positional encoding for the original and the repeated examples, as shown in Figure (b).

Theorems & Definitions (14)

• Definition 3.1
• Definition 3.2
• Definition 3.3
• Proposition 3.4
• Proposition 3.5
• Proposition 3.6
• Theorem C.1
• proof
• proof
• Lemma D.1