PEARL: Towards Permutation-Resilient LLMs
Authors: Liang CHEN
Li Shen
Yang Deng
Xiaoyan Zhao
Bin Liang
Kam-Fai Wong
Venue: ICLR 2025 Poster (2025)
Keywords: In-Context Learning
Large Language Models
Instruction Tuning
Robustness
Attack
Distributionally Robust Optimization
Optimal Transport
Sinkhorn Algorithm
Primary area: foundation or frontier models, including LLMs

Abstract: The in-context learning (ICL) capability of large language models (LLMs) enables them to perform challenging tasks using provided demonstrations. However, ICL is highly sensitive to the ordering of demonstrations, leading to instability in predictions. This paper shows that this vulnerability can be exploited to design a natural attack—difficult for model providers to detect—that achieves nearly 80% success rate on LLaMA-3 by simply permuting the demonstrations. Existing mitigation methods primarily rely on post-processing and fail to enhance the model's inherent robustness to input permutations, raising concerns about safety and reliability of LLMs. To address this issue, we propose Permutation-resilient learning (PEARL), a novel framework based on distributionally robust optimization (DRO), which optimizes model performance against the worst-case input permutation. Specifically, PEARL consists of a permutation-proposal network (P-Net) and the LLM. The P-Net generates the most challenging permutations by treating it as an optimal transport problem, which is solved using an entropy-constrained Sinkhorn algorithm. Through minimax optimization, the P-Net and the LLM iteratively optimize against each other, progressively improving the LLM's robustness. Experiments on synthetic pre-training and real-world instruction tuning tasks demonstrate that PEARL effectively mitigates permutation attacks and enhances performance. Notably, despite being trained on fewer shots and shorter contexts, PEARL achieves performance gains of up to 40% when scaled to many-shot and long-context scenarios, highlighting its efficiency and generalization capabilities.

Published as a conference paper at ICLR 2025
PEARL: TOWARDS PERMUTATION-RESILEER LLMS
Liang Chen' Li Shen2* Yang Deng' Xiaoyan Zhao' Bin Liang' Kam-Fai Wong**
The Chinese University of Hong Kong Shenzhen Campus of Sun Yat-sen University 3SMU (ichen, kfwong) bése. cuhk. edu. hk mathshenl i gmail. i com
ABSTRACT
The in-context learning (ICL) capability of large language models (LLMs) en- ables them to perform challenging tasks using provided demonstrations. However, ICL is highly sensitive the ordering of demonstrations, leading instability in predictions. This paper shows that this vuinerability can be exploited to design a natural attack diffccliffiuult for model providers detect-that achieves nearly 80% success rate on LLaMA-3 by simply permuting the demonstrations. Exist- ing mitigation methods primarily rely on post-procsssing and fail to enhance the model's 's inherent robustness to input permutations, raising concerns about safety and reliability LLMs. To address this issue, we propose Permutation- tesiiittonn-reillientt learning PI), a novel framework based on distributionally robust optimization (DRO), which optimizes model performance against the worst-case input permu- tation. Specifically, PEARL consists of permutation-prott pprooosal network (P-Net) and the LLM. The P-Net generates the most challenging permutations by treating it optimal transport problem, which is solved using entropy stroopy-constriieed Sinkhorn algorithm. Through minimax optimization, the P-Net and the LLM iteratively optimize against each other, progressively improving the LLM's robust- ness. Experiments on synthetic pre-training and real-world instruction tuning tasks demonstrate that PEARL effectively mitigates permutation attacks and enhances performance. Notably, despite being trained on fewer shots and shorter contexts, PF.RIL achieves performance gains of up to 40% when scaled many-shot and long-context scenarios, highlighting its efficiency and generalization capabilities.
INTRODUCTION
hallmark of human intelligence is the ability learn and execute new tasks by reasoning from a few examples. Mirroring this, in-context learning (Brown et al. 2020). as a crucial supplement to zero-shot prompting, has shown promising results across a spectrum of complex tasks (Cobbe et al. [021]: Chen et al. 2023a OpenAl et al., 2023). Despite these advancements, the in-context learning capabilities of L.LMs remain fragile. LLMs exhibit sensitivity to permutations of provided demonstrations (Lu al.) 2022 Zhao et al. 202 Reynolds & McDonell 2021), posing challenges for prompt engineering and leaving a significant gap in achieving human- like adaptability.
Most existing studies on ICL primarily aim to enhance normal-case performance (Min al. 2022 Wei et al.) 2023), with limited attention to improving permutation robustness. Current approaches to addressing this issue typically involve modifying training objectives to mitigate the limitations of transformers' unidirectional attention (Xiang et al. (2244) or designing permutation-cquivarinn architectures (Chen al., 2023b). However, these methods often lack scalability. Alternatively, decoding-stagee stage techniques such as output calibration (Zhao et al.) (202T) and demonstration order optimization (Lu et al. 2022) introduce additional computational overhead per inference call, further limiting their practicality. Thus, a critical need remains for methods that fundamentally enhance LLMs" inherent robustness to input permutations.
In this work, we conduct extensive experiments on LLaMA-3, an advanced open-source LLM, to reassess its vulnerability to permutation- attacks from an adversarial perspective ($2). Our findings reveal that even LLaMA-3 remains highly susceptible to simple attacks that alter only
"Corresoondngg Authors.

---

Published as a conference paper at ICLR 2025
CurDial
CurDial
Success (%6)
1.0
(%6
Perfermance
Parfermance
02
23-
Number of Shots
Number of Shots
Threshold (%)
Threshold (%)
Average
Shot (Exhoustive)
Shot austive)
Shot (Neural) Shot (Neural)
Worst
Shot 5 (Exhaustive)
Shot 4 (Neural)
dom
Figure 1: Performance and attack success rates of Llama-3 CurDial and TMW datasets. Left panels: Random, average and worst-case performance as a function of shot number. Right panels: Attack success rates for exhaustive and neural search attack methods at different thresholds.
the order of ICL demonstrations. These attacks preserve the semantic content of examples and introduce no adversarial modifications, yet they degrade performance with success rates exceeding 80%. Consequently, they are difficult for model providers detect but significantly undermine LLM performance, highlighting a critical vulnerability concern.
To counteract the vulnerability to input permutations, introduce novel Permutation-resiliee learning (PEARL) framework, which is based on distributionally robust optimization (DRO) (Ben-'In Tal al., 2011). Unlike standard empirical risk minimization training, adopted by most supervised fine-tuning (SFT) methods, which views each training instance merely in terms of its one or several permutations observed during training, DRO conceptualizes each instance as part of broader distribution that includes all conceivable permutations. This comprehensive set of all possible permutations termed the ambiguity set. By explicitly identifying and optimizing the worst-case within this ambiguity set, our strategy substantially enhances the resilience of LLMs against all different permutations. This paradigm shift-from considering training instances as single data points to viewing them within distribution of potential permutations equips the model to better prepare for and generalize to combinatorial input scenarios.
Specifically, PEARL operationalizes DRO as two-player game, consisting of a cotnuutttion-propppoo network (P-Net) as the adversary and the LLM as the target model. For each training instance, P-Net identifies a challenging permutation of given demonstrations, aiming maximize the LLM's loss. Conversely, the LLM strives to minimize its loss under the P-Net's 's manipulation. thereby performing well these difficult examples. P-Net treats the generation of the adversarial ICL permutation as an optimal transport (OT) (Monge 178T) problem between the distribution over input permutations and the distribution of challenging permutations for LLMs. We solve the OT problem using the Sinkhorn algorithm (Sinkhorn 1966) with an element-wiis entropy constraint designed to prevent trivial solutions. Through adversarial training (AT), both networks improve iteratively. Ideally, at convergence, the P-Net represents a uniform distribution across all permutations, as the LLM handles all possible permutations equally well.
We validate our method in two scenarios: (1) pretraining transformer to in-context learn linear functions (Garg et al. 2022), and (2) instruction tuning of LLMS on the Super-Natural Instructions (Wang et al. 2022). The results demonstrate that, unseen tasks, method consistently improves both the average and worst-case performance of LLMs across different permutations, effectively defending against permutation-be based attacks. Furthermore, despite being trained with much smaller configurations, our method generalizes effectively to many-shot ICI. and long sequences, achieving performance gains of 24% to 40%. These results highlight the efficiency and generalization capabili- ties of our approach. The code is available at httttps: //g/ thub.
REVISITING PERMUTATION VULNERABILITY IN LLMS
This section investigates performance fluctuations in SOTA open-source LLMs when presented with different permutations of given demonstrations, Additionally, from adversarial perspective, explore whether this vulnerability can be exploited to devise an effective attack on LLMs.
2

---

Published as a conference paper at ICLR 2025
Experimental Setups To conduct evaluations, we select two tasks from Super-Naturalllsstuettee (Wang et al. (0022, including Curiosity-
Dialog (CurDial) and TellMeWhy QA (TMW). We test 100 samples for each task, with each sample structured as a quadruple consisting of (instruction, demonstrations, input, output). The number of demonstrations (shots) ranges from two to six.
Following WWnn al. (2022), the performance is measured using the ROUGE-L (Lin) 2004). We analyze the permutation vulnerability of LLaMA-3-8B on two settings as follows:
1) Permutation Vulnerability on Different Number of Demonstrations We first examine the average and worst-case performance of the model across different permutations of input demonstra- tions and the effect of scaling the number of demonstrations. As shown in the left of Figure [ there is a notable observation: adding demonstrations is a double-edged sword. Increasing the number of demonstrations (shots) generally enhances the model's average performance due to richer contextual information. However, it can simultaneously worsen the worst-case performance. This suggests that while additional demonstrations provide useful context, the exponentially growing number of permutations (n!) increases the risk of the model performing poorly on certain input configurations.
2) Input Permutation as Attack We then consider a two-party adversarial scenario between a malicious user (attacker) and model provider (defender). The attacker aims to induce compromised responses from LLMs by permuting ICL demonstrations, making the attack less detectable, Given a task D = {Ps,, Fy y4)]. a sample successfully attacked if its relative performance degradation, induced by the attacker, exceeds threshold 5 E [0%, 100%]. Here, p is an ICL prompt with demonstrations. The set of all possible demonstration permutations TP (To,.- -nn-1): ). where [P] n!. Let g be a performance metric (e.g., ROUGE-L). The attack success rate (ASR) for task D is defined as:
ID
ASR(D,8)
1 T((m--)))) 2 26)
E
(1)
[D]
where I denotes the indicator function, |D] the size of the dataset, and ô is the threshold. The average performance of the i-th sample, (4s is defined by:
m1
Σ MTI, 21)))) [9[1 Pt; (.))) n! j-1
and W; is the compromised performance induced by the attack strategy adopted by the malicious user. Here, we analyze two attack method:
To calculate the upper bound of the effect the permutation- attack can achieve, we assume that the malicious user has unlimited attempts and conducts an exhaustive search. For each sample (Pr, 24, y)), this process involved testing all possible permutations of demonstrations in Q; and identifying the permutation that yields the poorest performance. In this case, the attacked performance is calculated as follows:
Exhaustive Search Attack:
min ((l Pi, (44 y/) IeFF
(3)
Neural Search Attack: To approximate the upper bound established by the exhaustive search when the number of attempts is limited, we employ meta-learning approach to optimize a permutation- proposal network (P-Net). As illustrated in Figure 3 (details are in the Methods section), during training, this network takes the standard sample (P:, Z4, 34) as input and outputs a permutation matrix II; The permuted samples (II, Pi, Ti; ys)) are then fed into the LM to maximize its loss function. During testing, the network generates the most challenging permutation IT; for each sample (p,, D., V.). Then the attacked performance is calculated as follows:
W; = (II: P:, Ty: vc), s.t. II; P-Net(p:, 2;, Vi)
(4)
As shown in the right of Figure T the results indicate that permutation attacks are effective and approachable. Leveraging this characteristic, the exhaustive search attack successfully attacks over 50% and 80% of the samples with achieved a successful rate close to this upper bound across different shots. These results demonstrate that this vulnerability poses real concern, even for advanced L.MSs like LLaMA-3.
50% on two datasets respectively, and the neural attack
Remark These deficiencies may directly stem from the fundamental limitations of standard Empir- ical Risk Minimization (ERM) training, which focuses on optimizing average performance while neglecting worst-case performance. We discuss this issue in depth in the next section and propose a method to address the model's improper behavior unseen but practically valid input spaces.
3

---

Published as a conference paper at ICLR 2025
3 PERMUTATION-RESILIA LEARNING (PEARL)
3.1 INSTRUCTION TUNING VIA DRO
Our objective is to train a LLM to perform well across all possible permutations of given demonstra- tions when prompted with few-shot instructions.
In supervised fine-tuning for few-shot learning, the LLM is trained to predict an output E )/ given an input x E X' and a few-shot instruction p E P. where p typically consists of a sequence of demonstrations, each being input-output pair. Let θ denote the parameter space of the language model, and let & θ x (P x X' x )) + R, be a nonnegative loss function measuring the discrepancy between the model's prediction and the true output. The standard approach is to find parameters E θ that minimize the empirical loss over the training data via empirical risk minimization:
ÒERM == arg min ( 0
(p...))~ (((00 2r y)]
(5)
where P denotes the empirical distribution derived from the training dataset.
Under appropriate assumptions, learning theory (Vapnik [999] Shalev-Shwartz & Ben- David 2014) guarantees that models trained via ERM perform well on the test distribution given sufficient training data. However, in practice, models trained using ERM often fail to generalize well to different permutations of the same set of demonstrations. This occurs because the training set covers only subset of all possible permutations of the demonstrations, and during testing, the model may encounter permutations not seen during training, leading to significant degradation in performance.
To systematically address the permutation sensitivity issue, we propose fine-tuning under the frame- work of distributionally robust optimization, which optimizes the risk under the worst-case distribution within a specified ambiguity set. Specifically, we aim to solve:
ODRO = arg min sup E(p.n))) [6(0; p.r,
(6)
DEE QuEQ
The ambiguity set Q is constructed the convex hull of all distributions obtained by permuting the prompts in the empirical distribution P. Specifically, we define:
e: Σ qm Qu g E AEA-I1 A where Qn = { (II p. x, v) (p. ~ p).
(7)
NEP
Here, II a permutation matrix that reorders the sequence of demonstrations in p. and IP denotes the set of all such matrices. The vector g lies in the IP| 1-dimensional probability simplex A 111--1-
To illustrate the advantages of DRO ERM in handling different permutations, consider the example in Figure 2 For 3- shot training example (p, y) with prompt p containing three demonstrations, there are six possible permutations denoted as (p" ,,, .), (y", -x. y), indexed from 0 to 5. P denotes the empirical distribution of permutations in training data, represented by blue bars. The bars show that permuta- tion indices O, 1, and appear in training data with frequencies, while permutations 2, and do not appear. P& represents the dis- tribution learned by the LLM, represented by purple curves. In panel (a), the ERM- trained model assigns higher probabilities to frequently occurring permutations (0, 1, and lower probabilities to less frequent ones 3,5, leading to poor performance unseen permuta- tions during testing. In contrast, panel (b) shows that the DRO-trained model distributes probabilities more uniformly across all possible permutations, as it explicitly considers them all (Equation (6) during learning. This demonstrates how DRO mitigates ERM's 's limitations by encouraging models to assign reasonable probabilities to all valid permutations, regardless of their frequency in training data.
(a) ERM
(b) DRO
Permutation Index
P
Peruutation Index
Figure 2: Comparison of models trained under ERM and DRO paradigms. The blue bars represent the empir- ical distribution i' of training data, showing different frequencies of six permutations in the training set. The purple curves denote the learned distribution Pe by (a) ERM and (b) DRO models, illustrating their different behaviors on less appeared but valid permutations.
4

---

Published as a conference paper at ICLR 2025
permutate
x2
[ok]
sample
x1
Rlg
P-Net
(ckg
LLM
0.0
1.0
x3
ya
(ck)
backward
Figure 3: An overview of the learning framework. The P-Net is small model incorporating the Sinkhorn operator, trained jointly with the LLM under the adversarial optimization algorithm. Note that the permutation matrix operates the input sequence's 's embeddings (simplified here text sequences for clarity). After training, only the LLM is retained while the P-Net is discarded.
3.2 LEARNING TO GENERATE PERMUTATIONS VIA P-NET
To enable DRO framework to function effectively, need to efficiently find the worst-case scenario within the ambiguity set (solve the max step in Equation (6)) Directly addressing this problem through exhaustive search is computationally infeasible due to the exponential search space.
We address this challenge by drawing inspiration from optimal transport, treating the problem as transportation between permutation distributions. We introduce the Permutation-proppsaa Network (P-Net), denoted as P-Net: (P x X' x y) + (IT),, which learns a distribution over permutations increase task difficulty for the L.M given input examples. As shown in Figure challenging permutations from this distribution to reorder the given demonstrations.
sample
Specifically, P-Net consists of two components: parameter part that extracts features and models the relationships between demonstrations, a non-parameter part using the Sinkhorn algorithm to build the distribution (II), and Gumbel sampling for differentiable sampling from (II A(T)))
Parameter component. The parameter component consists of a feature extractor and cross- relationship modeling layer. The feature extractor is an encoder model that takes an ICL prompt composed of n demonstration pairs p { y/ ); ) and predicting sample (Jr, y), and produces their representations as follows:
([CLS), (II,,))... v.)............ (CLS), (In,, Bn)) [CLS), (zr, u)) Encodor (hı, ba ha, hn+1),
where h, is the representation corresponding the i-th (CLS] token, which is often used to segment and extract the representation of sequences (Devlin et al. 2019b Lu et al. 202I).
After extracting the representations of demonstrations, have // (h he,..., hn) Raxh We then model the pairwise relationships among the demonstrations. Specifically, we design simple dross-demonstration layer obtain a relationship matrix R E R=*n that captures the pairwise relationships between each pair of demonstrations, defined as:
R = (HWH").
(9)
where W Rixh bh is weight matrix, and y denotes a nonlinear activation function.
The matrix R can be interpreted as adjacency matrix in graph theory, where demonstrations serve as nodes, and Ray represents the relationship between demonstrations and j. Specifically, we define Raj the potential increase in task difficulty for the LLM if demonstrations and are swapped: higher values of Ray indicate that swapping these two demonstrations may significantly impact prediction. Thus, this parameterized component models edge prediction process.
However, while R captures the potential for swapping between demonstrations, it is not yet suitable for sampling permutations because its elements can take any real values and do not necessarily form a valid probability distribution. To convert R into a distribution over permutations A (II) that we can sample from, introduce non-parameter component.

---

Published as a conference paper at ICLR 2025
Non-paraeeter component. The con-parmeteer component aims to transform the adjacency matrix R into a doubly stochastic matrix, representing probability distribution over permutations. Specifically, following Adams & Zemel (201 I); Mena et al. (2018), adopt the Sinkhom operator S(.) to obtain such matrices through an iterative process of row and column normalization:
S(R) lim (Te (T, (exp())))),
(10)
Loo0 T - ( T.(R) R @ R1,
(11)
where T, (R) and (R) represent the row and column normalization operators, respectively, Ø indicates element-wise division: and 1, is column vector of ones. As established by (Sinkhomn 1966), the Sinkhorn operator S strictly converges to doubly stochastic matrix as the number of iterations approaches infinity.
To ensure differentiable process when sampling permutations from the distribution, the Gumbel trick (Jang et al. 2017) is applied:
11 lim S ((R+ G)/T), G-y=
(12)
Gl)g Gl ~ U0,
(13)
Where G E _ is the Gumbel noise and т the temperature. As T approaches zero, the result approximates permutation matrix II. Hyperparameters studied in Appendix C
By regarding permutation generation as an optimal transport problem and implementing through P-Net, transform the input permutation distribution into target distribution. Next, introduce how P-Net is co-optimized with the LLM to make the target permutation distribution the most challenging for the current LLM.
3.3 ADVERSARIAL OPTIMIZATION
As illustrated in Figure we adopt an adversarial optimization framework to jointly train the LLM and the P-Net. Let 0 and denote the parameters of the LLM and P-Net, respectively. For each sample (p. x, y) drawn from the empirical distribution P. the P-Net generates an adversarial permutation IT that maximizes the LLM's ss loss. In response, the L.M aims to minimize its loss, adversarially influenced by the P-Net. The LLM's loss function is defined as:
Lim 0) E (.ae::))_....I-P-e(pppppp.... [e(;; (II p..,
(14)
To prevent the P-Net from collapsing trivial solutions, such producing uniform permutations that degrade demonstration semantics, we introduce an entropy-based regularization term:
Lem ()) (p....)) (f..,a)) __-P.NN((s:pp...)) [IIII))
(15)
where H(-) denotes the element-wise entropy function.
This results in a two-player min-max optimization problem with the following objective:
min max (Lim 0) (l)m()))
(16)
where 3 hyperparameter controlling the strength of entropy regularization.
We employ alternating optimization to iteratively update and Φ. The full training procedure is detailed in Algorithm
Algorithm 1: Adversarial Optimization Algorithm for PEARL
Input:
0, 6 (LLM, P-Net); 7e, Te (learning rates); m (inner steps); 3 (entropy coefficient) repeat
for
to do
(p. x. y) P:
// sample training examples // Generate permutations // Compute LLM loss
11 ~~ P-Net( (. p, p,
Lm 0) E(0: p.
(Ø) + H(T):
// Compute entropy regularizatio // Update P-Net
④ - Φ + no V (Tam 1 BTea)
end
0
no Tm ( (0.0) 0)
// Update LLM
until convergence;

---

Published as a conference paper at ICLR 2025
4 IN-CONTEXT LEARNING WITH LINEAR FUNCTIONS
4.1 DATASETS AND EVALUATION METRICS
We investigate in-context learning on linear functions f(x) = TT., where w R. follow-
ing Garg et al. (0022); Guo et al. (2024b). For each w. we construct each example p'
(11, (I1),.... di (z 1;-1) containing a input-output demonstration pairs and a query input
I4+1- A language model LM is trained to minimize:
Σ E(0:p', ((0:p" ((X++)))
min Bp
(17)
where E(-) is the MSE loss and is the number of demonstrations. During testing, we evaluated performance using the same MSE metric. We report the normalized squared error ((LM((p w Jquery) )2/ /d), where is the problem dimension. Detailed settings are in Appendix A.1
4.2 IMPLEMENTATION DETAILS AND BASELINES
Architecture and Training, We implement Ly using a GPT-2 base model (Radford et al. 2019 and train it from seratch generated dataset using the AdamW (Loshchilov & Hutter 2019). Key training parameters include a batch size of 128 and 500k training steps. In the PEARL frame- work, the P-Net is initialized BERT-base (Devlin et al.) 2019a) and also trained from scratch. Implementation details are in Appendix
A.2
Baselines. Consistent with Garg al. (2022), adopt an empirical risk minimization method
with curriculum learning (Bengio et al. 2209] Wu et al. (2020) (ERM+CL) to train the model. The training process gradually increase the number of demonstrations presented the model. allowing for progressive learning of more complex patterns and making the training more stable.
4.3 EVALUATION RESULTS
We evaluate the effect of permutations on the worst-case and average performance of different methods, as well as each method's defence capability against permutation attacks.
Table I: Normalized MSE across permutations.
ERM+CL
ERM+CL
Shot Method
Avg.
Worst.
ERM+=LLSn0t5
REARL
ERM+CI.
1.45
2.67
Attack Success
<
PEARL shat
PEARL
PEARL
0.86
(+40.7)
0.92 (+65.5)
ERM+CI.
1.20
3.34
PEARL
0.79
(+34.1)
1.11 (+66.8)
ERM+CI.
1,28
5.03
PEARL
0.87 (+32.0)
1.33 (+73.6)
Threshold (%)
Figure 4: Comparison of attack success rates.
As shown in Table the performance gap between average and worst-case performance across permutations for the baseline methods was significant, indicating substantial vulnerability permu- tations. Specifically, the worst-case performance of the baseline methods decreased dramatically compared to their average performance, with the relative performance drop increasing from 74.6% 3 shots to 84.1% at 4 shots, effectively losing most of the performance gains achieved by increasing the number of shots. In contrast, our method, PEARL,, not only improved the average performance but also significantly enhanced the worst-case generalization performance compared to the baselines. While the average performance gains tend to plateau as the number of shots increases, the worst-case performance gains continue to rise, increasing from 65.5% at 3 shots to 73.6% shots.
Figure depicts the proportion of successfully attacked samples in terms of (1) different attack pessimistic
success thresholds and (2) number of demonstrations (shots). The former considers scenarios (attacked samples drop large margin), while the latter examines larger input spaces. We observed that PEARL's advantage increased as the threshold grew. At 8 > 50%, the defence success
7

---

Published as conference paper at ICLR 2025
rate for PEARL across all shots was approximately double that of the baseline methods. This indicates that PEARL can effectively prevent pessimistic scenarios (samples attacked with large threshold). Moreover, PEARL's performance improved with increasing number of shots, suggesting better scalability compared to baseline methods.
INSTRUCTION FINE-TUNING OF LARGE LANGUAGE MODELS
5.1 EXPERIMENTAL SETUPS
Datasets. Our instruction tuning data derived from Super-Natural Instructions (Wang et al. 2022), which are part of the FLAN v2 benchmark (Chung We selected 17 representative tasks, comprising natural language generation (NLG) tasks and 8 natural language understanding (NLU) tasks. Following the methodology of Wang al. (022), we randomly designated datasets as held-out test sets and used the remaining 13 datasets for training. Each training dataset contains 150 examples, and each test dataset contains 100 examples, resulting in training set of 1,950 examples and test set of 400 examples, as summarized in Table
Table 2: Summary of datasets.
2024
Split
Category # Tasks Samples
NLG
Training
NLU
NLG
2
200
Testing
NLU
2
200
Evaluation Metrics. Following the practice in Super-Natural Instructions (Mishra al. 2022 Wang
ct al. 2022). we adopt ROUGE-L (Lin (2004) for reporting performance results, due to the diversity of our tasks and the open-ended nature of instruction tuning. We also report single "average" metric across all datasets, following the methodology in FLAN (Wei et al. 2022 2023).
Baselines and Models. We evaluate framework against several learning algorithms: Empirical Risk Minimization (ERM) (Min et al.) 2022), ERM with Demonstration Shuffling (ERM+DS) and ERM with Instance Mixup (ERM+IM) (Zhang al. 2018), and InfoAC (Xiang et al. 22024). We implement FLAN-large as the P-Net and evaluate across five LLMs: Llama3-88, LJama2-7B/133
Mistral-7B. and Gemma-7B. The implementation details are provided in Appendix
5.2 EVALUATION RESULTS
We evaluate PEARL from three perspectives: (1) comparison with training-Stage methods, generalization to diverse type of LLMs, and (3) scalability to many-shot in-context learning (Agarwal al. 2024) and long sequences.
Table 3: Average and Worst-Case Performance of Llama3-8B four held-out tasks: Common- senseQA (CSQA), Curiosity Dialogue (CurDial), CoLA, and Tell Me Why (TMW). Performance improvements (%) over ERM shown in blue. Worst-case performance tested using exhaustive search.
Average
CSQA
CurDial
CoLAA
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
ERM
57.3
58.0
54.0
57.9
43.4
62.0
51.1
ERM+DS
(-0.2)
(-1.6)
62.0
54.0
54.1
37.8
61.0
42.7
ERM+IM
(-6.6)
(-10.))
63.0
44.7
28.1
$7.0
56.3
49.4
INFOAC
55.7 (-2.9)
47.6
57.5
56.0
53.4
36.4
63.0
61.5
48.7
37.3
PEARL
62.9 (+9.8)
56.4 (+14.2)
65.0
62.0
60.3
50.7
71.0
68.0
55.1
44,8
ERM
57.8
38.3
57.7
47.0
61.4
25.9
61.9
52.0
50.3
29.4
ERM+DS
56.1 (-2.9)
39.7 (43.7)
60.0
46.0
54.1
25.4
60.0
56.0
50.3
31.5
ERM+IM
(-4.3)
39.8 (43.9)
59.0
46.0
54.6
28.0
57.6
53.1
50.0
31.9
INFOAC
(-2.6)
(+3.1)
49.0
55.2
24.3
62.1
55.8
48.4
28.8
PEARL
63.1 (+9.2)
46.9 (+22.5)
68.4
62.0
66.7
34.8
64.7
56.0
52.4
34.7
ERM
59.7
30.6
61.3
38.0
62.9
21.3
63.3
45.8
51.1
17.5
ERM+DS
57.7 (-3.4)
31.8 (+3.9)
63.3
40.0
57.3
17.6
60.1
52.0
49.9
17.8
FRM+IM
56.0 (-6.2)
32.4 (+5.9)
63.2
42.0
53.7
17.8
57.6
48.5
49.6
21.3
INFOAC
58.6 (-1.8)
33.0 (+7.8)
63.7
44.0
58.7
19.0
63.9
51.0
48.1
17.0
PEARL
63.1 (+5.7)
39.6 (429.4)
68.4
52.0
69,2
31.3
64.7
52.0
50.1
23.0
8

---

Published as a conference paper at ICLR 2025
Table 3 presents the comparative performance of PEARL against other learning methods. PEARL consistently improves both average and worst-case performance across all unseen tasks. As the number of shots increases, the worst-case performance gain relative to ERM progressively increases from 14.2% at two shots to 29.4% at four shots. Notably, while optimized for worst-case performance, PEARL also achieves superior average performance with gains of 5.7% 9.8%. This improvement may stem from the rapid convergence observed during Llama-7B's s fine-tuning, where the training loss plateaus within one epoch. The rapid convergence suggests that for advanced LLMs like LLama3, focusing on challenging permutations during training is more effective than using random ones an observation consistent with Xu al. (2024).
Extension Different LLMs
Scaling many-shot
40
Performance (%6
30
10
10
Mistral-18
Gemma-78
Uama2-78
Liama3-88 888
32
Number of Shots
Average Gain
Worst
Figure 5: Generalization performance of our method across different types of LLMs and many-shot settings. Left: Performance gains on 3-shot across different L.LMs (Mistral-7B, Gemma-7B, Llama 2-7B, and Llama3-8B). Right: Scaling behavior across many-shot settings (8, 16, 32, and 64 shots) and longer sequences (8k tokens) when trained with 5 shots and a sequence length 512 tokens.
To validate the general applicability of our method, we expanded our experiments to include three additional LLMs: Mistral-7B, Gemma-7B and Llama 2-7B. As shown in the left panel of Figure our method consistently improves worst-case performance by more than 10% in three-shot settings. Additional results for higher-shot settings are provided in Appendix D We also observed that different LLM families exhibit varying sensitivity to imput permutations, with Llama models being the most sensitive, followed by Gemma and Mistral. Despite these differences, the phenomenon remains significant, with performance drops exceeding 10% in most cases. Notably, our method achieves consistent worst-case performance improvements of over 10% for three or more shots, demonstrating its robustness across diverse model families.
We scale our evaluation to the many-shot ICL setting (up to 64 shots) and longer sequences (up to 8,000 tokens) after training with 5 shots and sequence length of 512 tokens. As shown in the right panel of Figure PEARL achieves substantial worst-case performance gains ranging from 24% 40% when generalizing larger shot numbers and longer sequences, despite being trained on a smaller setup. These results suggest that PEARL enables LLMs to learn robust features that generalize effectively both many-shot in-context learning and longer sequences, demonstrating the strong generalization capability of our method. Detailed results are provided in Appendix
Table 4: Shot Efficiency: Average Performance with and without PEARL.
# Shots
2
4
8
16
32
64
ERM
57.3
59.7
61.8
66.9
67.4
68.1
PEARL
62.9
63.1
66.5
70.5
70.0
70.4
As we scale to the many-shot setting, we also observe notable trends in shot efficiency, which quantifies the number of shots baseline model would require to match the average performance of PEARL-trained model. As shown in Table PEARL-trained models achieve comparable average performance while requiring four times fewer shots, highlighting the efficiency of our approach.

---

Published as a conference paper at ICLR 2025
6 RELATED WORK
Order Sensitivity in In-context Learning Despite the huge success of ICL, its robustness to demonstration permutations remains an unresolved challenge (Zhao et al. 2021).
Most Most training-stage methods focus on improving general performance in ICI. (Min et al. 2022 Wei et al. 2023) while neglecting the lack of robustness to the permutations of demonstrations. Recent studies suggest that this phenomenon stems from the autoregressive nature of transformer language models (Chen et al. 2023b Xiang et al. 2024). InfoAC (Xiang et al. 2024) introduces contrastive learning during tuning break the autoregressive constraint and enable bidirectional token visibility; however, their approach achieves limited success and is restricted to classification tasks. Preliminary work of Chen et al. (2023b) shows the DeepSet architecture exhibits better permutation invariance than transformer; however, this MLP-based new architecture is too small to solve complex language modeling tasks. Our approach falls within the category of training- -tagge methods but proposes a general learning framework that enhances permutation robustness in LLMs without modifying the Transformer architecture or its autoregressive objective, thereby preserving scalability.
Inference-stage methods can be categorized into four types: (1) demonstration selection (Chng & Jia 2023 Peng et a 2024), which improves normal-case performance but lacks worst-case guarantees under permutations; (2) output calibration (Zhao al. [202]] Li et al. 2023) Guo et al., (0224)), which are effective for classification but is less applicable to generation tasks due to sequence calibration challenges; (3) order optimization (Lu al. 2022), which finds the best ordering during inference but suffers from exponential complexity; and (4) prediction ensembling: A recent work (Zhngg al. 2024) transforms n-shot ICL into n one-shot predictions and ensembles results-effective for classification but harms generation. In summary, inference-stage methods mitigate order sensitivity often introducing additional inference overhead. Moreover, most methods target classification and underperform on generation tasks. In contrast, our training-stage solution complements inference-stage methods, enhancing LLM robustness without additional inference costs while remaining broadly applicable various tasks.
via
Distributionally Robust Optimization. Distributionally robust optimization optimizes the objective function over ambiguity sets, often defined as balls centered on the empirical distribution (Ben-Tr -all et al. 2013] Lam & Zhou 2015) Duchi et al. 2016 MMyato et al. 2018). Prior applications of DRO have primarily addressed distributional shifts, including label shift (Hu et al.) 2018) and data source shift (Oren et al.] 2019) and group shift (Sagawa et al. 2020). To the best of our knowledge, we are the first to apply DRO to enhance the ICL robustness of LLMs by defining the ambiguity set over all possible permutations of the empirical distribution, thereby providing performance guarantees.
Optimal Transport Optimal transport is fundamental mathematical discipline established by Monge (181): [Katorovich] (1942). It defines a metric for measuring distances between probability distributions, known as the Wasserstein distance, and has been widely employed in machine learning for distribution matching (Montesuma et al. 2024 Xiao et al.) 2024). Our work extends the concept of learning permutation structures through neural networks, as explored in (Mena et al., 2018) for learning to sort numbers or solve jigsaw puzzles. However, we apply OT in the context of LLMs and design a neural network, P-Net, equipped with the Sinkhorn operator to generate challenging permutations, enabling LLMs to undergo DRO training and thereby enhancing their ICL robustness.
7 CONCLUSION
We introduce a novel cemmutation-rssileene learning framework, PEARL, designed enhance the robustness of L.LSs against different permutations. PFARL employs permutation- prooposal network, which leverages the Sinkhorn algorithm to generate challenging permutations, optimized under the DRO to systematically improve LLM's robustness. Through empirical evaluations on both synthetic tasks and real-world citructtiontuunnng tasks, our framework has demonstrated effectiveness in mitigating the permutation- based attacks and enhancing average performance.
While PEARL primarily focuses on improving in in-context learning, it provides a general framework for handling set-structured inputs with inder-independent elements, such multiple documents, images, or videos. We hope this work inspires further research on permutation-resilient learning, contributing to the development of more robust and trustworthy language models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We sincerely thank the anonymous reviewers for their valuable feedback and constructive suggestions, which have helped improve the quality of this work. This work is partially supported by Hong Kong RGC GRF No. 14206324, CUHK direct grant No. 4055209, and CUHK Knowledge Transfer Project Fund No. KPF23GWP20. Li Shen is supported by Shenzhen Basic Research Project (Natural Science Foundation) Basic Research Key Project (NO. JCX11002221222224400411 and CCF-DiDi GAIA Collaborative Research Funds.
REFERENCES
Ryan Prescott Adams and Richard S. Zemel. Ranking via sinkhorn propagation, 2011.
Rishabh Agarwal, Avi Singh, Lei M. Zhang, Bernd Bohnet, Luis Rosias, Stephanie Chan, Biao Zhang, Ankesh Anand, Zaheer Abbas, Azade Nova, John D. Co-Reyes, Eric Chu, Feryal Behbahani. Aleksandra Faust, and Hugo Larochelle. Many-shot in-context learning, 2024. URL https
//arxiv.org/ant aaaaabs44000 4004 11018
A. Ben- Tal, D. den Hertog, A. D. Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59:341-357, 2013.
Aharon Ben-Tal Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Ro- bust solutions of optimization problems affected by uncertain probabilities. Advanced Risk Port- folio Management@ Research Paper Series, 2011. Utt.ht http:///aaaaa.eeeeeeeeeeeeeeeee
Sxg/CorpusID:7550 7617993
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In International Conference on Machine Learning (ICML), pp. 41-48, 2009.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin. Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
Ting- Yun Chang and Robin Jia. Data curation alone can stabilize in-context learning, 2023. URI htttss:/aaxivv www.or/////... /ttgg/bb//aat 2tt.:///222..13..3 10378
Liang Chen, Yang Deng, Yatao Bian, Zeyu Qin, Bingzhe Wu, Tat-Seng Chua, and Kam-Fai Wong. Beyond factuality: A comprehensive evaluation of large language models as knowledge generators. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 6325-6341, Singapore, December URL
2023a. Association for Computational Linguistics. doi: 10.
wrw.oooyoorr//223333 www mimaann. in. 390/
Yonggiang Chen, Binghui Xie, Kaiwen Zhou, Bo Han, Yatao Bian, and James Cheng. Positional information matters for invariant in-context learning: case study of simple function classes, 20236. URL
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoe V. Le, and Jason Wei. Scaling rnsteuetionfinetuuedd language models. Journal of Machine Learning Research, 25(70):1-53, 2024. URL http:// jmir.. r. www..oorgrgpppoee htgg/apprrs/v55/
23-0870. html
11

---

Published as a conference paper at ICLR 2025
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchrr Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL http::tts/////.. httppps://axixyva arxiv.org/
abs/211 14168
Jacob Devlin, Ming-Wey Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers); pp. 4171-4186, Minneapolis, Minnesota, June 2019a. Association for Computational Linguistics. doi: 10.18653/v1/N19- 1423,. URL : http://aaaaaattoo.....o...o...aa / wwg/N19/g// 1423
Jacob Devlin, Ming- Ming-Wei
Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019b. URL httptttss///axii..rrrt abs/1810 0. 04805
J. Duchi, P. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. arXiv, 2016.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. In Advances in Neural Information Processing Systems, volume 35, pp. 30583-30598. Curran Associates, Inc., 2022.
Qi Guo, Leiyu Wang, Yidong Wang, Wei Ye, and Shikun Zhang. What makes a good order of examples in in-context learning. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics: ACL 2024, pp. 14892-14904, Bangkok, Thailand, August 2024a. Association for Computational Linguisties. doi: 1/228255//1/2224, fin.:////.iiiiinnn--aaa-aa 884
Andings-acl..88 URL https: http:////osot//////aaaoooooooooo
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with repre- sentations, In The Tweifth International Conference on Learning Representations, 2024b. URL net/ httt:/ooummmid-iwwwweee
httts:://ppereet
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. International
Conference Learning Representations, 2022. URL
forrmm?
[ddnzeVVeeeFYF9
W. Hu, G. Niu, I. Sato, and M. Sugiyama. Does distributionally robust supervised learning give robust classifiers? In International Conference Machine Learning (ICML), 2018
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations, 2017. URL httpLttttp///oeeeeevieee
net/ forum?ormm? 11=
Leonid Kantorovich. On the transfer of masses. In Doklady Akademii Nauk, volume 37, pp. 227-229, 1942.
H. Lam and E. Zhou. Quantifying input uncertainty in stochastic optimization. In 2015 Winter Simulation Conference, 2015.
Hongjing Li, Hanqi Yan, Yanran Li, Li Qian, Yulan He, and Lin Gui. Distinguishability calibration to in-context learning. In Andreas Vlachos and Isabelle Augenstein (eds.), Findings of the Association for Computational Linguistics: EACL 2023, pp. 1385-1397, Dubrovnik, Croatia, May C 2023. Association for Computational Linguisties. Lingutttio doi: 10.18655/ findings-eaal..122 URL attps://aclaanthoooyyyy...
Chin-Yews Lin, ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Ou, pp. 74-81, Barcelona, Spain, July 2004. Association for Computational Linguistics.
https:::: / ac www.annoooooyy...oooooooom org/ W04 1013
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https
05101
12

---

Published as a conference paper at ICLR 2025
Shuqi Lu, Di He, Chenyan Xiong, Guolin Ke, Waleed Malik, Zhicheng Dou, Paul Bennett, Tie- Yan Liu, and Arnold Overwijk. Less is more: Pretrain a strong Siamese encoder for dense text retrieval using a weak decoder. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 3780-2791, Online and Punta Cana, Dominican Republie, November 2021. Association for Computational Linguistics. doi: 10. 10.18653/v1/200 1/202/////22.emmmlppp. wet.0/1eemmlpmmmne comnlp-main.22 main.20n. URL ht pattss:/aallanthologgs
org/ 2021 2tt.aaaannnoooooooorg2222222222. emnlp ma mnain... 220
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 Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 8086-8098, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 1/255//1//2222.alllloooo 556. URL httpt::tp//t....t..
pttttp::/accaattaie http://aaaaanthoo....ooo org 2022 awl-llllllngg 556
Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permutations with gumbel-sinkhon networks. In International Conference Learning Representations, 2018.
Sewon Min, Mike Lewis, Luke Zettlemoyer. and Hannaneh Hajishirzi. MetalCL: in context. In Marine Carpuat, Marie-Catherine de Marneffe, and Ivan Vladimir Meza Ruiz (eds.), Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. pp. 2091-229999 Seattle, United States, July 2022. Association for Computational Linguistics.
Learning to learn
Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. Cross-task general- ization via natural language crowdsourcing instructions. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computaaionnl Linguistics (Volume 1: Long Papers), pp- 3470-3487, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 1/2255//1/22222..allllo 244. URL wttp:attw...ooooooo2.2.....eeoe.
244
T. Miyato, S. Maeda, S. Ishii, and M. Koyama. Virtual adversarial training: regularization method for supervised and supervised learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2018.
Gaspard Monge. Memoire la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences de Paris, 1781.
Eduardo Fernandes Montesuma, Fred Ngolè Mboula, and Antoine Souloumiac. Recent advances in optimal transport for machine learning, 2024. URL http:/:///xxww..////aaaa 16156
OpenAI, :, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valcrie Balcom, Paul Baltescu, Haiming Bao, Mo Bavar- ian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro. Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gibson, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Loneijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga. Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin. Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Lukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Hendrik
13

---

Published as a conference paper at ICLR 2025
Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike. Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reiichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo. Hyeonwoo Noh, Long Ouyang, Cullen O' Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipee de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond. Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher, Felipe Petroski Such, Natalie Summers. Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine Thompson. Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Preston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang. Alvin Wang. Ben Wang, Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report, 2023.
Y. Oren, S. Sagawa, T. Hashimoto, and P. Liang. Distributionally robust language modeling. In Empirical Methods in Natural Language Processing (EMNLP), 2019.
Keqin Peng, Liang Ding, Yancheng Yuan, Xuebo Liu, Min Zhang, Yuanxin Ouyang, and Dacheng Tao. Revisiting demonstration selection strategies in in-context learning. In Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 9090-9101, Bangkok, Thailand, August 2024, Association for Computational Linguistics. doi: 10. 10.118653/11/222224c-ooog URL httpt:tp/////m aa/aacaaaaaa http//aaaanattgg...ooooo
www.:222...loooono. 492
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the few-shot paradigm. In Extended Abstracts of the 2021 CHI Conference on Human Factors in Computing Systems, CHI EA '21, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450380959 doi: 00.14453411763..45176.. URL httpp:////// /rg/
10. 20.1145/3411766 3451760
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020. URL htthtt
//openreview.oo //openreview.net/Errmm
ryxGuUrFvS
Shai
and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, USA, 2014. ISBN 1107057132.
Richard Sinkhorn. A relationship between arbitrary positive matrices and stochastic matrices. Canadian Journal of Mathematics, 18:303-306, 1966. doi: 10.4153/CIMM1196--333.
Vladimir Vapnik. The Nature of Statistical Learning Theory. Springer: New York, 1999.
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Anjana Arunkumar, Arjum Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, Eshaan Pathak, Giannis Karamanolakis, Haizhi Gary Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby
14

---

Published as a conference paper at ICLR 2025
Kuznia, Krima Doshi, Maitreya Patel, Kuntal Kumar Pal, Mehrad Moradshahi, Mihir Parmar, Mirali Purohit, Neeraj Varshney, Phani Rohitha Kaza, Pulkit Verma, Ravsehaj Singh Puri, Rushang Karia, Shailaja Keyur Sampat, Savan Doshi, Siddhartha Mishra, Sujan Reddy, Sumanta Patro, Tanay Dixit, Xudong Shen, Chitta Baral, Yejin Choi, Noah A. Smith, Hannaneh Hajishirzi, and Daniel Khashabi. Super-naturalinnstaciins Generalization via declarative instructions on 1600+ nlp tasks, 2022. URL TRL/tttss///axxx.t http://awwxx...ooo//////22.2.. 07705
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe Le. Finetuned language models are zero-shot learners. International Conference Learning Representations, 2022. htttt::::/pperee//eeee
ld GEZrGCozzz
Jerry Wei, Le Hou, Andrew Lampinen, Xiangning Chen, Da Huang. Yi Tay, Xinyun Chen, Yifeng Lu, Denny Zhou, Tengyu Ma, and Quoe Le. Symbol tuning improves in-context learning language models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 968-979, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v 1/201/2222ea wet.3mmmmmllmmm..... URL
emnlp main. 61
Xiaoxia Wu, Ethan Dyer, and Behnam Neyshabur. When do curricula work? arXiv preprint arXiv:2012..0317. 2020.
Yanzheng Xiang, Hanqi Yan, Lin Gui, and Yulan He. Addressing order sensitivity of in-context demonstration examples in causal language models. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics: ACL 2024, pp. 6467- 6481, Bangkok, Thailand, August 2024. Association for Computational Linguistics, doi: 10.18653/ v1/2024. v1/2024.findinssaadl findings-acl.386. URL https: // / ant 386
htt///acccaaooooooooooooo wwwhhooog....oogo/ooooooo findi wan-nngg---e acl
Weiwei Xiao, Yongyong Chen, Qiben Shan, Yaowei Wang. and Jingyong Su. Feature distribution matching by optimal transport for effective and robust coreset selection. Proceedings of the AAAI Conference on Artificial Intelligence, 38(8)99196 Mar. 2024, doi: 10. 1609/aaai v3888 28771.
URL https:://jjs aaai. wwr./orgg/.n.... hhtp:pppppaaaaaaaaaaaaaaaaa htt.aaaaaaa....aaaaaaaaaaaaaaaaa
Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, Qingwei Lin, and Daxin Jiang. WizardLM: Empowering large pre-trained language models to follow complex instructions. In The Twelfth International Conference on Learning Representations, 2024. URL
net wwt.w//wwwweee/oreaae.a...a. (at.orummdd-Ceh330010 risk minimization, 2018. ott.://eeee..eerrroooooooooooood
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz mixup: Beyond empirical
Kaiyi Zhang, Ang Lv, Yuhan Chen, Hansen Ha, Tao Xu, and Rui Yan. Batch-ICL: Effective, efficient, and order-agnostic in-context learning. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics: ACI. 2024, pp. 10728- 10739, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v1/202,, findings-acl.638. URL http:pss///ccaaatooloo.oooooo
faraa..iiningggg acl. 638
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before Improving few-shot performance of language models. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 12697-12706. PMLR, 18 24 Jul 2021. URL ht..pttss:t. http::/://oooceedeeg/e/ mlr. hdr.press/vi33 zhao21c..ccooolll html
15

---

Published as a conference paper at ICLR 2025
APPENDIX
A DETAILED SETUP OF ICL WITH LINEAR FUNCTIONS
A.1 DATASETS CONSTRUCTION
We investigate training language model to perform in-context learning linear functions, following
Garg et al. 2022>; ; Guo et al. (2024b). The function class is defined as F { f(x) = w
F. where d is the imput dimension. Each data sample is constructed follows:
(a) Function sampling: A weight vector w NO Ia) sampled, defining a linear function
fx)
w x.
(b) Input sampling: Inputs 21, .2....,,kk+, ~ NO, La) are independently drawn.
(c) Output generation: For each input, the corresponding output is computed as f(x,,) w TTx
for: 1,2,.... k
The input prompt p' consists of i demonstrations and the (i + 1)-th example the query: p' (x1, f (x1). x2. (x22),... ... Zt, (z;), r'+1). We trained a language model Lg, parameterized by 0, to minimize the expected loss over all input prompts:
Σ &
f(0:''((E=+))) (((H))))
min
(18)
where I(-) is the mean squared error (MSE) loss. During testing, evaluated performance using the same MSE metric. We report the normalized squared error (LLM(p) w x'quer)² /2//)////// /dd, where is the problem dimension.
A.2 IMPLEMENT DETAILS
Architecture. Following Garg al. 22022).
implement Lg using GPT-2 architecture (Radford 2019)) with 12 layers, 8 attention heads, and a hidden dimension of 256. The model takes as input sequence of vectors in its embedding space and predicts the next vector in the sequence within the same space.
Training. We pre-train the model from scratch generated dataset of 40k linear functions using the AdamW [Loshchilov & Hutter 2019). We employ a batch size of 128 and trained for 500k steps, selecting the best checkpoint based on validation set performance. In the PEARL framework, randomly initialize the P-Net with a BERT-base-sizee transformer encoder, also pre-training it from scratch. During testing, we sample novel functions to assess the model's 's ability infer new weights w through in-context demonstrations.
B DETAILED SETUP OF INSTRUCTION FINE-TUNING
B. DETAILS OF DATASETS
The details of datasets used in instruction tuning is presented in Table
B.2 BASELINE AND IMPLEMENTATION DETAILS
To evaluate the performance of our trained model, compare it with other learning algorithms.
Empirical Risk Minimization (ERM) Min et al., 2022): : Standard approach minimizing the average
loss over the training dataset, adopted by mainstream instruction tuning models such as FLAN (Chung
ct al. 2024), Natural Instructions (Mishra et al. 2022 Wang et al. 2022). and MetaICL (Min et al.
2022
ERM with Demonstration Shuffling (ERM+DS) (Zhang et al. 2018): Enhances ERM by randomly shuffling the order of in-context demonstrations within each sample at each training step. This introduces robustness by exposing the model to different permutations demonstrations during training. It can be considered form of epoch-level data augmentation.
16

---

Published as a conference paper at ICLR 2025
Table Details of datasets used in instruction tuning from natural instructions.
Task ID
Task Name
Source
Category
1297
QASC Question Answering
QASC
Question Answering
442
COM_ QA Paraphrase Question Generation
COM QA
Question Rewriting
908
DialogRE Identify Familial Relationships
DialogRE
Speaker Relation Classification
288
Gigaword Summarization
Gigaword
Title Generation
582
Natural Questions Answer Generation
Natural Questions
Question Answering
151
TOMQA Find Location Easy Clean
TOM QA
Question Answering
1714
ConvA13 Sentence Generation
ClariQ
Dialogue Generation
379
AGNews Topic Classification
AG News
Text Categorization
639
MultiWOZ. User Utterance Generation
MultiWOZ
Dialogue Generation
209
Stance Detection Classification
StarCon
Stance Detection
1516
IMPPRES Natural Language Inference
IMPPRES
Textual Entailment
589
Amazon Food Summary Text Generation
Amazon Reviews
Summarization
1285
KPA Keypoint Matching
ArgKP
Text Matching
ERM with Instance Mixup (ZRM+ al. 2018): Incorporates Instance Mixup technique during each training step. For each data point, we generate multiple augmented versions by randomly selecting different in-context demonstrations. We perform multiple forward passes compute the loss for each augmented version, average these losses, and then perform a single backward pass using the averaged loss. This approach provides finer-grained data augmentation compared to demonstration shuffling. Notably, by comparing this baseline with our method, we contrast min-mean optimization (ERM+IM) with min-max optimization (our method).
InfoAC:
(Xiang et al. 2024) is a training-stage method that employs contrastive learning to enable earlier tokens access information from later tokens, amining to mitigate the order sensitivity of ICL inherent in autoregressive LM.
By including these baselines we provide a comprehensive evaluation of our proposed method.
As for the proposed PEARL framework, we select the LLaMA3-SB model as our LLM and the FLAN-large encoder as the P-Net. Both models are fine-tuned using LoRA (Hu al. 2022). with the number of finetuned parameters of P-Net being 1/20 that of the LLM. We train the models on instruction dataset for two epochs using a single NVIDIA A40 GPU, with a batch size 16, resulting in total of 246 training steps. The optimizer used was AdamW. W. The learning rates for the P-Net and the L.LM are set to 10-4 and 3 x 10-1, respectively. For the Sinkhorn algorithm, we use 80 iterations, temperature parameter of 0.1, and entropy constraint coefficient 3 1.0.
B.3 DETAILS OF HYPERPARAMETER SETTINGS
In this section, we provide comprehensive overview of the hyperparameter settings used in our experiments (Table(6). The hyperparameters can be categorized into three groups: (1) basic LLM training parameters, such learning rate and batch size; (2) LoRA configuration parameters; and (3) P-Net optimization parameters. These byperparameters were selected based on average validation performance and kept consistent across comparative experiments to ensure fair comparison.
C ANALYSIS OF HYPERPARAMETERS IN INSTRUCTION FINETUNING
We conduct analysis understand the impact of key hyperparameters on P-Net learning and our overall framework. Our analysis focuses on two main aspects: the effect of the entropy constraint strength, and the influence of iteration number and temperature in the Sinkhorn algorithm.
Influence of Entropy Regularization in OT We examine the impact of the entropy regularization coefficient OT, testing values of 0.3, 1.0, 3.0, and 10.0 (Figure 77. At a low coefficient (0.3), P-Net's gradient norm remained small, indicating minimal learning and potential generation of simple semantic overlaps to satisfy adversarial training requirements. Concurrently, the LLM's gradient norm struggled to decrease. The gradient norm for N-Net peaked 1.0, suggesting optimal learning conditions. As coefficients increased to 3.0 and 10.0, P-Net's 's gradient norm decreased again, suggesting excessive restrictions. The range of 1.0-3.0 provided an ideal balance, encouraging
17

---

Published as a conference paper at ICLR 2025
Category
Hyperparameter
Value
Learning rate
3e-5
Batch size
16
LLMs
Max sequence length
512
Weight decay coefficient Epoch
0.1
2
Rank
8
Alpha
32
LoRA
Dropout
0.1
P-Net target modules
9.
q_proj. k __projj. v_proj.
LLMs target modules
o_proj. gate_proj.
up up_proj, down_proj
Temperature
0.1
Iteration coefficient
80
Entropy constraint Noise
1.0
P-Net
0.3
Learning rate
le-4
Batch size
16
Max sequence length
512
Table 6: Hyperparameter settings used our main experiment.
Figure 6: Impact of number of iterations and temper- ature on the average/worst-case performance.
PM
Temperature
# Iter.
0.03
0.1
0.3
80
55.7 / 40.0
55.7 / 40.0
55.4/
200
55.7 40.0
55.8/ 40.0
55.8 40.6
Figure 7: Impact of entropy coefficient.
P-Net to extract meaningful information from the LLM without oversimplifying or overcomplicating the task. In contrast, the LLM's s gradient norm decreased consistently with increasing coefficients, indicating a distinct response to entropy regularization.
Effect of Sinkhorn Algorithm Parameters We investigate the interplay between two critical parameters in the Sinkhorn algorithm: number of iterations and temperature. Intuitively, these parameters positively correlated; higher iteration counts typically allow for higher temperatures. Our experiments, however, reveal an unexpected robustness to parameter variations. With the entropy regularization coefficient fixed 1, we vary the number of iterations (80, 200) and temperature (0.03, 0.1, 0.3). As presented in Table surprisingly, these substantial parameter changes result in minimal performance variation. This suggests that the Sinkhorn algorithm in our framework less sensitive to these parameters than initially hypothesized, potentially indicating wider range of stable configurations for practical applications.
D EXTENDED INSTRUCTION FINETUNING ACROSS DIVERSE LLMS
We expanded our evaluation more LLMs: Mistral-7B. Gemma-7B and earlier generations such as Llama2-7B and 1lama2-13B, as detailed in the tables from Table 7 to Table (TO).
Sensitivity to Permutations Across LLM Families Our analysis reveals that different LLM families exhibit varying degrees of sensitivity to permutations. The sensitivity ranking, from highest to lowest, is as follows: Llama, Gemma, and Mistral. Notably, all examined families showed significant performance declines, typically exceeding 10 %.
18

---

Published as a conference paper at ICLR 2025
Table Instruction fine-tuning results for Mistral-7B evaluated on four held-out tasks. Performance gains (%) over the ERM baseline are indicated in blue.
Average
CSQA
CurDial
CoLA
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
ERM
64.1
58.1
67.0
64.0
54.6
41.8
81.0
78.0
53.7
48.5
PEARL
67.0 (+4.5)
62.4 (+7.5)
68.0
66.0
59.4
49.0
82.0
78.0
58.4
56.7
ERM
66.6
56.1
67.0
62.0
63.7
38.9
80.0
76.0
55.6
47.3
PEARL.
69.5 (+4.3)
62.8 (+12.0)
70.0
66.0
70.1
60.1
83.6
78.0
54.1
47.0
ERM
66.7
50.4
68.9
60.0
67.6
47.8
74.2
52.0
55.9
41.6
PEARL
68.3 (+2.5)
57.1 (+13.4)
69.9
62.0
71.6
54.8
74.9
66.0
56.8
45.5
ERM
67.9
50.7
67.5
56.0
70.7
52.6
56.0
57.4
38.2
PEARL 70.2 (+3.4)
58.1 (+14.5)
70.4
64.0
76.7
59.3
66.0
60.4
43.0
Table 8: Instruction fine-tuning results for Gemma-7B evaluated on four held-out tasks. Performance gains (%) over the ERM baseline are indicated blue.
Average
CSQA
CurDial
CoLA
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
ERM
66.2
59.5
71.0
70.0
59.1
46.1
77.0
70.0
57.8
52.0
PEARL
66.3 (+0.0)
60.7 (+2.0)
74.0
68.0
47.3
39.2
82.0
78.0
61.7
57.6
ERM
64.7
52.5
70.7
64.0
67.1
45.2
60.0
50.5
40.7
PEARL
68.4 (45.8)
59.3 (+13.0)
74.7
68.0
59.2
42.5
78.7
76.0
61.0
50.6
4
ERM
65.0
46.5
65.0
54.0
71.4
41.1
72.5
58.0
51.1
32.9
PEARL
67.2 (+3.4)
52.5 (+13.0)
71.4
60.0
60.7
38.9
75.9
66.0
60.8
45.2
ERM
64.3
46.3
65.9
54.0
73.4
48.3
65.6
50.0
52.3
32.9
PEARL
66.3 (43.1)
51.0
(+10.2)
70.3
60.0
63.4
43.6
71.3
60.0
60.2
40.4
Table 9: Instruction fine-tuning results for Llama2-7B evaluated on four held-out tasks. Performance gains (%) over the ERM baseline are indicated blue.
Average
CSQA
CurDial
CoL.A
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
2
ERM
56.6
46.3
56.0
50.0
61.3
50.2
58.2
42.0
50.7
43.1
PEARI.
57.4 (+1.5)
46.5 (+0.4)
58.0
48.0
55.2
44.7
62.0
48.0
54.4
45.4
ERM
58.2
34.0
52.7
34.0
64.0
36.4
66.0
36.0
50.1
29.4
PEARL
59.6 (42.3)
40.4 (+19.1)
56.3
40.0
66.2
46.2
67.0
42.0
48.7
33.5
ERM
58.9
19.9
60.0
26.0
68.1
24.4
60.2
14.0
47.3
15.1
PEARL
60.5 (+2.7)
31.6 (+59.1)
61.2
40.0
69.4
40.1
62.4
24.0
48.9
22.4
5
ERM
61.9
25.8
59.0
32.0
74.2
43.9
65.7
10.0
48.6
17.1
PEARL
(+1.6)
32.1 (+24.7)
62.4
38.0
73.3
43.4
64.8
24.0
51.0
23.0
Table 10: Instruction fine-tuning results for Llama2-133 13B evaluated on four held-out tasks. Perfor- mance gains (%) over the ERM baseline are indicated in blue.
Average
CSQA
CurDial
CoLA
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
2.0
ERM
66.3
56.6
56.0
46.0
72.6
56,2
83.0
76.0
53.4
48.0
PEARL
67.9 (42.4)
60.7 (+7.3)
64.0
58.0
73.8
64.2
81.0
76.0
52.6
44.4
3.0
ERM
65.7
46.2
55.7
38.0
76.4
51.3
77.7
56.0
53.1
39,6
PEARL
68.5 (+4.2)
50.3 (+8.7)
62.7
44.0
81.0
58.4
76.7
56.0
53.5
42.6
4.0
ERM
65.8
33.2
58.2
28.0
79.6
41.6
73.7
38.0
51.8
25.0
PEARL
66.4 (+0.9)
40.2 (+21.1)
63.3
42.0
80.4
45.5
69.4
42.0
53.1
29.1
19

---

Published as a conference paper at ICLR 2025
Adaptation of the Proposed Method In scenarios with three more examples, method consistently demonstrated substantial improvements. often enhancing worst-case performance by more than 10%. These results confirm the robustness and effectiveness of our approach.
E SCALING TO MANY-SHOT IN-CONTEXT LEARNING
We evaluate the scalability of PEARL by extending analysis to many-shot scenarios, testing performance with 8 to 64 in-context examples (Table IT.. Notably, despite being trained solely on 5-shot demonstrations, PEARL exhibits strong generalization to settings with substantially more examples. Using Llama3-8B as our base model, we compare PEARL and ERM training approaches four held-out tasks, Our analysis reveals persistent performance advantages of PEARL over the ERM baseline across all shot regimes.
Table 11:
Performance evaluation across 8-, 16-, 32-, and 64-shot settings comparing PEARL and ERM learning algorithm for Llama3-8B on four held-out tasks, with gains (%) relative to the ERM.
Average
CSQA
CurDial
CoLA
TMW
Shot
Method
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
Avg.
Worst.
ERM
61.8
21.3
61.4
36.0
68.3
22.7
62.7
16.0
54.8
10.6
PEARL
66.5 (47.6)
29.7 (+39.2)
67.7
44.0
77.1
28.7
65.0
32.0
56.2
14.0
16
ERM
66.9
21.3
67.3
36.0
76.5
31.4
67.2
8.0
56.5
9.7
PEARL
70.5 (+5.3)
26.3 (+23.7)
70.9
46.0
83.9
37.5
70.1
12.0
56.9
9.8
ERM
67.4
19.3
32.0
77.8
30.7
56.1
PEARL
70.0 (+3.8)
26.4 (+36.4)
70.0
44.0
82.6
40.3
70.6
12.0
56.6
9.1
64
ERM
68.1
20.6
68.1
38.0
76.9
27.7
72.2
8.7
55.0
8.0
PEARL
70.4 (43.5)
28.2 (436.7)
69.5
46.0
82.9
38.9
74.2
19.6
55.1
8.1
F BEST-CASE PERFORMANCE
Although our methodology was initially designed to optimize for pessimistic (worst-case) scenarios, we have also included an evaluation of the best-case performance for both PEARL and ERM to provide balanced perspective. The results are shown in the Table 12
Table 12: Best performance comparison between ERM and PEARL
#Shot
Method
Average
Gain
CSQA
CurDial
CoLA
TMW
ERM 64.1
68.8
64.4
64.1
59.2
PEARL
68.8
7.2%
73.4
69.2
70.3
62.1
ERM PEARL
72.8
70.3
85.0
65.6
70.3
77.0
5.7%
73.4
87.9
79.7
66.9
4
ERM PEARL
82.9
81.3
92.4
78.1
79.7
84.3
1.7%
82.8
93.6
81.2
79.5
5
ERM
86.8
84.4
95.3
81.3
86.2
PEARL
89.3
2.9%
87.5
96.5
85.9
87.3
Surprisingly, the results show that across all datasets and in every shot condition, PEARL's best performance consistently exceeded that ERM. This indicates that our method not only optimizes performance in worst-case scenarios but also slightly enhances best-case performance.
20


Reviews:
[{"review_id": "js36SmuM8c", "reviewer": "Reviewer_SLuu", "summary": "The authors propose a method for improving the robustness of ICL to permutations in the demonstrations. The idea is to pose the problem as a min-max problem akin to distributionally robust optimization, where the LLM is trained to be robust over hard permutations generated by the adversarial distribution. The adversarial distribution is generated with another neural network and is learned jointly with the LLM, with alternating updates that are reminiscent of GANs. The authors present experiments that re-validate the existence of poor demonstration orders (using exhaustive search and search with their neural network), and demonstrate that their method, PEARL, outperforms two baselines (shuffling and mixup) on learning linear functions and instruction fine-tuning.", "questions": "+ Are the baselines in the paper supposed to correspond to methods explored in previous literature? I feel I've seen demonstration shuffling in other papers, and instance mixup suggests it is based on the mixup augmentation, but no citations or descriptions are given for either. \n+ How novel is the attack method? Currently, I feel the contribution of the attack is not as strongly posed as it could be---there is some discussion on other approaches being exponential in cost, but no explicit experimental comparison that touts the concrete benefits. The current section just shows that the attack works, but the fact that it works is not surprising given the previous work. This could be a more significant result if compared more concretely with respect to the appropriate prior work. \n+ Where do the authors see their work in this space, given the other methods for improving robustness to permutations and/or quantifying the impact of position? Are these comparable, or can they be discussed in any capacity? I would not personally require that these need to be added as additional rows in the experiment, as I fundamentally disagree with the \"moving goalposts\" that the field has become accustomed to. However, I do believe some amount of earnest and thoughtful discussion on how these approaches relate to each other would be fair to include, since the latest two were released in early 2024.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The paper combines distributionally robust optimization using transport maps with the min-max problem of learning over worst-case perturbations, appears to be a unique approach. \n+ The attacker model for finding worst-case perturbation looks to have a favorable computational complexity. \n+ The robustness approach improves upon random and mixup-based baselines.", "weaknesses": "Most of this feedback is centered around how this work is situated against / compares to existing work in the literature. \n\n+ While the authors appear to be aware of some other works studying the fragility of ICL to demonstration order, I felt the paper did not situate their work relative to the existing studies. There is a non-trivial body of work on specifically studying demonstration ordering that has come out since the \"Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity\" paper by Lu et al. 2021, and this body feels like it should be discussed as a section in related work. Also related is Chen et al. 2023 (same problem but perhaps more theoretical in nature). \n+ Similarly, the experiments do not compare how the approach compares to prior work in this space. For example, Chang & Jia (2022) quantify the influence of examples and their position and use this information to stabilize ICL.   Xiang et al. (2024) and Zhang et al. (2024) explicitly have methods that try to solve the permutation robustness problem for ICL. \n+ Section 3 is dedicated toward \"unveiling the vulnerability of LLMs to input permutations\". The phrasing makes it sound like this is new contribution, but perturbation ordering at this point is not a new phenomenon. \n\nChang, Ting-Yun, and Robin Jia. \"Data curation alone can stabilize in-context learning.\" arXiv preprint arXiv:2212.10378 (2022).\nChen, Yongqiang, et al. \"Positional Information Matters for Invariant In-Context Learning: A Case Study of Simple Function Classes.\" arXiv preprint arXiv:2311.18194 (2023).\nZhang, Kaiyi, et al. \"Batch-icl: Effective, efficient, and order-agnostic in-context learning.\" arXiv preprint arXiv:2401.06469 (2024).\nXiang, Yanzheng, et al. \"Addressing Order Sensitivity of In-Context Demonstration Examples in Causal Language Models.\" arXiv preprint arXiv:2402.15637 (2024).", "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": "vqck1qazCh", "reviewer": "Reviewer_ZWPA", "summary": "The paper proposes a permutation-resilient learning framework (PEARL) aimed at improving the robustness of Large Language Models (LLMs) to the order in which examples are presented in in-context learning tasks. The framework is based on adversarial fine-tuning, where the adversary of the LLM tries to find a permutation of the inputs most detrimental to the LLM’s performance. \n\nThe paper starts by empirically establishing that, for Llama 3 8B, there is indeed a significant variation of in-context learning performance depending on the ordering of the inputs. Then, they propose adopting a Distributionally Robust Optimization (DRO) framework for fine-tuning LLMs to mitigate this. \n\nIn their setting, DRO requires finding permutations where the LLM performs poorly. Since an exhaustive search would be intractable, the authors propose using a neural network (referred to as P-Net) to propose hard permutations. The network outputs are then passed through the Sinkhorn algorithm, which converts them into a doubly-stochastic matrix, and a differentiable sampling step (which uses Gumbel noise), yielding a permutation matrix. The authors then propose an adversarial optimization routine for jointly training the P-net and fine-tuning the LLM. The P-net is used only for training, and can be discarded at inference time.\n\nThe authors then evaluate PEARL in the established toy setting of in-context learning of linear functions, and in a real-world LLM fine-tuning setting. In both settings, they find that PEARL improves both average- and worst-case model performance, compared to existing approaches such as raw empirical risk minimization (ERM) or ERM combined with curriculum learning.", "questions": "1. As presented in the methods section, the matrix W does not need to be positive definite, meaning that $(h_1, h_2) \\mapsto h_1^T W h_2$ need not define an inner product in the model’s embedding space. Do you have an intuition or empirical results as to whether constraining W to be positive would harm performance?\n\n1. Do you have an explanation for why the ERM+DS and ERM+IM baselines make average-case performance worse across all numbers of shots in Table 3?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Empirical results on LLM fine-tuning: the results indicating improved average- and worst-case performance in standard ICL benchmarks indicate that PEARL can be useful for practitioners as a part of their fine-tuning pipeline.\n\n1. Effectiveness in the few-shot setting: In addition, the gains above are already noticeable with a small number of shots (2, 3 or 4), indicating the method does not require a very large number of in-context examples to be beneficial (if this were the case, it could hinder the method’s practical usefulness).\n\n1. Practical approach to working with permutations: the proposed method for sampling hard permutations is elegant, as it reduces an intractable search problem to learning an attention-like interaction between token embeddings of ICL examples, followed by a differentiable sampling routine involving the Sinkhorn algorithm. The adversarial optimization routine enables the P-net to be discarded at inference time, not requiring any pre- or post-processing of model inputs and outputs. Meanwhile, many adversarial robustness methods do incur additional inference-time overhead.\n\n1. Clarity and presentation: the paper is clearly written and easy to follow, despite minor points referenced below.", "weaknesses": "1. The authors work only with Llama 3 8B. It would be relevant to assess the model’s generalizability to include other models; in particular Llama 2 7B and 13B (from the previous generation), Mistral 7B v0.2 and Gemma 7B.\n\n1. Lack of evaluations for many-shot settings: recent work in adversarial robustness has highlighted the vulnerability of LLMs to many-shot adversarial attacks (Anil et al. 2024). Hence, it would strengthen the paper to include additional evaluations where the number of shots is much larger than 5 (e.g. 32, 64 and 128 shots). In particular, the setup used by Anil et al. 2024 could enable the authors to derive direct consequences of their work when it comes to model safety. In addition, the many-shot setting is also relevant for agentic workflow and RAG applications. Hence, including such evaluations could further enhance the applicability of the findings for practitioners.\n\n1. Unclear connection with Optimal Transport: the authors use the Sinkhorn algorithm, but, as someone without a background in optimal transport, it is unclear that this alone qualifies the method as learning an optimal transport map between the uniform distribution on permutations and the “distribution of hard permutations” (for which a precise definition is not given). It would be helpful if the authors could clarify this, e.g. in the appendix.\n\n1. Minor points about formatting:\n\n    - The use of the absolute value in Equation (1) confused me at first, as it seemed to suggest that the ASR metric also factors in cases where permuting the inputs increases performance. This is of course not the case, since the worst-case performance is never greater than the average case performance. Still, I believe it would improve clarity to remove the absolute value in Equation (1), since its argument is anyways always non-negative.\n\n    - In Algorithm 1, the authors refer to “ascending the gradient” of functions labeled $L_D$ and $L_G$. I found this confusing, as the naming of the functions seems to suggest these are loss functions, in which case one would descend their gradients.\n\n    - Also, the nomenclature for the functions in 4.3 is different from that of Algorithm 1. I believe it would help clarity if these matched.", "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": "N6jXsaXx3M", "reviewer": "Reviewer_cDtR", "summary": "The paper address the problem that the order of demonstrations in few-shot learning influences the task performance. The authors propose an adversarial optimisation training scheme, whereby an additional model mines hard permutations while fine-tuning the original model to perform well on them.", "questions": "1. The separator line between 2-shot and 3-shot in Table 3 is offsetted.\n2. Do you need a separate P-Net for every number of demonstrators? Eq. (4) seems to hint at that. If that's the case, I think the authors should discuss the implications of this approach on scalability and practical implementation.\n3. If one wants demonstrator-wise permutation invariance, can’t one just mask the examples in a way that no example can attend to another example, and thus making the model permutation-invariant by design?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper offers a simple, straightforward, intuitive and, most importantly, well-performing solution to the problem of permutation-sensitivity.\n2. The paper is clearly written and well-presented.\n3. It is not immediately obvious how one can handle the combinatorial explosion for the hard permutation mining with a neural network but the authors propose an elegant solution using the Sinkhorn operator and Gumbel sampling.", "weaknesses": "1. The evaluations seem to be restricted to the 3-, 4- and 5-shot cases. However, these settings are quite small and one could even enumerate them and try each permutation without needing the P-Net at all. Currently, models have very large contexts and can have hundreds if not thousands of demonstrators (e.g., Many-Shot In-Context Learning, Agarwal et al., 2024). It is not clear from the paper whether such larger sets of demonstrations also exhibit such permuntation-sensitivity. Furthermore, it is not clear whether the proposed approach can scale to the many-shot setting (or more than 5 examples) and what the benefit would be in that case. I think a key question that needs to be answered is whether the proposed method is better than simply adding extra demonstrators (in random order).\n2. The paper does not compare how the proposed method fares against baselines such as selecting the most performant permutation (without specially fine-tuning the model to be permutation-invariant) and selecting the demonstrators that result in maximal performance, e.g. the (Lu et al., 2022) reference from 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": "tJBorKJILs", "reviewer": "Reviewer_n5jf", "summary": "This paper considers the tendency for the performance of LLMs to decrease given changes to the ordering of in-context (IC) examples. The authors propose an approach based on distributionally robust optimization (DRO), which they then run on both synthetically generated data as well as standard benchmarks. They show that performance against adversarial modifications of the ordering of the IC examples tends to increase when fine-tuning using their DRO approach.", "questions": "* The authors consider average and worst-case orderings of IC examples. But what about best-case ordering? This also seems like a reasonable baseline. Given that performance significantly drops between the average and worst orderings, should we expect performance to increase given a best ordering?  This almost seems like a more realistic model of LLM use than the adversarial setting: A user wants to maximize their performance, and one way of doing that is by finding a good ordering of IC examples. \n* Another missing (easy) baseline is the impact of picking a *random* permutation in Figure 1. Or, is this how the average performance was calculated?\n* Line 171 -- is the upper bound \"simulated,\" or is this the actual upper bound? My understanding is that if you were to enumerate all of the permutations of the IC examples, you would be able to calculate the *true* upper bound.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The paper has a relatively strong set of experiments. If you agree that the threat model is worth studying, then you would certainly conclude that the method provides clear robustness against this kind of attack.\n* Aside from a few minor typos, the paper is well-written and easy to follow.", "weaknesses": "**Threat model**\n\n* \"This paper shows that this vulnerability can be exploited to design a natural and completely imperceptible attack. . .\" -- This claim seems untrue. The attack is perceptible; the attack changes the order of the IC examples, and a keen user could very well notice this. Or, if the argument is that the user never sees the IC examples, then of course it's imperceptible. But *any* attack that is hidden from the user by construction would be imperceptible, so this doesn't seem like a distinguishing aspect of the algorithm.\n* \"Consequently, there remains a significant need for methods that can fundamentally enhance LLMs’ inherent ability to manage the vast combinatorial space of possible input permutations.\" -- I'm not sure I'm convinced. Why is the need so pronounced, and in what realistic setting do you imagine that an adversary will have access to the IC examples? It seems quite unlikely that an adversary would have access to the same context window as a user. Moreover, even if the IC examples go in the system prompt, it's unclear in what settings an adversary would have this kind of access, particularly if one tends to use proprietary models, which tend not to expose their system prompts to users. Indeed, this accounts for the vast majority of use cases, so calling the need \"significant\" seems too strong. My advice: The authors should motivate their threat model and clearly explain why it corresponds to a real-world setting. Perhaps a demonstration of this attack happening in the wild, especially in the first few sections of this paper, would make sense.\n\n**Experiments.**\n* \"We select two tasks. . . We test 100 samples from each task.\" -- Is 100 samples sufficiently many? I'd be concerned that you may not get stable results (i.e., the variance would be relatively high) on datasets this small. Did you run multiple independent trials? How were these 100 examples selected?\n* It would be worth using more than one LLM, especially in the experiments described in Figure 3. The authors use Llama 3.1, which is indeed one of the strongest one-source models. It seems that one could also simulate this attack on proprietary models by putting these examples in the context window before the user's prompt. Is this the case, or have I misunderstood something?  I understand that the defense proposed later on requires fine-tuning the model, but even here, it would be worth extending the score of the analysis to include other open-source models to get a better sense for how this vulnerability impacts models outside of the Llama family.\n* The drop in performance seems to be a factor not only of $\\tau$, but also of the number of examples, since (as the authors point out), larger $n$ would tend to mean that the lower bound on performance is lower.\n* In eq. (1), how are $\\mu_i$ and $w_i$ calculated? Is this meant to be a general overview of the method, or something more specific to the experiments in Figure 1. Indeed, it may also be worth characterizing $R$ as a function of $D$ and $\\tau$, i.e., $R = R(D, \\tau)$.  \n* I didn't understand Figure 2. In the ERM problem in eq. (2), the optimization  $\\theta$ is learned, but the text indicates that both ERM and DRO learn \"probability distributions of models.\" This doesn't make sense to me, because neither (2) nor (3) deal with a distribution over models. For this reason, the plots in Figure 2 are hard to interpret; what is $\\hat{P}$ (as far as I know, this isn't defined in the text), is this different from $\\hat{p}$ (referred to in the caption), and what is $P_\\theta$? It would be worth explaining the example more as well. It's unclear why only indices 2,3, and 5 are considered here.  Moreover, the bars in both plots also look exactly the same, which I'd assume isn't the right message to take away.\n* The paper would benefit from a clearer explanation regarding how the hyperparameters are chosen.\n* It would be worth describing how the Rogue-L metric is calculated and why it's a good metric to use here.\n\n**Proposed method.**\n* It'd be worth saying more about why this instantiation of the problem is NP-Hard. Is it simply because I would need to enumerate all of the options in $\\mathcal{Q}$ (assuming that's what $\\mathcal{Q}$ is, since it's rigorously defined in the text)?\n* The [CLS] token seems to appear from nowhere. What does CLS stand for?\n* It would be worth describing more why the proposed approach corresponds to *adversarial* training. The approach seems to involve minimizing a certain loss function, rather than the typical zero-sum formulation of adversarial training which is standard in the computer vision/adversarial examples literature.\n* The problem is a tad bit underspecified. The authors should be precise about what the dimensionality of $X$ is and what $N$ is in the formulation on page 6.  It's also unclear why the authors want to add randomness (c.f., \"To introduce randomness and control...\"); why is randomness needed in this algorithm, and why, in particular, is the current method the right way to introduce it? Also, why do we need to \"control the discreteness?\" I am having a hard time understanding what this means and/or why the authors are trying to do this.\n* The authors say that their formulation allows the user to transform \"the input permutation distribution into any permutation distribution.\" Again, the paper should be more clear about whether we are dealing with distributions over perturbations or perturbations themselves. You'd imagine that there is a primal view of (3) where we deal over the objects themselves, rather than distributions over these objects. \n* In (8), it's unclear where the sinkhorn operator factors in. And how does this formulation connect to the DRO formulation in (3). From the perspective of the reader, this problem in (8) seems to appear out of thin air, and I'd imagine that readers will get stuck at this point.\n* Unless I missed something, the variable $G$ is undefined in Algorithm 1. Is that the P-Net?\n* \"To prevent the P_Net from exploiting trivial solutions--such as outputting even matrices that average out the semantic content...\" -- I'm not sure what this means. What is semantic content here? Isn't the P-Net outputting a permutation?\n* I didn't understand the steps between 9 and 10. Consider that if we plug in the losses to\n\n$$\\min_{\\theta,\\phi} L(\\theta){\\text{LM}} + L(\\phi)_{P-Net}$$\n\nthen objective becomes:\n\n$$ \\mathbb{E}[\\ell(\\theta; \\phi; (\\Pi \\cdot p, x, y)) - \\ell(\\theta, \\phi; (\\pi\\cdot p, x, y)) + \\beta L(\\theta, \\phi)_\\text{ent}] $$\n\nwhich is equal to \n\n$$ \\mathbb{E}[ \\beta L(\\theta, \\phi)_\\text{ent}] $$\n\nsince the first two terms cancel.  Could the authors clarify this, because it seems like the loss reduces to an entropy penalty.\n\n**Discussion of related work.**\n* Outside of the related work section, there is very little discussion of related work. It would be worth justifying claims using citations, e.g., \"Under appropriate assumptions, learning theory guarantees that models trained via ERM perform well on the test distribution given sufficient training data.\" Since this wasn't studied in the paper under review, it would be worth citing learning theoretic papers, which, if the reader had inclination, would enable looking into this matter more closely.\n* It would also be worth citing representative works from papers that seek to generalize to data that is outside of the training distribution. A particularly related piece of related work is the GroupDRO paper (https://openreview.net/pdf?id=ryxGuJrFvS), which uses a similar approach in classification to the approach outlined by the authors here.\n\n**Miscellaneous.**\n* \"Most existing studies on ICL primarily aim to enhance the general performance on few-shot learning.\" -- The paper would benefit from more precise language. \"General performance\" is fairly vague, and it's a bit unclear how one would distinguish this from the setting considered by the authors. Does this paper not also consider the \"general performance\" of ICL?\n* Figure 1 -- the colors are too difficult to tell apart. It would be worth considering colors outside of the blue family.", "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)