Rethinking Invariance in In-context Learning
Authors: Lizhe Fang
Yifei Wang
Khashayar Gatmiry
Lei Fang
Yisen Wang
Venue: ICLR 2025 Poster (2025)
Keywords: In-context Learning
Permuation Invariance
Primary area: foundation or frontier models, including LLMs

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 \emph{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.

Published as a conference paper at ICLR 2025
RETHINKING INVARIANCE IN IN-CONTEXT LEARNING
Lizhe Fang Yifei Wang?* Khashayar Gatmiry²
Lei Fang
Yisen Wang
State Key Lab of General Artificial Intelligence,
School of Intelligence Science and Technology, Peking University
2 MIT CSAIL
3 School of Economics, Peking University
4 Institute for Artificial Intelligence, Peking University
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 perfor- mance 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 (ImvILL, 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, most benchmark datasets, showcasing superior generalization capabilities across varying input lengths. Code available at http://ghhbbbbooo//KKUUUMMM/
INTRODUCTION
In-Context Learning (ICL) has shown to be a key emergent property of large language models (LLMs) (Brown et al. 2020). By utilizing sequence of examples as the context, LLMs can be adapted quickly and accurately to new tasks without parameter tuning (Wang et al. 2024 Kossen
ctt al. 2024; Wang et al. 2025). Despite its impressive potential, ICL exhibits a crucially unusual behavior: sensitiviry to the order of context examples (Lu et al., 2022, Zhao et al. 2021 Xie et al.
2021 Agrawal et al. 2022). Although context examples are independent, the order which they are presented can dramatically influence ICL predictions, with variations from about 90% to 50% on the SST-2 dataset (Lu al. 2022).
is easy to note that the auto-regressive (AR) nature of LLMs is the root order sensitivity. AR- LLMs often utilize a so-called causal mask in the attention module, which breaks the permutation invariance property of the de facto Transformer architecturd] As the context examples are intrinsically equivalent under different permutations, a model that respects this data symmetry tends to enhance both learning and generalization (Sokolić al. 2016 Bietti et al. |021 Tahmasebi & Jegelkal 2023). Therefore, recent works have proposed several variant algorithms of ICL to achieve the invariance by modifying the Transformer architecture (e.g., Prefix ICL. (Raffel et al. 2020), PCW (Ratner et al.)
2022), and BatchICL (Zhang al. 2024 ) However, they often perform inferior to non-invariant counterparts like AR ICI., as we extensively observed in practice shown in Figure [
We note that although desirable, the invariance property alone insufficient for good ICL performance (e.g., a model with constant output f(-) e invariant yet provides no useful information). Therefore, to ensure the performance of ICL, we need to satisfy the following two properties while making ICL invariant: 1) Information Non-leakage: it prevents the query from accessing its answer, thereby avoiding shortcuts and enabling dense learning signals for ICL. by allowing the prediction of every context example in the input. 2) Context Interdependence: Each context example interacts with all preceding examples. As the sequence lengthens, more information is provided, thereby enhancing
"Equal Contribution.
i 'Corresponding Author: Yisen Wang (yisen. wang @ pkk..edu.cn). Besides, sequential positional encodings (PEs) of the prompt also introduce order sensitivity.

---

Published as a conference paper at ICLR 2025
prediction accuracy. However, existing methods less compromise these properties when making ICL invariant (Table resulting the lack of well-performing invariant ICL method.
Motivated by the analysis above, we de- sign effective Invariant In-context Learning (InvCCL) algorithm that main- tains these essential properties, ensuring both invariance and high performance. InvICL addresses the issuc order sen- sitivity (invariance), not only avoiding information leakage but also enhancing context interdependence beyond what is achievable with AR-LLMs. To facili- tate practical implementation, velop a fully parallel version of InvICI, capable of obtaining all Leave-One Onee-Onn--uut (L.OO) embeddings and predictions single forward pass using novel LOO- type attention mask. Empirically, InvICI. outperforms existing invariant ICL ver- sions, and even surpasses AR-ICI. (non- invariant) on most tasks both synthetic and real-world datasets.
AR ICL
Prefix ICL
60
BatchiCL
PCW
also de-
SST2
RTE
Dataset
Figure 1: Performance of existing ICL algorithms der the settings of (Zhang et al. 2024). including auto- regressive (AR) ICL, Prefix ICL (Raffel et al. (2020). BatchICI. (Zhang et al. 2024) and PCW (Ratner al.) 2022). Task prompts are removed for fair comparison.
We summarize our contributions as fol- lows:
We undertake a comprehensive exploration into designing invariant ICL. algorithms, high- lighting the importance of preserving information non-lcakage and context interdependence.
We propose InvICI, which synergizes the goals invariant ICL. algorithms by utilizing leave-one-out embeddings to achieve invariant predictions and information non-leakage while maximizing context interdependence.
Empirically, InvICL indeed achieves superior performance across range of tasks on both synthetic and real-world datasets.
Table 1:
Comparisons of different ICL types (details in Section information non-leakage, context interdependence, and performance.
2
permutation invariance,
ICL Type
Invariance Non-leakage
Interdependence
Performance
Auto-regressive Prefix (full attn.)
(partiall)
A (baseline)
A-
lng-of-Example InvICL (ours)
X
A-
A/
2
PRELIMINARIES
Consider classification task with few i.idd. training examples D = {x; := (x), where denotes the input and y; denotes the classification target. An ICL algorithm takes these training examples (a.k.a. context examples) together as imput and then predicts a new test example x,- A general formulation of is
8n. Sr) f(x;, Ve......... Xn: Vn. x,);
where ÿ; denotes the label prediction for x, Note the predictions for context example, {y.)a are optional but they are generally available for AR-LIMs.
popular model choice for ICL. is Transformer (aawwani al. 2017). where the self-attention layer is the elementary module. Denote H (hi,... h,) be the input hidden state of self-attention layer, outputs
H- AHW.P. .., where A softmax (HW,HW)) MM)
2

---

Published as a conference paper at ICLR 2025
Quer
conton careples
Query
Quenry
Query
52
11
22
22
.22
43
33
33
2t
Zt
223
22
Key
51 2g R3
21
de Key
Excodinnng
(a) Auto-regressive ICL
(b) Prefix ICL
(c) BBg-of-Exammp
(d) Invariant
Figure 2: The attention masks of four types of ICL, corresponding different types of ICL methods.
where W 2s WR, W., P denotes the query, key, value, and projector matrices, respectively. M {0. 00) -00 nxn an attention mask. For standard (or full) self-attention layer, M zero matrix, while a causal self-attention layer utilizes the following causal mask:
0
0
M
As result, the softmax attention A only has nonzero weights in its lower triangular terms. Notably, the form Eq. ( can be generalized other attention types, as will be discussed later.
Revisiting existing Yransformer-basee ICL algorithms, they can be categorized into three families depending on their aggregation scheme over the context examples: Auto-regressive ICL, 2) Prefix ICI., and 3) Bag-of-Exampa ICL.
Auto-regressive ICL (AR ICL), A naive way to perform ICL is to adopt the original auto-regressive Transformer (Radford al. 2018), which admits the following aggregation rules
hxc aggr (((b., hy. hxr
[n + 1].
(4)
where hx; hy:: he denote the encodings of x,, Vi. (X&, yk), respectively. Here let X,+1 X for notation simplicity. Therefore, every example hş only attends to the previous ones h = (h....., bb J. which introduces sequential order to the input examples. As former examples have smaller context. later examples in the sequence enjoy higher accuracy, shown in Liu al. (2022):
Wu et al. (202)). Figure 2(a) illustrates the implementation by applying causal mask M, which is exactly the form in Eq.
Prefix ICL. To fully utilize the information every context example, the causal mask is discarded in Prefix LM (Raffel et al. 2020). Therefore, it aggregates over all context examples as
hx + aggr (hx, by. 1 & E [n]]
(5a)
hx, + aggr (((,x, by. ) hx.)
(5b)
Figure 2(b) illustrates the implementation by modifying the attention mask M in Eq. (2). where utilize full attention among the context examples {x, Ja= and causal attention on the test example X;
Bag-of-Exxe Example ICL (BoE ICL). In addition to the two conventional designs above, there is new variant of ICL. Methods like PCW (Ratner et al., 2022). SAICL (Cai al., (023), and BatchICI. (Zhang et al.] (0224) encode each context example yi) independently (without considering other context examples), similar to the "hag-of-word" representation. Its aggregation rules can be
formulated
hx hyd aggr ((hx, hy,)) ). V. h [n];
(6a)
hx, f aggr {(hx, by, hx.
(6b)
Figure 2(c) illustrates implementation (PCW (Ranner al. (022 by modifying the attention
mask M. restricts attention to occur only within each context example X;, [n], preventing cross-attentions between them, while retaining attention between the test example X, and context examples.

---

Published as a conference paper at ICLR 2025
3 THE PROPOSED INVARIANT IN-CONTEXT LEARNING (INVICL)
We begin with formalizing the desiderata of invariant ICL (Section 3.1). and explore how to meet all these desiderata (Section Next, we introduce how to implement our proposed InvICI. method in practice (Sections) 3.3
3.1 INVARIANT ICL AND ITS DESIDERATA
We begin with a formal characterization of three important desiderata in invariant in-context learning.
1) Invariance. In an ICL task, we have the prior knowledge of data symmetry that the n context examples & are independently identical distributed (i.i..). We define an iCL algorithm that preserves this symmetry property as an invariant ICL algorithm:
Definition 3.1. An ICL algorithm f is said be (permulation) invariant if its last prediction f satisfies fi( Xn, xx) = (x, Xi,, Xx) for any (f1....,ān) E Sa. a permutation of
[n]
2) Information Non-leakage. During training, AR-LLMS learn to dynamically predict each interme- diate context example x; based on its previous tokens x as its own context, leading to n prediction tasks that provide rich learning signals for ICL. To achieve this, an essential architectural inductive bias is the causal mask, which ensures that the prediction of each query x;, such as ý:, does not have access to its ground-truth answer y;; otherwise, the prediction task would become trivial. We believe this principle should be generally adhered when designing ICL algorithms. We name this the information non-leakage principle, formally described below.
Definition 3.2. An ICL algorithm f has no information leakage if its prediction of every example x; is invariant to its label y, (with others fixed), i.e., f(...,X,, Yi....... ) = ((.... x, y.....). ); holds for any two labels y:, E V, where f(-); denotes the i-th element of f()).
3) Context Interdependence. Another advantage of AR-LLMS is that they allow the encoding of each example to depend on other (previous) examples. These examples provide the context for better encoding of x,, which in turn improves the prediction of future examples when x; serves as their context, We name this property as context interdependence. Unlike the information non-leakage principle, this property requires that the prediction of each example x; should flexibly depend on as many other context examples as possible.
Definition 3.3. An ICL algorithm is context-interdepeneer for if the prediction of x; dependent on other examples. Formally, for any j * i, there exists xy) such that [(.., Xi, Vi...., Yj........) ); * Ni,.,,..,i,......... XY,,.........,,,,,"""".......... Y........ with other examples fixed.
Limitations of Previous Methods. Through a close examination, we find that no existing ICL methods satisfy all these principles: 1) AR ICL avoids information leakage and has partial context interdependence, but its sequential structure breaks permutation invariance; 2) Prefix ICL maintains permutation invariance and full context interdependence, but leaks information; 3) BoE ICL achieves permutation invariance and prevents information leakage through independent encoding, but it sacrifices context interdependence and limits the flexibility of context representations. These properties are summarized in Table [ design an ICI. algorithm that achieves all three properties.
Motivated by the limitations of previous methods, we aim to
3.2 A PRINCIPLED DESIGN OF INVARIANT ICL
In this section, we we explore how to design ICL algorithms that preserve all three principles: permutation invariance, information non-leakage, and context interdependence. In a Transformer, the only interaction among different examples occurs in the self-attention layer (Eq. 2 Conceptually, self-attention can be viewed as a message- -aasgggg-painggggg scheme on digraph of n examples, denoted as G, with the adjacency matrix defined by the attention score matrix A E Raxn,
A softmax (HW (HW)) TT + +M)
where M (0 2n is constant mask matrix. Under this definition, Ai represents the message passed from the j-th example to the i-th example. The message passing then updates with H + AH (informal) using A as the propagation matrix. Here, we only consider the graph of context examples |- we always want the test example &, be fully aware of the context examples,
4

---

Published as a conference paper at ICLR 2025
making these edges trivial. A straightforward way to design invariant ICL algorithm, commonly used
by existing works (such as Prefix ICL (Raffel et al. 2020), PCW (Ratner al. [022), and SAICL
(Cai et al. 2023 ), is to modify the attention mask M as it is the only controllable factor in Eq. ( Therefore, in the following, we discuss how to design the attention mask M to meet these desiderata.
All the proofs are in Appendix D
Permutation Invariance by Three-choice Mask. Intuitively, permutation invariance requires the attention mask M to exhibit some form of symmetry. Notably, both the Prefix and BoE masks (Figure 2/b, c)) satisfy permutation invariance, while the causal mask (Figure 2(a) does not. The following proposition explores whether other attention masks can also achieve this property.
Proposition 3.4. Given input matrix H = (hi,... .... TR) E TRaxd with the features of the context examples only. The permutation invariance of ICL outputs (Definition 3.I holds if and only if the attention mask on the context examples, M, belongs to M = (M1, M2, 0 where
0
-00
-00
30
0
0
1
M =
00
Proposition 3.44 demonstrates that achieve permutation invariance, the attention mask on the context example must fall into one of the three choices in M: 0 corresponds to full attention in Prefix ICI, M, corresponds BoE ICL,; and the attention score before softmax under M2 is the linear combination of that of M, and 0 (only cross-attention between tokens without self-attention).
Information Non-leakage by Lower Triangular Mask. According to Zheng et al. (0118), ensuring information non-leakage equivalent to guaranteeing the message- message-passing process through the graph is acyclic (except for self-loops). This imposes the following restriction on the attention mask M.
Proposition 3.5. An ICL algorithm realizes information non-leakage if and only if it is possible to reorder context examples such that the attention mask on context examples M is lower triangular.
Combining the conditions for attention masks outlined in Propositions 3.4 & 3.5 (belong to M and lower triangular), we find that the attention mask on context examples must be a diagonal matrix, as concluded in the following proposition.
The message-passing scheme respects permutation invariance and information non-leakage if and only if the attention mask on context examples M is diagonal.
Proposition 3.6.
Therefore, we conclude that if an ICL algorithm preserves both permutation invariance and infor- mation non-leakage, its attention mask not only can be, but has to be in the form depicted in Figure [[c] Specifically, must take the form of rag-of-examples (BoE) ICL, encoding each example individually before aggregation as in Eq. (6), denoted as:
hx,
BoE (f(bx. 11 h
(8)
However, discussed in Section
3.1 BoE lacks context interdependence.
Context Interdependence through Pre-encoding. While context interdependence cannot be imple- mented within a single propagation step among context examples, it can still be achieved by encoding each context example with the context of other samples, a process the three principles simultancously, the pre-encoding step must also adopt the form of a BoE ICL scheme, where it aggregates independent encodings of all other samples (i.e., leave-one-out encoding):
term pre-encoding. To ensure
h BoE (1(bx, hy hipk hx.) k. [n)
where hx,, hy: are the independent encoding (similar to Eq. (6a)). Therefore, we arrive at a two-stage ICL method as follows. First, encode each context example with leave-one-our (LOO) BoE encoding as Eq. (9) Then, in the second stage, we utilize these contextual encodings to predict the test examples as in Eq. (8) This approach guarantees the three desiderata of invariant ICL.
Symmetric Positional Encuding. As minor point, ensure the symmetry of the model, it is also necessary to incorporate permutation invariance into the positional encoding. We adopt an independent position encoding scheme that treats each example as an independent sequence. It is also applicable BoF ICL and Prefix ICI. for ensuring permutation invariance. Details in Appendix A.1
5

---

Published as a conference paper at ICLR 2025
Finally, reach proposed InvICL (Invariant In-context Learning). The propagation pro- cess for InvICL is outlined in Algorithm where h the encoding of x, the k-th layer of Transformer.
Algorithm Invariant In-context Learning
Require: {hxx h embedding of context examples; hx ( embedding of the ICL query 1: for k == 1 to #TranfformerLaeess do
2: for 1 to n do
3: Compute the independent encoding of context examples:
hhx; hy;
(k) (k)
aggr((hx ( (k-1) hy, )) (where h = hx, (0) 4: end for
5: for 1 to n do
6: Compute the leave-one-out pre-encoding of the i-th context example: (h (K) (8)
aggr( (1) by:
((1-1)))a ))zei hx, 1
7: end for
8: Update (h) = aggr ((b
(k-1) hy. (it-"")) }
9: end for
3.3 PARALLEL IMPLEMENTATION
we have developed truly invariant ICL algorithm achieving the three desiderata. However, a significant drawback of the encoding scheme in Algorithm iis its computational cost. For each sequence of context examples, it requires LOO forward passes to pre-encode each example, plus an additional forward pass for the final prediction. This results in a total of n + 1 forward passes for a single prediction. In contrast, AR ICL, BoE ICL, and Prefix ICL can all be implemented in parallel using a single forward pass by modifying the attention mask the form illustrated in Figure aa b, c).
In Section 3.2
Parallel Computation via Unrolling. To address the computational cost issue, we propose parallel implementation for InvICL, leveraging the cfain-of-thought idea from LLM reasoning (Wei al. 2022). While implementing InvICL within a single forward pass of the input sequence (x1, Xn Is challenging. this difficulty can be overcome by unrolling the input sequence twice. As illustrated Figure 2(d) we duplicate the context examples twice as ($....., Xn, X1.... X,, xx) and perform two-step forward process in parallel to encode the context examples. In the first step, we perform a
(k) (h in Algorithm 1p. In the second step. apply
BoE-style encoding of each context example
LOO-style attention mask to obtain the LOO encodings each example (h, (h{*) in Algorithm 1) that the LOO encodings (b(()) predict the test
aware of all other context examples. At last, example x,. This unrolling scheme enables us to accomplish InvICL in a single forward pass, which results in the same complexity order (nn²) as the baselines.
EXPERIMENTS
4.1 SYNTHETIC SCENARIO
To evaluate the in-context capability of InvICL, we conduct a series of experiments inspired by Garg et al. (2022). Taking the linear regression task for example, we train model to perform linear regression using in-context learning, i.e., the model takes the sequence (xI, Xn, ((x,), ), x,)
as input and predicts g(x;) where g is linear function. Detailed experimental settings are provided in Appendix A.3 We compared the ICL performance across four models: 1) Auto-regressive (AR) (Radford et al. 2019); Prefix (Raffel et al., 2020); 3) -feemmls- (BoE) (Ratner et al. 2022):
and 4) NoPE (t.e, removing the positional encoding) (Kazemnejad et al. 2024). The MSE loss was reported for models trained over various epochs, as illustrated in Figure The key insights from our experiments are as follows:
InvICL Converges Fast. At 50k epochs, only InvICL demonstrates good ICL performance (Figure 3(a). while other models perform well at later epochs (Figure

---

Published as a conference paper at ICLR 2025
20
1.0
Fodio 12
2.4
PriSe.c CL
2.4
5.0
(a) 50k Epochs
(b) 200k Epochs
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.
InvICL Has a Strong Length Extrapolation Ability. The models are trained with a sequence length of 40. As shown in Figure 3(b) when the test sequence length exceeds 40, it is clearly that InvICL > AR ICL Prefix ICL = NoPE> BoE ICL in terms of performance. InvICL. On one hand, this result
This indicates the strong length generalization capability confirms the conventional conclusion that model that respects the data symmetry enjoys better generalization capability. On the other hand, it highlights that preventing information leakage and maintaining context interdependence are crucial for an invariant ICL algorithm.
We further conduct experiments in out-of-distribution tasks and other function settings in Appendix B.I. and present the trend of loss as it changes with the training epochs in Appendix B.4 Both experiments demonstrate that InvICL's out-of-distriben performs AR ICL. Additionally, in Appendix B.5 conduct linear probing experiments further demonstrate how the architecture of InvICL impacts the model's internal representations.
in-context performance consistently out-
4.2 REAL-WORLD DATASETS
In this part, we conduct experiments to evaluate the capacity of InvICL on real-world datasets. Since ICL tasks generally different from the pertaining one and some ICL methods introduce new masking schemes for aggregation (significantly different from the masking in pretrained model), short finetuning of the pretrained model on the ICL tasks using these new ICL methods is necessary to fully utilize the pretrained model's s capacity for ICL. (Min et al. 2022b Wei et al.. 2021 [yer et al.) 2022 Cai et al. 2023). Here, we follow MetalCL (Min et al. 2022b) to do the short finetuning and evaluation.
As in MetalCL, utilize 142 tasks including text classification, question answering (QA), natural language inference (NLI), and paraphrase detection. For each training iteration, we first sample a task T, from the C' meta-training tasks, and then sample k 4++1 training examples (X1, y.).... (x++1, Vk+1) from Tr Given the model parameter the training objective is maximizing prediction accuracy of Yk+1 under the formatting of ICL: max LCu (y yk+1). where Lee is the cross-entropy loss, and Y641 is the in-context prediction defined in Eq. I We evaluate the meta-trained models on the 7 settings of MetalCL. For each setting, we test two cases: 1) all target tasks; 2) target tasks in the training unseen domains (OOD generalization). More details are in Appendix [A.4]
Baselines. Following MetalCL, we use GPT-2 Large (762M) (Radford et al. 2019) as base model, and also includes GPT-Neo 2.7B (Black et al.) (221)) and Pythia-2.8B (Biderman et al./ (023) (Appendix
B.2
For non-invariant methods, select AR ICL (Radford et al. 2211 and NoPE (Kazemnejad 2024). For invariant methods, we select Prefix ICL et al. 2020) and three types of BoE
al.
ICL
Appendix A.2): PCW (Ratner et al. 2022). SAICL (Cai et al., (023). and BatchICL (Zhngg 2024). We adopt context examples for training and evaluation.
al.
Results. As shown in Table 2 compared to non-invariant methods, InvICI outperforms in 4 out of 7 tasks in the "All target task" setting and all the tasks in the "Target tasks in unseen domains" setting. This indicates that permutation invariance is indeed crucial property for ICL algorithm, which incorporate the inductive bias of symmetry into the model architectures, resulting in an extraordinary improvement on performance, especially when generalizing to OOD tasks.
2 Although NoPE alone is invariant, it still utilizes an auto-regressive I.I which breaks the invariance.

---

Published as a conference paper at ICLR 2025
Table 2: The in-context learning performance with language models based on GPT-2 Large. We changed the causal mask and positional encoding implement different types of ICL models. The models are finetuned under the framework of MetalCL (Min et al. 2022b)
HR
NON-CLASS
NON-QA -QA
NON-NLI NNLI
METHOD
-CCASS
QA
PPARA
All forget nasks
Non-doverton
AR ICL, RRADFORD ET
2018
43.dan no
43.4
40.2
44.0
37 37.9.
30.3
41.9
NoPE
KAVIERNEJUET 2024
41.7gg
30.020.
40.3gam
44.5gn
36 06.62n..e
26.8
38.8
37.0 37.0gn
Invoriont
PCW (BoE) RAINER 2023
39.7
37.7.00.s
40.8.o.:2
37.7
40.Ta3.52
35.101.63
38.1-0.0
SAICE (BoF)
43.4gn as
43 43.2an
37.510.24
46-1 a0
40 2.3.0
33-921-4
BATCHICL MaMer
$1.729 zt
25. 25.4
27.120022
2an.:2
35.9g
30.7
(RAIEEL
40.3:.0.0 o.
69.6.
10.1
43.6s0.12
3G.8 30.8.4.35
45.4.3.65
34.9.2.0
INVIC1/OURS
15.3
42.9
28.44
45.3.
38.3
51.6.1..4
34.701.50
42.4
Tirger
domains
Non
AR
(RADFORR
31.6
35.7.0.50
21.1
66-5an se
30 2
80.3g s0
34.1g
43.6 6s
NoPE NAAZEANEAAUE 2024
32.9g1 iz
23. 23.4g0.
26.921.44
63.62.1.7
38, 2.st.4
33.2 33.22...0
32.0gb
35.8.0 it
Invoriont
PCW (BoE) (RATNER FT
2022
35.6
31.3
26.9 26.921.39
65.3
33 13.77
66.7g3.u
34.gn
42.0g1
SAICL (BoET
30.7
29.7
26.441441
56.210.50
41 41.Dass.0
64.312221
37.1e1.0
40.8
BATCHICI (RoF) CHANIi
2024
24.2gom
12.320
23.0g001
$1.9 31.9g:
29.4 4.
37 ag..
36.821.
20.3.20
PRERX ICL (RAFELE)
31.0a2.1
38.0s1.53
29.6
63.8
36.
52 52.0a2.54
34,0=0az
40.1=3.11
INvICL/OUst)
44.4
38.8
67.6.2.22
42.6a1.4
81.8.0as
37.5
48.4
46
Table 3: The inference time of different models.
Accuracy
-2
80
Method
Inference time (ms)
28
AR ICL
21.9
PCW (BoE ICL)
21.7
36
InvICL
Prefix ICL
22.0
Number of Demonstrations
2
1b
InvICL
22.0
Figure 4: The length generalization behavior of InvICL and AR ICL on HR-LR setting. The models are meta- trained by sequences with 8 context examples.
Compared to invariant methods, InvICL outperforms 5 out of 7 tasks in the "All target task" setting and 6 out of 7 tasks in the "Target tasks in unseen domains" setting. Although being permutation invariant, these baselines exhibit poor performance (none of them surpasses AR ICL on average). This highlights the crucial properties of information non-leakage and context interdependence implemented by InvICI.
Length Generalization. The generalization ability to different input lengths is crucial property of the language model. In the context of ICL, the ability to adapt to varying numbers context examples can be perceived as dimension of its length generalization capability. However, in the main experiments, the number of context examples remains consistent throughout both the training and evaluation phases. Hence, we vary the number context examples, as illustrated in Figure 4 We observe that InvICL is much more robust than AR ICL when the length of the test data differs from that of the training data, indicating its strong capability for length generalization.
we claim that our parallel implementation of InvICL has the same computational complexity order as full self-attention and AR self-attention. In Table 3 we empirically verify this by evaluating the inference time of different ICL models, showing that InvICL enjoys roughly the same inference speed as other models. Besides, question worth considering is the memory cost of InvICL since it duplicates the input sequence. We find that when the inputs size of the GPT-2 Large model increases from 512 to 1024, the GPU memory overhead increases by 14% (from 4.2 GB to 4.8GB). We consider this acceptable given the clear improvements in performance.
Computational Cost. In Section 3.3
Ablation Study. In Table 4 we conduct ablation study demonstrate the effect of the two components of InvICI: the invariant mask and the symmetric positional encoding. The experiments show that either component is important for invariant ICL. Additionally, in Appendix B.3 we demonstrate that the effectiveness of InvICI. not due to its doubled input.
8

---

Published as a conference paper at ICLR 2025
Table 4: Ablation study of invariant mask and symmetric positional encodings (PE) ICL perfor- mance and order sensitivity.
METHOD
HR-*L (1)
SENSITIVITY (L)
AR ICL
43.4
0.25
+SYM PE
38.4-5.0
0.30
+NV MASK + BOTH (INVICL)
44.8
0.10-0.15
45.1
0.00
Permutation Invariance. In Table we demonstrate the permutation invariance of InvICL. Fol- lowing Chen et al. (2022), we measure the order sensitivity as the frequency that the prediction is changed under random permutation. We observe that both the invariant mask and PE are important for achieving invariance, and lower sensitivity indicates better performance.
5 DISCUSSION
The Mechanism behind InvICL's Strong Length Generalization Ability. We consider that the mechanism primarily stems from InvICL achieving invariance. As mentioned in the introduction, previous studies have found that respecting data symmetry in models helps improve generalization. For example, (Sokolić al. (2016) demonstrated that when the input data exhibits invariance under certain transformations (such as rotation or translation). utilizing an invariant classifier can achieve lower generalization eπror compared to regular classifier. Bietti et al. (2021), Tahmasebi & Jegelka (2023) concluded that encoding invariances into model improves the effective number of samples, thereby enhance generalization ability. These theoretical results could help explain why InvICL demonstrates stronger length generalization ability.
Theoretical Complexity of InvICI. Suppose there are n context examples and test example (considering the examples attention units), and let M E ( (n+1) be the attention mask defined in Figure 2(d) The complexity of InvICL determined by the number of "0" elements in M. The attention computation for InvICL includes: 1) Independent self-encoding of the first-time input (corresponding to M[:n, :n), which requires n self-attention calculations: 2) LOO pre-encoding (corresponding to M[n: 2n, :2n)), which requires n² calculations; 3) Aggregation to the test example (corresponding M[2n+1, 2n+ll), which requires o-+ calculations. In total, there n 2n attention calculations, which is of the same order as Prefix ICL (π12 + 1) and twice that of AR ICL (n2/2 + 3n/2 2 1).
The ICL Training Objective. In the synthetic experiments, utilize the ICL objective to train the Transformers, which does not align with how LLMs are pre-trained. However, our paper focuses on improving the ICL capability of LLMs, rather than investigating the reasons behind the emergence train the model using the ICL. objective to demonstrate that InvICL
of ICI. ability. Therefore, can achieve stronger ICL capability compared to traditional AR ICL. This is also aligned with the objective we use in the real-world experiments.
Theoretical Understanding InvICI from Optimization Perspective. Previous studies have established the duality between ICL and the gradient descent algorithm, demonstrating that under specific parameterizations, ICL can implicitly implement gradient descent. In Appendix C we build upon this line of research and prove that InvICL, under the same parameterizations, can also approximately perform gradient descent, thereby highlighting the theoretical potential of InvICI.
CONCLUSION
In this paper, by distilling the advantages of regresregeessve language models, we identified two additional desiderata for invariant ICL: information non-leakage and context interdependence. Since existing invariant ICL algorithms cannot achieve these desiderata simultaneously, we proposed novel invariant ICL scheme called Invariant In-context Learning (InvICI.) which accomplishes these goals concurently. We also proposed an efficient parallel implementation of InvICL. Empirically, we show that InvICI. outperforms both invariant and non-invariant ICI. methods on most tasks, and demonstrates good length generalization abilities. These results sparked the unique ad advantages of the principled design of invariant ICI.

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
Yisen Wang was supported by National Key R&D Program of China (20222D0103300) National Nat- ural Science Foundation of China (92370129, 62376010), and Beijing Nova Program (20230484344, 20240484642). Yifei Wang was supported part by the NSF Institute TILOS, and an Alexander von Humboldt Professorship
REFERENCES
Sweta Agrawal, Chunting Zhou, Mike Lewis, Luke Zettlemoyer, and Marjan Ghazvininejad. In- context examples selection for machine translation. arXiv preprint wrxiv:22/2.024331 2022.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv: 2306,00297, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. ICIRR 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.0463) 2023.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle OrBri,, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In ICML, 2023.
Alberto Bietti, Luca Venturi, and Joan Bruna. On the sample complexity of learning under geometric stability. In NeurlPS, 2021.
Sid Black, Gao Leo, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, March 2021. URL https://dol
3rg/10.5281. zenodo. www.eeennnnd...... 5297715
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
Tianle Cai, Kaixuan Huang, Jason Lee, and Mengdi Wang. Scaling in-context demonstrations with structured attention. In ICML 2023 Workshop on Efficient Systems for Foundation Models, 2023.
Yanda Chen, Chen Zhao, Zhou Yu, Kathleen McKeown, and He He. On the relation between sensitivity and accuracy in-context learning. arXiv preprint arXiv:2209.00 077666 2022.
Yongqiang Chen, Binghui Xie, Kaiwen Zhou, Bo Han, Yatao Bian, and James Cheng. Positional information matters for invariant in-context learning: A study of simple function classes. arXiv preprint ar arXiv:2311.18194, 2023.
Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in- context? language models secretly perform gradient descent meta optimizers. arXiv preprint arXiv:2212. 10559, 2022.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning, arXiv preprint wrXiv:2388.06112 2023.
Deqing Fu, Tianqi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. In NeurIPS 2023 Workshop Mathematics of Modern Machine Learning, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. In NeurlPS, 2022.
10

---

Published as a conference paper at ICLR 2025
Srinivasan Iyer, Xi Victoria Lin, Ramakanth Pasumuru, Todor Mihaylov, Daniel Simig, Ping Yu, Kurt Shuster, Tianlu Wang, Qing Liu, Pumit Singh Koura, et al. Opt-iml: Scaling language model instruction meta learning through the lens of generalization. arXiv preprint arXiv:2212.12017, 2022.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy, The impact of positional encoding on length generalization transformers. In NeurIPS, 2024.
Jannik Kossen, Yarin Gal, and Tom Rainforth. In-context learning learns label relationships but not conventional learning. In ICLR, 2024.
Jiachang Liu, Dinghan Shen, Yizhe Zhang, William Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 100-114, 2022.
Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In ACL, 2022.
Sewon Min, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Noisy channel language model prompting for few-shot text classification. In ACL, 2022a.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannanch Hajishirzi. Metaicl: Learning to learn in context. NAACI., 2022b.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(5):548--5511 2020.
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin dcyton-Brown, and Yoav Shoham. Parallel context windows improve in-context learning of large language models. arXiv preprint arXiv:2212. /0947, 2022.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent? arXiv preprint arXiv:2310.08540, 2023.
Jure Sokolić, R. Giryes, G. Sapiro, and M. Rodrigues. Generalization error of invariant classifiers. In AISTATS, 2016.
Behrooz Tahmasebi and Stefanie Jegelka. The exact sample complexity gain from invariances for kernel regression. In NeurIPS, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurlPS, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICML, 2023.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.0588 05858, 2023.
Qixun Wang, Yifei Wang, Yisen Wang, and Xianghua Ying. Can in-context learning really generalize to out-of-distribution tasks? In ICLR, 2025.
Yifei Wang, Yuyang Wu, Zeming Wei, Stefanie Jegelka, and Yisen Wang. A theoretical understanding of self-correction through in-context alignment. In NeurIPS, 2024.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe V Le.. Finetuned language models are zero-shot learners. In ICLR, 2021.
11

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, al. Chain-of-thoughh prompting elicits reasoning in large language models. In NeurIPS, 2022.
Zhiyong Wu, Yaoxiang Wang. Jiacheng Ye, and Lingpeng Kong. Self-adaptive in-context learning. arXiv preprint arXiv:2212.10375, 2022.
Yanzheng Xiang, Hanqi Yan, Lin Gui, and Yulan He. Addressing order sensitivity of in-context demonstration examples in causal language models, arXiv preprint arXiv:2402.15637, 2024.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In ICLR, 2021.
Kaiyi Zhang, Ang Lv, Yuhan Chen, Hansen Ha, Tao Xu, and Rui Yan. Batch-icl: Effective, efficient, and order-agnostic in-context learning. arXiv preprint wrXiv:2401.0646) 2024.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In ICML, 2021.
Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and E. Xing. Dags with tears: Continuous optimization for structure learning. In NeurlPS, 2018.
12

---

Published as a conference paper at ICLR 2025
A IMPLEMENTATION DETAILS
A.1 SYMMETRIC POSITIONAL ENCODING
In this paper, we mainly focus on the absolute positional encoding which is used in the GPT family, As shown in Figure 5] we adopt an independent position encoding scheme that treats each example independent sequence, which follows the design in (Ratner et al. (022). For each context example X,, we always allocate the positional encoding as it starts from the first position. Denote the maximal sequence length among x; as mmax- For the test example xt, we assign its positional encodings starting from the index max This implementation applicable to BoE ICL, Prefix ICL, and InvICL.
X
X
X]
X
X
Tolansing
Toboes
Cetuuding
Cedentinn
Tolne
Embending
Tox
Drossdiing
(a) Symmetric PE for standard input
(b) Symmetric PE for the duplicated input of InvICL
Figure 5: The symmetric positional encoding applied in our work. p: refers to the learned absolute positional embeddings that are added the token embeddings 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).
A.2 BAG-OF-EXAMPLES ICL
We introduce the implementation detail of two BoE ICL methods mentioned in the main text, PCW (Ratner et 2022), SAICL (Cai et al. 2023) and BatchICL (Zhang al. 2024).
PCW (Parallel Context Window), PCW is a work originally aimed at improving the acceptable length of language models. Denote N be the maximal length of a language model, and n > N be the input length. PCW divides the input into context windows with length NN, and separately puts them method (similar to Figure 2(c) where each block
into the LM. Finally, it utilizes a "bag-of-windons in the mask refers to context window) generate the predictions. We note that by considering each context example as a window PCW, can implement the Bag-of-Examplee ICL algorithm.
SAICL (Structured Attention for ICL). SAICL is a method proposed to improve the inference efficiency and order-sensitivity of in-context learning. Similar to PCW, they also encode the context examples independently but are also aware of the test example. The original method is based on T5 (Raffel et al. (020). language model with the ancoder-decodee architecture. We transfer its design to the GPT Tamily by directly taking its attention mask and use the symmetrie PE proposed above.
BatchICL. Instead of conducting N-shot encoding for all context examples, BatchICI. utilizes N separate 1-shot encodings for each context example. It then aggregates the intermediate hidden states of the respective last token, which is subsequently incorporated into the forward computation of zero-shot query to generate the final prediction. We basically follows all the setting introduced in the original paper. As for the layer to extract the aggregated vector, we simply takes the 15-th layer, since they found that any intermediate or later layer is a fair choice.
A.3 SETTING OF THE EXPERIMENTS ON LINEAR REGRESSION TASKS.
T + b) as the linear function class. Let Dg be a distribution on G. and Dx be a distribution on X. In the training phase, iteratively sample a random function g E from De, and sample i.i.d. XI,..,, 1k...,14411 from Dx. Then, we produce prompt in the ICL manner P = (x1))(x1),... Xk. g(xx), x(+x) 1). and train a model 0 to output (x,x)))-., (xx))(((x++))
Denote G = (g X E Rd + gx)
(P) (as equation Eq. (I)) where ((x,) is the prediction for g(x)).
13

---

Published as a conference paper at ICLR 2025
The training objective is
min ED
((((x)) 9(x,))
(10)
k
1=0
where l is the MSE loss. In the experiments in Section and Dg w ~ ((0 Ia), 6 0.
4.1
we set d 20, A: 40, Dx -N Ia).
The architecture selection follows (Garg et al. 2022), where 12-layer GPT-like Transformer decoder is utilized. We implement the four model types by using the symmetric attention mask and PE.
A.4 IMPLEMENTATION DETAILS OF EXPERIMENTS ON REAL- HEAL-WORLD DATA.
Evaluation. Following MetalCL Min et al. 2022b), we consider evaluation settings: I) HR-LR, which means training with high resouree data and testing on low resource data; 2) X-x (X=(Clasiffication QA) )). which means training and testing on the same task type, but with no overlap in tasks; 3) Non-X-X (X=(Classffication, QA, NLI, Paraphrase, J, which means training and testing on different task type. The last settings the most challenging, which require strong generalization capacities of language models (Min et al.) 2022b). For each setting, we make evaluations both on all target tasks and subset that contains target tasks in the unseen domains of the source tasks, e.g.. medical, financial, and climate. This setting also challenges the out-t--diitributtin generalization capability of models.
Truncation.
Since MetalCI. (Min et al. 2022b) truncates the training sequence when it exceeds the maximum input length of the LM, and the ICL prompt sequence is duplicated in our implementation of InvICI. the training sequences differ between InvICL and other methods because of different truncate rates. As shown in Table [] there is significant gap in the dataset size between the standard input and the duplicated input under the truncating setting. make the comparison fair, we apply the same truncate rate in InvICL to the normal training sequence so that all the methods share the same training Additionally, reduce the number of context examples in the training phase from 16 to to control the truncate rate of InvICL to the same level as standard ICL.
Table 5: Ratio the remaining data between different input types under the truncating setting of MetalCL (Min et al. 2022b). Here the number of context examples is set to 8.
HR L.R
CLASS CCLSS
NON-CLASS CLASS
QA
NON-QA +QA
NON-NLI -NLI
NON-PARA
INPUT TYPE
QQ
PPARA
Remaining ratio of training dataset
STANDARD
70%
90%
71%
59%
80%
85%
85%
DUPLICATED 53%
79%
55%
40%
62%
75%
71%
Direct & Channel. Besides the standard ICL paradigm, MetalCI. (Min et al. 2022h) adopts new inference paradigm called noisy channel ("Channel") (Min et al. 2022a) and achieves better experimental performance. Contrary to the standard ICI. paradigm (also called "Direct" in (Min et al. (022b)) that takes (1........ Xn, Ea.yn, Yn: Xx) as input, the Channel paradigm takes
X1, Yn. X.. Ye) and try generate x,. Note that in order to generate the prediction, Channel ICL needs to perform n forward passes conditioned on each of the n labels y, and regard the label with minimum perplexity as the prediction. This will, on the one hand, increase the computational complexity and, on the other hand, reduce its applicability the generative tasks where the label space is large, e.g., Question Answering. Therefore, we adopt the "Direct" setting in our experiments, i.., the standard ICI. paradigm.
B ADDITIONAL EXPERIMENTAL RESULTS
B. SYNTHETIC EXPERIMENTS ON OTHER SETUPS
In this section, we conduct additional synthetic experiments on more functions and out-of-diittibuutin setups, further showcase the generalization capability of InvICI.
Other function settings. We consider two other function settings proposed by (Garg et al. 2022 sparse linear regression and decision tree, to illustrate the ability of InvICL to learn algorithms to solve other tasks. Results are given in Figure
14

---

Published as a conference paper at ICLR 2025
1. Sparse linear regression. In this task, random linear function y w is sampled be predicted. yet the efficient has only 3 Although it is also a linear regression task, its optimal algorithm is no longer least squares but Lasso, which involves solving the least squares objective with an 11-norm regularizer for the weight vector. This demands the in-context learners to learn an algorithm different from that in linear regression to solve this task. Following the experimental settings in our paper, we test the performance of AR ICL and InvICL which are trained with 200k epochs. We can still observe the consistent results of our paper that InvICL possesses fast convergence (InvICL converges while AR ICL does not).
coordinates out 20 dimensions.
2. Decision tree. We follow the setting in Garg et al. 2022 where the class of depth decision trees with 20-dimensional inputs is considered. We evaluate the performance of AR ICL and InvICL that are trained with 200k epochs. We find that although AR ICL performs better than InvICL for short inputs, as the length of the input sequence increases, InvICL gradually outperforms AR ICL, indicating the strong extrapolation ability of InvICL.
1.2
Inw
InVICL
squares ertor
0.8
Lassa (alpha-0.13
0.8
0.6
0.4
0.4
0.2
0.0
0.0
0
50
100
150
200
0
50
100
150
200
in-context examples
in-cantext examples
(a) Sparse linear regression
(b) Desicion tree
Figure 6: ICL performance on sparse linear regression and decision tree.
We consider three 2022 Chen al.] ) showcase the generalization capability of InvICL out-of-distribttion (OOD) tasks. We consider a distribution shift between the training and test datasets. The training data remain consistent with Section A.3 However, for the test data, we apply the following modification:
onttof-distribution Setups.
setups proposed by (Garg et al.
1. Add random noise to the labels by altering 0 to ~((, 1).
2. Scale the data sampling by altering Dx = NN to Dx = 32)
3. Sample the data 2y from random 10-dimensional subspace (out of 20 dimensions).
we report the testing MSE loss with the models trained for respectively 50k and 200k epochs. We omit Prefix ICI. and BoE ICL for their poor performance. We find that InvICI continues the advantages mentioned earlier, i.e., the fast convergence and strong extrapolation ability, indicating its strong capacity on OOD tasks.
In Figure 7
B.2 REAL-WORLD EXPERIMENTS BASED ON GPT-NEO AND PYTHIA
We also conduct experiments with models based on GPT-Neo 2.7B Black et al. 2.8B (Biderman et al. 2023) with other result is similar to what is demonstrated in the main text: InvICL outperforms the baseline in most of the tasks and especially performs well in the OOD settings. This indicates the applicability of InvICL to different base models.
2021 and Pythia-
unchanged, as shown in Table 6 and 7 The
Besides, we note that the three LLMs (GPT-2, GPT-Neo and Pythia) studied in our work utilize three different kinds of PE trainable PE, Alibi and Rotary PE, respectively. Therefore, our design of symmetric PE is applicable to wide range of PEs.
15

---

Published as a conference paper at ICLR 2025
1.50
1.50
InviCI
InviCL
1.25
1.25
Sqqaares
squared ener
3.00
Least Squares
0.75
0.75
0.0
0.50
0.00
0.00
In-context examples
In-context examples
(a) Random noise, 50k epochs.
(b) Random noise, 200k epochs.
12.5
InvICL
12.5
InvICL
10.0
Squares
10.0
Scuares
7.5
7.5
5.0
5.0
2.5
2.5
0.0
0.0
150
150
200
in-context examples
In-context examples
(c) Scaling, 50k epochs.
(d) Scaling, 200k epochs.
2.2
PWCC
12
InvILL
Least Squarees
sQuared errer
Least Squaress
0.0
0.8
0.4
0.4
0.0
0.0
50
100
150
200
in context examples
In-context examples
(c) Half subspace, 50k epochs.
(f) Half subspace, 200k epochs.
Figure 7: ICL performance on OOD tasks. The training dataset remains consistent with Section 4.1 but change the distribution of the test dataset. Random noise: changing the distribution of the linear bias from b = 0 to ~ NO 1). Sealing: changing the sampling distribution of x; from Dx (O0, Ia) Dx (0, T2). Half subspace: Sample the data 2; from random 10-dimensional subspace (out of 20 dimensions).
Table 6: The in-context learning performance on GPT-Neo 2.7B.
CLASS CLASS
NON-CLASS CLASS
NON-QA +QA
NON-NLI -NII
NON-PARA PPARA
METHOD
HR -
-QA
AVG.
All target tasks 40.1
AUTO-REGRESSIVE ICI.
45.8
41.2
46.4
36.8
45.2
33.1
41.2
INVICL(OURS)
46.1
40.2
40.2
48.6
35.8
44.7
33.7
41.3
Target tasks domains
AUTO- IITO-REGRESSIVE ICI.
39.1
33.1
31.8
66.5
34.7
56.7
33.1
42.1
INVICL(OURS)
39.6
33.9
32.7
68.1
31.4
56.9
36.0
42.7
B.3 ABLATION STUDY FOR INVICL
In this section, we conduct experiments to test the baselines (AR ICI. PCW, Prefix ICL. using the same duplicated data as InvICL. As shown in Table InvICL still outperforms the baselines when they given the doubled input as InvICI does.
16

---

Published as a conference paper at ICLR 2025
Table 7: The in-context learning performance Pythia-2.8B.
CLASS -CLAS
NON-CLASS -CLASS
QA -AQ
NON-QA -QA
NON-NLI -NNLI
NON-PARA -AAAA
METHOD
AvG.
larget
AUTO-RECRESSIVE ICL
31.3
22.3
27.8
33.4
33.7
29.7
37.6
30.8
INVICL(OURS)
31.5
26.3
28.5
33.0
35.6
28.0
40.2
31.9
Target tasks wnseen domains
AUTO-REGRESSIVE ICL
20.8
21.0
21.0
43.5
39.7
33.5
34.2
30.5
INVICL(OURS)
20.9
24.2
21.1
44.6
43.7
33.5
38.6
32.4
Table 8: Ablation study of using doubling input for the baseline methods. We report the result HR-LR. InvICL still outperforms the baselines.
METHOD
DOUBLED INPUT
ORIGINAL INPUT
AR ICL
43.8
43.4
PCW (BOE ICL)
40.6
39.7
PREFIX ICL
41.7
40.3
INVICL
45.1
B.4 DETAILED RESULTS FOR SYNTHETIC EXPERIMENTS
In this section, we provide detailed results for the synthetic experiments in section 4.1 In figure
g we demonstrate the еттоr curves of AR ICL and InvICL at different training epochs. In figure different training epochs when the number of context examples is 100. Both experiments demonstrate that InvICL's OOD in-context performance (length > 40) consistently outperforms AR ICL across all epochs. Specifically, as shown in figure 9 in the early stages of training, the error of InvICI. decreases rapidly, while the error of AR ICL only shows significant reduction after approximately 100k epochs. Furthermore, after 200k epochs, the error of InvICL stabilizes, whereas the error of AR ICL increases.
9 we present the
B.5 LINEAR PROBING EXPERIMENTS
In this section, conduct linear probing experiments based the synthetic setting, further explore how the architecture of InvICL impacts the model's internal representations. For a pre-trained model on the synthetic linear regression dataset, we freeze the model parameters and trained single linear layer on the hidden states the 3rd, 6th, 9th, and 12th layers, respectively.
As shown in Figure 10, the linear probing error of InvICL is consistent and close to the error curve of the pre-trained model across all tested layers. In contrast, for AR ICI., only the error curve of layer 12 converges to that of the pre-trained model. This indicates that InvICL encodes task features in the model much faster than AR ICI.. We believe this is closely related to its context interdependence property, which allows it to utilize richer context example information for encoding.
C THEORETICAL UNDERSTANDING INVICL FROM AN OPTIMIZATION PERSPECTIVE
In this section, we further characterize the advantages of InvICL from an optimization perspective.
InvICL Can Approximately Implement Gradient Descent. Consider linear regression task with instances (X, y), where X consists of row vectors x, E IRd and consists of labels yı E R, E [m]. The goal is to find the optimal weight vector w that minimizes the LSE loss L(w) = |[ww 2 standard gradient descent (GD) algorithm updates the weights iteratively as follows:
W = W x/X (Xwi-1-9 y),
(11)
where stands for the iteration step, and 77 is the step size.
Consider the ICL-style model input, formulated Z = (Z) Zn, Z1. Zn Zt),
where zj Xj .J E [n]] are the context examples, and z; = an arbitrary test example. Here we
8j
17

---

Published as a conference paper at ICLR 2025
1.2
Squared error
0.8
83k
50k
10ck
1004
0.6
125k
150k
1204
0.4
0.4
300K
300k
0.2
400k
0.2
400k
500k
500K
0.0
0.0
50
100
150
200
250
in-context examples
in-context examples
(a) InvICL
(b) AR ICL.
Figure Intermediate results for InvICL and AR ICL on the linear regression setting. The line colors represent the models trained with different epochs.
1.2
1.0
0.4
022
0.0
500
Training epoch (K)
Figure 9: The squared error different training epochs. We set the number of context examples to 100.
duplicate the input as required by InvICL and expect the model to predict
X₁
at the last token.
The theorem below illustrates the evolution of the prediction through the Transformer layer of InvICL. Theorem C.1. With the attention weight matrices configured as in (Von Oswald et al. 2023 iee.
W = Wg ( Idxd 0 w.
0dxd Wo
)
0
P
(12)
InvICL implements the following weight updating rule: at the l-th layer of the Transformer, the last (E)
X₁
token outputs zi
where
X, W:
W. W T (Xwi-1
-) - A
(13)
Here, Aw; = X X - diag XX (XW y).
Theorem C. shows that under specific parametrization. the weight updating rule implemented by InvICL (Eq. (13) is very similar to that standard GD (Eq. (TI), differing only by a second-orde -oooed term. For gradient descent to converge, the step size should be at most the inverse of the largest eigenvalue of X T Under this condition, the term 1/2 Pwww1 is dominated by 1 (Xw,-1 - y). ensuring that InvICL has the potential to closely approximates the standard GD algorithm in this linear regression task.
Discussion to Other ICL Methods. We provide comprehensive comparison of all the ICL methods considered in this paper from the optimization perspective: under the parametrization as in Eq. (12) 1) AR ICI emulates the online GD algorithm (with a constant learning rate) (Ding et al.) (023). which is not guaranteed to converge; 2) Prefix ICL implicitly implements the standard GD algorithm under specific set of parameters for attention (Von Oswald et al. 2023 Ding al. 2023): and 3) BoF ICL.
18

---

Published as a conference paper at ICLR 2025
1.2
Squared error
0.8
Layer 3
Layer
0.6
0.4
0.4
0.2
0.2
0.0
0.0
50
100
150
200
250
in-context examples
in-context examples
(a) InvICL.
(b) AR ICL.
Figure 10: The linear probing results on InvICL and AR ICL.
can only update the weight vector of the test (last) example (not the context examples) without context interdependence. This leads constant gradient computed at the initial point, causing fail to converge (detailed discussion is in Appendix D.T). Compared with these ICL algorithms, InvICL has several distinct advantages: 1) InvICL surpasses AR ICL in terms of convergence to optimal solutions; 2) Similar to Prefix ICL, InvICL approximately implements the standard GD algorithm while avoiding information leakage; and 3) Unlike BoE ICL, InvICL effectively incorporates context interdependence, allowing emulate more efficient GD algorithm. These advantages underscore the theoretical superiority of InvICL, which synergizes information non-leakage and context interdependence within an invariant ICL framework.
Practicality of Theorem C.1 Theorem C.1 is existence proof which illustrate that the Transform- ers have the potential to implement complex optimization mechanisms like gradient descent. In fact, The actual weight may not be strictly follow its parametrization. However, empirical studies including Von Oswald al./ (2023): von Oswald al. (2023), have shown that pre-trained Transformers exhibit behaviors akin to gradient descent in certain scenarios, thereby providing empirical evidence for the theory.
D PROOFS
D.1 PROOF OF THEOREM C.I
Proof. We mainly adopt the setting of (Von Oswald al., 2023) and (Ding et al. 2023
).
Let Z xj
(Z1,..., 22 (2n+1)
R (24+1)x(2++1) be the duplicated input of InvICL, where 2z
xx
Rd. y, R. and Znti for [m]. Consider the linear self self-attention layer in the scheme of InvICI. Given the query, key, value matrix W, 9, We, W, E [dd++))((+1) and the projection matrix
P e ((4+))((++1) the updating rule of the layer follows:
2j < 2j + PW 2. (z, W Wgzz),
Znzj 2ntj PW. Σ (z W W Wyzn+z),
(14)
Z2n+1 &- Z2n+1 + PW.
Σ Znti
W W W22z++1),
where j [n].. Following the setting of (Von Oswald et al. 2023) and DDing et al. 2023). we let
W W = ( Laxa 0
Oaxd (0)
) P nTI.
0
W =
-1
(15)
19

---

Published as a conference paper at ICLR 2025
Now, hope to what kind of iterative algorithm naturally be implemented by InvICL. Before that, we first give the L2 loss after doing one step of gradient descent
X(w 7) (Xw y)) - yll
Xw - xXX (Xw y)
(16)
" II nx XXW ()(X--y)ll -y)l??
To compare InvICL with the conventional attention heads for ICL linear regression, here investigate the convergence properties of the leave-one-out scheme Eq. (17) viewed as an optimization algorithm for solving the regression problem, and compare it to that of gradient descent. It turns out that if we use the same weighting strategy as (Von Oswald al. (023) but with InvICL, then we
obtain a similar iterative algorithm for in-context linear regression according to which the last row of Z evolves, but the update rule transforms into
Wi = We-1 1/X (Wi-1-yy)
(17)
where
-XX (Xw-1- 7|} (x, We-1 3/b y/)]
(18)
the label updated by the leave-one-out scheme. This equation is obtained by first perform gradient descent step w.r.t. the whole dataset with gradient update n) (Xw y) and then minus the term w.r.t the i-th data point (x w y:).
Expanding Eq.
we get that the global update becomes
We W,-1 x (Xw;-1 -yy)
= We-1- nx (Xwe-1-3) - nXX (XW--1-9) -y)
-nix, (x] We-1 - -1i)
(19)
-Y) + + (Xw--1-))
W
~//x T Diag(XXX (X
This delivers Eq. 13
Remark. BoE ICL, since the context examples cannot interact with each other, the GD algorithm implemented by it can only update the weight vector w of the test (last) example, but not the context examples. Particularly, this means the gradient update process w, W g (Wo, {xi, Wi }), where is the update function of BoE ICL. This means that the gradients are always computed at the initial point of the algorithm, thus the algorithm cannot converge.
D.2 PROOF OF PROPOSITION 3.4
Proof. We will first demonstrate that the attention score matrix A needs to adhere to a specific form when constrained by the attention mask M, in order to guarantee the permutation equivariance of embeddings the context examples. Subsequently, we will establish that this requirement equivalent to the permutation invariance of the ICL prediction with respect to the context examples.
Lemma D.I. Given an input matrix H (hi, ..., hn) E Brxxd and its attention score matrix A E Rn defined in Eq. (7. Denote SA(H) AHW, , be the self reff-attention operation, where A defined in Eq. Then, SA(H) is permutation equivariant to (h:} isf the atlention mask M is equal to
00
x
0
x
or O.
0
-OC
Proof. Denote T E Raxx a permutation matrix on the row vectors of H. This implies that T E (0, 1) jnx and I.T IT. Tl, 1,, Then the permutation equivariant condition can be
20

---

Published as a conference paper at ICLR 2025
stated TSA(H) = SA(TH). Since SA(H) softmax (HW (HW+)) M) HW.P, the condition can be expanded as
T softmax
+ M) HW.P
(20)
softmax (THW, W H T + M) THW.P.
It can be easily verified that 1) the permutation and softmax operations commutative, and 2) T is orthogonal. Therefore, the above equation can be transformed to
softmax (THW, (HW&) + TM) HW.P softmax (HWW, W. H + MT) HW,P. P.
(21)
This is equivalent to
TMT-1=1 M
(22)
for arbitrary permutation matrix T. Next, we will discuss what kind of matrix M satisfies this condition. For notation simplicity, we denote Jii,j) as the permutation performed only between the i-th and j-th index.
Assume M.,- o1- Taking T T(i.j). from Eq. (22) we have M,.j c1- By iterating over j. we have M3 = C1 for every [n].
Assume M. C2 i i Taking T T(t, k), * j. from Eq. 22 we have My.j = C2;
taking T = Tú, k), k o i, we have M,,, = Hence, by iterative applying permutations in this way, we can conclude that M3.1 c2 for every
In conclusion, In + c2(1nxn L.). Since the elements an attention mask can only take the value of either 0 or -o, M can only be one of the three forms demonstrated in Lemma D.I all attention mask is meaningless).
(an
Now we prove the equivalence between the desired permutation invariance property and the equivari- ance property discussed in Lemma D.I] As the permutation invariance property involves the ICL prediction, which relies the test embedding h, it is necessary incorporate it into the discussion. We denote the full input matrix of ICL as H (h,,... h h,) (((+1)xd, and the corresponding matrices in the self-attention process A, M.
Lemma D.2. Let the output embeddings of the Transformer be H' (hi,..., h/ h(). Then, h', .....hh, h²,, equivariant to the permutation of invariant to the permutation of (hi,..., h.,....,,,, iff (h',... (hj.... h.).
Proof. First, we need to extend existing results to the circumstance of the full input H. Consider the attention mask M E ((++1)x(x+1) of the full input, whose n x n submatrix at the upper-left is equal to M, ie., M M = M. From the condition in the Proposition we have that M = 0
Besides, it is evident that Proposition 3.5 also satisfies for M, we have M -00 Other
variables can naturally extended.
In Lemma D.I we have proved that the equivariance property equivalent to the attention mask M being one of three specific forms. Now prove the contrapositive statement of Lemma D.2
If (h/.... h(, is not equivariant the permutation of (h....., ), by Lemma D.I. the mask on context examples M must satisfy either 1) Ei * j, M2 F Mjj- or 2) 3i + j. k 7 M,j * Ma We separately demonstrate that these properties will break the property of permutation invariance. For the following circumstances, we uniformly let W W. W P I.41. Denote the
embedding of h, after k celf-attention layer h, (k) Then, the embeddings updated
h (k+1)
Σ
[s(b] (k) hna²)) Myylh,a
(23)
where s(-, (() is the similarity function calculated by their inner product and softmax normalization, which is defined in 2
21

---

Published as a conference paper at ICLR 2025
Bi * MM # M,j- Without loss of generality, since the elements of M only take the value
of either fty. we let M1, = M22 -o. Then we construct the input matrix as
h1 = e1, h2 = eg, h; = 0(i > 2), h, = 0, where e, denotes the i-th unit vector (i E [d).
Since M22 = M2
00, following Eq. 23), we find that h (1) c And since
M11 0, we have h| = c201 + C301
Now we permute the first and second examples, L.e. h1 e2, h2 e1- Although we find that the first update of the test embedding remains unchanged since Eq. (23) is permutation invariant for hy. the second update differs. Since we have h² er and h
cyej +C2e1, the aggregation h (2) changes. Therefore, the property of permutation invariance is broken.
Di * j,k # I, M, # MR. Without loss of generality, let M,j = 0, Ma = -00. We construct h, = e1; hy = e2, h = 0. Then, we have hn = C1 e1 + 0202, and h h(2)) = 0301 Similar to the above process, we can prove that h, is not permutation invariant
w.r.t. the index exchange (i.j) ++ (k;
In conclusion, any attention mask M that violates Lemma D.I will break the property of permutation invariance. Thus Lemma D.2 is proved.
Finally, by combining Lemmas D.I and D.2 we can deliver Proposition 3.4
D.3 PROOF OF PROPOSITION 3.5
Proof. Consider the case that G has no self-loops. Since G' is digraph with no cycles, it is a directed acyclic graph (DAG), According the graph theory (Cormen et al.) (0022). DAG can be topologically ordered, which means in this ordering, any vertex is not reachable from later vertices in the graph. Therefore, if we reorder the vertices in this way. we have A; = O for i S j. which infers that A is strictly lower diagonal. Since the original graph allows self-loop, which corresponds the diagonal elements, the adjacency matrix is lower triangular. This is equivalent to that the attention mask on context examples M is lower triangular.
E RELATED WORK
The order-sensitivity of ICL. The phenomenon that ICL is sensitive to the permutation of context examples has been observed in several works. (Zhao et al., (021) and (Lu al. 2022) used GPT-3 to perform in-context learning on classification tasks such as SST-2 and observe a high variance w.r.t. the permutation of the context examples. Besides, (Xie et al. 2021) and (Agrawal et al. 2022) found the same phenomenon on a generated synthetic dataset and machine learning tasks, respectively. Additionally, (Chen et al. 2022) empirically showed that the order-sensitivity is negatively correlated to the performance of ICL. To address this issue, (Zhao et al. 2021) proposed calibration module to make the output distribution consistent with prior knowled ledge. (Lu et al., 2022) filtered out the best prompt ordering by investigating their calibration on a generated set. (Xiang et al.] (024) utilizes contrastive learning align representations of in-context examples across different positions, resulting in the alleviation of order sensitivity. Besides, there are works that focuses on implementing the concept of permutation invariance from an architectural perspective. For example, SAICL (Cai et al.) 2023) proposed a structured attention mechanism that achieves permutation invariance. However, their work is based on improving the ICL performance of T5 (Raffel et al.) 2020). a language model based on encoder-decoder architecture, which did not address the order-sensitivity issue of regr-seeee LMs. BatchICL (Zhang al.) (024)) is the work that is most relevant to us. Instead of conducting N-shot encoding for all context examples, it utilizes N separate 1 1-shot encodings for each context example. It then aggregates the intermediate hidden states of the respective last token, which is subsequently incorporated into the forward computation of a zero-shot query to generate the final prediction. In this way, the model achieves permutation invariance since the encoding of the context examples are independent.
The connection between ICL and Gradient Descent. Early stage formal theoretical investigation of the linear regression in-context learners includes (Akyürek et al. [2022] Von Oswald et al. (003). They first showed that Transformers learn in context via gradient descent, where layer performs
22

---

Published as a conference paper at ICLR 2025
one gradient update. In subsequent work, (von Oswald et al. 2023) further argued that Transformers strongly biased towards learning to implement gradient based optimization routines. (Ahn et al. 2023) showed Transformers can learn to implement preconditioned Gradient Descent, where the preccondt
adapt to the data. (Bai et al. 2023) provided detailed constructions for how Transformers can implement range learning algorithms via gradient descent. (Dai et al.)
2022) conducted experiments on NLP tasks and concluded that Transformer based language models performing ICL behave similarly models fine-tuned via gradient descent; however, concurrent work argued that real-world LLMs do not perform ICL via gradient descent (Shen et al. (023). (Fu et al. (023) argued that Transformers actually learn to perform in-context learning by implementing a higher-order optimization method, not gradient descent. Predictions made by different Transformer layers match iterations of higher-order optimization methods better than they match iterations of gradient descent.
23


Reviews:
[{"review_id": "u1lW29rPZi", "reviewer": "Reviewer_k8s4", "summary": "This paper proposes a new permutation-invariant attention mechanism InvICL, which is adapted from Bag-of-Example (BoE) but admits full context-interdependnce. The authors enable the context interdependence by duplicating the context examples and allowing leave-one-out (LOO) attention from the duplicated tokens to original tokens. By doing so, IncICL enjoys the properties of (i) permutation-invariant, (ii) information non-leakage, (iii) context interdependence. Evaluations show that InvICL outperform baselines (Prefix, AR, BoE, NoPE) on tasks that adimit a permutation-invariant nature.", "questions": "1. Is the loss taken over only the test token or all tokens (autoregressive loss) in the input for Prefix ICL in the training stage? If it is the latter, then it does not make too much sense since for Prefix ICL the context example can attend to its label directly.\n\n2. Can the test token attend to itself in authors' implementation in Prefix ICL for linear regressions? From figure 2 it seems to be attending to itself, but I am not sure if that is the case in the experiments. Other papers [1, 2] using Prefix attention avoid such attention since it could provide incorrect signals to multilayer transformers (the corresponding $\\hat y$ is inaccurate within the forward passes).\n\n3. Does the positive encoding in the footnote of page 1 refer to positional encoding?\n\n\n[1] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, Max Vladymyrov. Transformers Learn In-Context by Gradient Descent. ICML 2023.\n\n[2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The idea is novel and interesting. The motivation is clear as InvICL is proposed by the desired three properties.", "weaknesses": "1. Experiment results need more analysis and interpretations. The authors find InvICL shows better length-generalization capabilities, whose mechanism is unclear to me. \n\n2. More experiment results needed. Most results do not have a  reported std. Besides, it would be benificial if there is a figure of squared error curves where the x-axis is the training epochs. Currently there are only results from 50k and 200k epochs.\n\n3. The results of Prefix ICL for linear regression are a bit weird to me. The reported squared error seems to fairly high (around 0.5) even when the number of examples is 40. In another paper [1], the reported squared errors are much lower than 0.1 for 5-layer transformers with number of context examples 10 and data dimension 10. \n\n[1] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, Max Vladymyrov. Transformers Learn In-Context by Gradient Descent. ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mb0qYirvDK", "reviewer": "Reviewer_p8nF", "summary": "This paper presents InvICL, a new approach to in-context learning that achieves three key properties: permutation invariance, information non-leakage, and context interdependence. The authors propose a parallel implementation using a duplicated input sequence and modified attention patterns. They provide a theoretical analysis showing that InvICL approximates standard gradient descent and demonstrates improved performance over baseline methods (GPT-2 Large 762M, GPT-Neo 2.7B, Pythia-2.8B). The work tries to understand and improve in-context learning, though there are some practical concerns.", "questions": "How would InvICL perform in few-shot settings without MetaICL finetuning? This would be particularly relevant for scenarios where finetuning isn't practical.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper identifies and formalizes three important properties for ICL that weren't previously unified. The authors demonstrate why these properties matter. \n- The theoretical analysis is simple but straightforward. The authors prove that InvICL approximates standard gradient descent (Theorem 4.1) and show how this leads to better convergence properties compared to other ICL variants.\n- The experimental results look interesting. The method shows strong performance across multiple settings - synthetic tasks (Figure 3), out-of-distribution scenarios (Figure 7), and real-world datasets (Table 2).", "weaknesses": "- The practical applicability of the method raises some concerns. The paper relies on MetaICL finetuning, which is computationally expensive for modern large language models. I wonder if there are any training-free methods. \n- The efficiency implications are concerning. Doubling the input sequence length (as shown in Figure 2d) increases memory usage. In Section 5.2, “We find that when the inputs size of the GPT-2 Large model increases from 512 to 1024, the GPU memory overhead increases by 14% (from 4.2 GB to 4.8GB)”. However, when the context is long, e.g., 64k, the overhead will be super large. Moreover, the attention pattern used likely makes it incompatible with FlashAttention optimizations, which could further impact practical deployment.\n- The experimental setup feels somewhat dated by using GPT-2 Large as the primary model. While the authors include some results with GPT-Neo 2.7B and Pythia-2.8B in the appendix, evaluating more recent models like Phi-3.5 or Llama 3.2 would better demonstrate the method's relevance to current architectures.\n- There are some inconsistencies in the presentation. Figure 1 shows results for PCW while Table 1 uses different terminology (Inv ICL), making it difficult to track the method comparisons. It confuses me to understand the relative performance of different approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KFDMwilwpq", "reviewer": "Reviewer_fSCp", "summary": "The paper presents a new method to achieve higher performance out of In-Context Learning with no additional model training. It presents conditions under which a well defined notion of invariant In-context learning should perform well highlighting the gap in prior works that this paper fills with its proposed method. The paper gives theoretical and empirical justification for the superiority of its method.", "questions": "- **Definition 3.3**: if j > i, and this condition still holds, then isn’t it sort of contradicting information non-leakage (not technically with your definition as it only requires non-leakage to the corresponding label; but in my opinion, if an x_i can influence another x_i or y_i later in the sequence, it should count as information leakage)? Should the condition be j < i instead of j \\neq i? This is intuitively the case for AR LLMs because they only depend on previous context and you have mentioned it in point 3 (“These examples provide the context for better encoding of x_i, which in turn improves the prediction of future examples when x_i serves as their context.”) Why is context interdependence important, or why should I believe that context interdependence should be useful/necessary for AR LLMs, when they can not attend to future demonstrations?\n- External modification of attention masks (through aggregation via BoE) is a deviation from how these models are trained to read those activations. How do you think this impacts the model’s internal representations? What could be the effects on predictions apart from pure performance numbers (like does it increase hallucination)? Would like to see some discussion or simple experiment along this line.\n\nPS: I hope to increase my rating if the authors make an effort to discuss my concerns.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper presents a clear step up in terms of theoretical ideas as well as empirical evidence of improvement compared to prior works attempting to optimize ICL.\n- Paper is very well written, clearly presents and distinguishes its contribution.\n- Although the method requires more computation in theory, the authors achieve parallelism and same order of computation as standard ICL with a smart trick. The authors also talk about the additional memory requirements. [Point being that the paper addresses a plethora of relevant points surrounding its method].", "weaknesses": "**The significance of section 4 and section 5.1 are unclear.** Please see this [ICML 2024 paper](https://arxiv.org/abs/2310.08540) that talks about how training transformers with ICL objective may be incompatible with real ICL in LLMs that do not train explicitly for ICL with fixed ICL prompt format.\n- Theorem 4.1 shows that if we put weight matrices in a particular format, we can simulate InvICL with transformers. But, is there reason to believe that trained transformers end up with similar weights?\n- Similarly, results from section 5.1 assume that transformers are trained with ICL objective, which is not true for LLMs. \n\nI believe these points merit a discussion in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C1aXx6iNUu", "reviewer": "Reviewer_LHUi", "summary": "This paper proposes InvICL, a new invariant ICL algorithm that satisfies not only invariance, but also information non-leakage and context interdependence. Then the authors provide both theoretical and empirical evidence to prove the effectiveness of the proposed algorithm on ICL.", "questions": "- Even with InvICL, is there a *learning plateau* during training?\n- If we actually train with InvICL, does the weight matrices aligns with Theorem 4.1?\n- Why is adding only symmetric positional embedding not beneficial?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is overall well written and well presented. \n- The paper present reasonable desiderata for ICL algorithms, and the proposed method complements prior works on those criteria.\n- The paper present some theoretical intuition on why the proposed method should work better than its counter parts.\n- Authors conduct various experiments to prove the effectiveness of the proposed algorithm.", "weaknesses": "- One might view the proposed method as concatenation of the prior works (PrefixICL and (variant of) PCW)\n- As with the previous works, it lacks connection with *actual* ICL performed by LLMs.\n- Although the inference time of InvICL does not differ much from AR ICL, is it also true for training? To my understanding, in InvICL (and other invariant ICL algorithms), $\\hat{g}(x_i)$ and $\\hat{g}(x_{i+1})$ cannot be computed on one forward pass. \n- For real-world dataset experiments, is AR ICL also fine-tuned on ICL tasks? If not, comparing ICL specific fine-tuned model and general LLM on ICL tasks seems unfair.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)