PEARL: Towards Permutation-Resilient LLMs
Under review as a conference paper at ICLR 2025

PEARL: TOWARDS PERMUTATION-RESILIENT LLMS

Anonymous authors
Paper under double-blind review

ABSTRACT

The in-context learning (ICL) ability of large language models (LLMs) enables
them to undertake challenging tasks using provided demonstrations. However,
it is prone to instability: different orderings of demonstrations can significantly
influence predictions, revealing LLMs’ limitations in processing combinatorial
inputs. This paper shows that this vulnerability can be exploited to design a natural
attack that is imperceptible to the model provider and can achieve nearly 80%
success rates on the SOTA open-source model, LLaMA, by simply permuting the
demonstrations. In light of this, how to overcome the ordering sensitivity problem
is an important issue for improving the performance of LLMs. However, current
mitigation methods focus on post-processing and fail to enhance models’ inherent
robustness to the vast space of possible input permutations. To overcome this issue,
we propose a novel Permutation-resilient learning framework (PEARL) based on
distributionally robust optimization (DRO), which optimizes model performance
against the worst case among all possible permutations. Specifically, PEARL
consists of a hard permutation mining network (P-Net) and the LLM. The P-
Net identifies the most challenging permutations by formulating the task as an
optimal transport problem, which is solved using an entropy-constrained Sinkhorn
algorithm. Through minimax optimization, the P-Net progressively generates
harder samples to enhance the LLM’s worst-case performance. Experiments with
synthetic data and instruction tuning tasks demonstrate that the PEARL framework
effectively mitigates permutation attacks and improves overall performance.

1

INTRODUCTION

A hallmark of human intelligence is the ability to learn and execute new tasks by reasoning from a
few examples. Mirroring this, in-context learning (ICL) (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.,
2021; Chowdhery et al., 2023; OpenAI et al., 2023). Despite these advancements, the ICL capabilities
of large language models (LLMs) remain fragile. LLMs exhibit sensitivity to permutations of
provided demonstrations (Lu et al., 2022; Zhao et al., 2021; Reynolds & McDonell, 2021). This
fragility underscores a significant gap in achieving human-like adaptability.

Most existing studies on ICL primarily aim to enhance the normal-case performance on few-shot
learning (Min et al., 2022; Wei et al., 2023), with limited attention to improving permutation
robustness. Current strategies addressing this issue in few-shot learning generally fall into two
categories: 1) Output Calibration (Zhao et al., 2021), which proves effective for classification
tasks but is less applicable to generation tasks, and 2) Order Optimization (Lu et al., 2022), which
focuses on finding the optimal sequence of few-shot demonstrations during inference but suffers from
exponential computational complexity. 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.

In this work, we first conduct extensive experiments on LLaMA-3 to revisit the vulnerability of
latest LLMs to permutations of ICL (§3). Our empirical analysis reveals that even state-of-the-art
open-source LLMs, such as LLaMA-3-8B, are still highly susceptible to a simple permutation-based
attack that merely alters the order of ICL demonstrations. Remarkably, these attacks, which do not
modify the semantic content of the examples or append any malicious suffixes, can achieve success

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

054
055

056
057
058
059
060
061

062
063
064
065
066
067

068
069
070
071
072

073
074
075
076
077
078

079
080
081
082
083
084

085
086
087
088
089
090

091
092
093
094
095

096
097
098
099
100
101

102
103
104
105
106
107

rates exceeding 80%. Consequently, these attacks are less noticeable to model providers but highly
effective against LLMs, highlighting a critical vulnerability of LLms.

To counteract the vulnerability to input permutations, we introduce a novel Permutation-resilient
learning (PEARL) framework, which is based on distributionally robust optimization (DRO) (Ben-Tal
et 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 a broader
distribution that includes all conceivable permutations. This comprehensive set of all possible
permutations is 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 a distribution of potential permutations— equips the model to better prepare
for and generalize to combinatorial input scenarios.

Specifically, PEARL operationalizes DRO as a two-player game, consisting of a hard permutation
mining network (P-Net) as the adversary and the LLM as the target model. For each training instance,
P-Net identifies a hard permutation of given demonstrations, aiming to maximize the LLM’s loss.
Conversely, the LLM strives to minimize its loss under the P-Net’s perturbations, thereby performing
well on these challenging examples. P-Net frames the identification of the most adversarial ICL
permutation as an optimal transport (OT) (Monge, 1781) problem between the uniform distribution
over permutations and the distribution of currently challenging permutations. We solve the OT
problem using the Sinkhorn algorithm (Sinkhorn, 1966) with an element-wise 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 widely used scenarios: (1) pre-training a transformer to in-context
learn linear functions, and (2) instruction finetuning of LLMs on real-word tasks. Comprehensive eval-
uations demonstrate that compared to ERM-based training, our method consistently and substantially
improves both the average and worst-case performance of LLMs across all possible permutations
and effectively defends against permutation-based attacks. Notably, in practical instruction tuning
scenarios, our method achieves superior results with only hundreds of LoRA parameter updates,
highlighting its exceptional effectiveness and efficiency.

2 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 training-stage
methods focus on improving general performance in ICL (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.,
2023; Xiang et al., 2024). InfoAC (Xiang et al., 2024) introduces contrastive learning during fine-
tuning to 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., 2023) 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.
Inference-stage methods can be categorized into four types: (1) demonstration selection (Chang & Jia,
2023; Peng et al., 2024), which primarily enhances normal-case performance without guaranteeing
worst-case performance under permutations; (2) output calibration (Zhao et al., 2021; Li et al., 2023;
Guo et al., 2024a), which proves effective for classification tasks but is less applicable to generation
tasks due to sequence calibration challenges; (3) order optimization (Lu et al., 2022), which aims to
find the best ordering during inference but suffers from exponential computational complexity; and
(4) prediction ensembling: a recent work (Zhang et al., 2024) proposes to transform an n-shot ICL
into n one-shot predictions and ensembles the results—this is effective for classification but leads
to decreased performance on generation tasks. In summary, In summary, inference-stage methods
aims to circumvent order sensitivity by pre/post-processing without fundamentally enhancing the
robustness of LLMs to different orders. Moreover, most methods are designed for classification tasks
and show reduced effectiveness on generation tasks. To the best of our knowledge, our work is the

2

Under review as a conference paper at ICLR 2025

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

Figure 1: Performance and attack success rates of Llama-3 on 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.

first to solve this problem from an adversarial perspective. We propose a novel distributionally robust
optimization (DRO)-based learning algorithm to enhance the inherent robustness of LLMs against
order perturbations and solve it using the Sinkhorn operator. Our approach complements existing
inference-stage methods and generalizes across diverse task categories.

Distributionally Robust Optimization. In distributionally robust optimization (DRO), ambiguity
sets are often defined as divergence balls centred on the empirical distribution of data pairs (x, y),
which act as regularizers for small radii (Ben-Tal et al., 2013; Lam & Zhou, 2015; Duchi et al.,
2016; Miyato et al., 2018). However, larger radii can result in excessively conservative sets. Prior
applications of DRO have 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). In contrast, this study is the
first to apply DRO to in-context learning robustness, defining the ambiguity set through all possible
permutations of the empirical distribution that requires ICL performance guarantees.

Optimal Transport. Optimal transport (OT), a foundational mathematical discipline established by
(Monge, 1781; Kantorovich, 1942), provides a metric for measuring distances between distributions,
commonly known as the Wasserstein distance or Earth Mover Distance. It has been applied as a
tool for manipulating probability distributions. In our study, the hard Permutation mining Network
(P-Net) is designed to act as a conduit for transportation between two discrete measures, leveraging
entropy-constrained OT (Cuturi, 2013), also referred to as the Sinkhorn distance, to enable the
derivation of a differentiable loss (Genevay et al., 2018). 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 Sinkhorn operator to generate challenging permutations for
LLMs to perform adversarial training.

3 REVISITING PERMUTATION VULNERABILITY IN LLMS

This section examines the severity of performance fluctuations in LLMs in response to different
permutations of given demonstrations. Additionally, from an adversarial perspective, we explore
whether this vulnerability can be exploited to devise an effective attack on LLMs.

Experimental Setups To conduct evaluations, we select two tasks from Super-NaturalInstructions
(Wang et al., 2022), including Curiosity-based 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 (Wang et al., 2022), the performance is measured using the ROUGE-L (Lin, 2004).
We adopt LLaMA-3-8B for evaluation due to its widespread use. 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 1, 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

3

65432Number of Shots0.20.30.40.50.6PerformanceCurDial65432Number of Shots0.20.30.40.5PerformanceTMW020406080100Threshold (%)6065707580859095100Attack Success Rate (%)CurDial020406080100Threshold (%)2030405060708090100Attack Success Rate (%)TMWAverageWorstRandomShot 4 (Exhaustive)Shot 5 (Exhaustive)Shot 6 (Exhaustive)Shot 4 (Neural)Shot 5 (Neural)Shot 6 (Neural)Under review as a conference paper at ICLR 2025

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

information. However, it can simultaneously worsen the worst-case performance. This suggests
that while more demonstrations provide beneficial context, the exponentially increasing number of
possible permutations (n!) introduces a higher likelihood of a possible input configuration on which
the model performs poorly.

2) Input Permutation as Attack We then consider a two-party adversarial scenario (Zou
et al., 2023; Rao et al., 2024; He et al., 2024) between a malicious user (attacker) and a model
provider (defender). The attacker seeks to induce compromised responses from LLMs solely by
permuting the ICL demonstrations, making the attack less noticeable to the model provider. We
measure the effectiveness of such attack by reporting the attack success rate (ASR). Given a task
D = {(pi, xi, yi)}, we define a sample (pi, xi, yi) successfully attacked if its relative performance
degradation induced by a attacher exceeds a threshold δ ∈ [0%, 100%]. Here, pi represents an
ICL prompt containing n demonstrations. We denote the set of all possible permutations of the pi
demonstrations as P = {Π0, . . . , Πn!−1}, where |P| = n!. Let g be a performance metric function
(e.g., ROUGE-L). The ASR on the task D is defined as:

ASR(D, δ) =

1
|D|

|D|
(cid:88)

i=1

I(cid:0)(µi − ωi)/µi ≥ δ(cid:1)

(1)

where I denotes the indicator function, |D| is the size of the dataset, and δ is the threshold. The
average performance of the i-th sample, µi, is defined by:

µi = EΠ∼P[g(Π · pi, xi; yi)] =

1
n!

n!
(cid:88)

j=1

g(Πj · pi, xi; yi)

(2)

and ωi is the compromised performance induced by the attack strategy adopted by the malicious user.
Here, we analyze two attack method:

• Exhaustive Search Attack: To calculate the upper bound of the effect the permutation-based attack
can achieve, we assume that the malicious user has unlimited attempts and conducts an exhaustive
search. For each sample (pi, xi, yi), this process involved testing all possible permutations of
demonstrations in Qi and identifying the permutation that yields the poorest performance. In this
case, the attacked performance is calculated as follows:

ωi = min
Π∈P

g(Π · pi, xi; yi)

(3)

• Neural Search Attack: To approximate the upper bound established by the exhaustive search
when the number of attempts is limited, we employ a meta-learning approach to optimize a
hard permutation mining network (P-Net). As illustrated in Figure 3 (details are in the Methods
section), during training, this network takes the standard sample (pi, xi, yi) as input and outputs a
permutation matrix Πi. The permuted samples (Πi · pi, xi, yi) are then fed into the LM to maximize
its loss function. During testing, the network generates the most challenging permutation Πi for
each sample (pi, xi, yi). Then the attacked performance is calculated as follows:

ωi = g(Πi · pi, xi; yi),

s.t. Πi ∼ P-Net(pi, xi, yi)

(4)

As shown in the right of Figure 1, 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 δ = 50% on two datasets respectively, and the neural attack
achieved a successful rate close to this upper bound across different shots. These results demonstrate
that this vulnerability poses a real concern, even for advanced LLMs like LLaMA-3.

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 behaviour on unseen but practically valid input spaces.

4 PERMUTATION-RESILIENT LEARNING (PEARL)

4.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.

4

Under review as a conference paper at ICLR 2025

In supervised fine-tuning for few-shot learning, the LLM is trained to predict an output y ∈ Y given
an input x ∈ X and a few-shot instruction p ∈ P, where p typically consists of a sequence of
demonstrations, each being an input-output pair. Let Θ denote the parameter space of the language
model, and let ℓ : Θ × (P × X × Y) → 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
θ ∈ Θ that minimize the empirical loss over the training data via empirical risk minimization (ERM):

ˆθERM := arg min
θ∈Θ

E

(p,x,y)∼ ˆP [ℓ(θ; (p, x, y))]

(5)

where ˆP denotes the empirical distribution derived from the training dataset.

Under appropriate assumptions, learning theory (Vapnik, 1999; 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
a subset of all possible permutations of the demonstrations, and during testing, the model may
encounter permutations not seen during training, leading to a significant degradation in performance.

To systematically address the permutation sensitivity issue, we propose fine-tuning under the frame-
work of distributionally robust optimization (DRO), which optimizes the risk under the worst-case
distribution within a specified ambiguity set. Specifically, we aim to solve:

ˆθDRO = arg min
θ∈Θ

(cid:110)

sup
QΠ∈Q

(cid:111)
E(p,x,y)∼QΠ [ℓ(θ; (p, x, y))]

(6)

The ambiguity set Q is constructed to capture all distributions obtained by permuting the prompts
in the empirical distribution ˆP . Specifically, for each possible permutation Π ∈ P, we define the
permuted distribution QΠ by applying Π to the prompt p of each data point in ˆP :
(cid:110)(cid:0)Π · p, x, y(cid:1) (cid:12)
(cid:12) (p, x, y) ∼ ˆP
(cid:12)

, Π ∈ P,

QΠ :=

(7)

(cid:111)

where Π is a permutation matrix acting on the sequence of demonstrations in p, and P denotes the set
of all possible permutation matrices. The ambiguity set Q is then defined as the convex hull of these
permuted distributions:

(cid:26) (cid:88)

Q :=

qΠ QΠ

(cid:12)
(cid:12)
(cid:12) q ∈ ∆|P|−1

(cid:111)
,

(8)

Π∈P
where q is a probability vector belonging to the |P| − 1-dimensional simplex ∆|P|−1.
By considering all possible permutations of the prompts in the empirical distribution, Q encompasses
all distributions that could arise due to prompt permutations. This formulation allows DRO to identify
the worst-case distribution within Q (the sup step in Eq. 6) and optimize the model’s performance
against it (the arg min step), thereby enhancing robustness to permutations in the input data.

To illustrate the advantages of DRO over
ERM in handling different permutations,
consider the example in Figure 2. For a 3-
shot training example (p, x, y) with prompt
p containing three demonstrations, there
are six possible permutations denoted as
(p0, x, y), . . . , (p5, 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 0, 1, and 4 appear in training
data with frequencies, while permutations
2, 3, and 5 do not appear. Pθ represents
the distribution learned by the LLM, repre-
sented by purple curves. In panel (a), the
ERM-trained model assigns higher proba-
bilities to frequently occurring permutations (0, 1, 4) and lower probabilities to less frequent ones
(2, 3, 5), leading to poor performance on unseen permutations during testing. In contrast, panel

Figure 2: Comparison of models trained under ERM
and DRO paradigms. The blue bars represent the em-
pirical distribution ˆP of training data, showing different
frequencies of six permutations in the training set. The
purple curves denote the learned distribution Pθ by (a)
ERM and (b) DRO models, illustrating their different
behaviors on less appeared but valid permutations.

5

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

012345Permutation IndexP(a) ERM012345Permutation Index(b) DROPUnder review as a conference paper at ICLR 2025

(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 limitations by encouraging models to assign reasonable probabilities to
all valid permutations, regardless of their frequency in training data.

4.2 P-NET: LEARNING TO PERMUTE VIA OPTIMAL TRANSPORT

To enable our DRO framework to function effectively, we 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.

To overcome this challenge, we model the problem as the Otimal Transport (OT) from the distribu-
tion of permutations of the input data to a target distribution that is challenging for the current LLMs.
To implement this, we design a neural network called the Hard Permutation Mining Network
(P-Net), P-Net: (P × X × Y) → ∆(Π), which maps an input example to a distribution of challenging
permutations. As illustrated in Figure 3, we can sample a hard permutation from this distribution to
reorder the demonstrations into a more challenging version.

The P-Net consists of two components: a parameter part that extracts features and models the
relationships between demonstrations, a non-parameter part using the Sinkhorn algorithm to build
the distribution ∆(Π), and Gumbel sampling for differentiable sampling from it (Π ∼ ∆(Π)).

Parameter component. The parameter component consists of a feature extractor and a cross-
relationship modeling layer. The feature extractor can be a small pre-trained model that takes an ICL
prompt composed of n demonstration pairs p = {(xi, yi)}n
i=1 and a predicting sample (x, y), and
produces their representations as follows:

([CLS], (x1, y1), . . . , [CLS], (xn, yn), [CLS], (x, y)) Transformer

−−−−−−→ (h1, h2, . . . , hn, hn+1) ,

(9)

where hi is the representation corresponding to the i-th [CLS] token, which is often used to segment
and extract the representation of sequences (Devlin et al., 2019b; Lu et al., 2021).
After extracting the representations of n demonstrations, we have H = (h1, h2, . . . , hn) ∈ Rn×h.
We then model the pairwise relationships among the demonstrations. Specifically, we design a simple
cross-demonstration layer to obtain a relationship matrix R ∈ Rn×n that captures the pairwise
relationships between each pair of demonstrations, defined as:

R = g

HW H ⊤(cid:17)
(cid:16)

,

(10)

where W ∈ Rh×h is a weight matrix, and g denotes a nonlinear activation function.
The output matrix R ∈ Rn×n can be interpreted as an adjacency matrix in graph theory. Viewing
the demonstrations as nodes in a graph, the relationship between nodes i and j is represented by the
edge Rij. Here, we define Rij as the potential increase in difficulty for the LLM if demonstrations
i and j are swapped; a higher value of Rij indicates that swapping these two demonstrations may
significantly increase the task’s difficulty.

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 ∆(Π) that we can
sample from, we introduce a non-parameter component that employs the Sinkhorn operator.

Non-parameter component. As Sinkhorn operator S is a well-established method in optimization
theory that transforms a square matrix into a doubly stochastic matrix—also known as the Sinkhorn
distribution—which represents a distribution over permutations (Sinkhorn, 1966; Adams & Zemel,
2011; Mena et al., 2018), we can use it to transform R into the distribution of permutations ∆(Π).
We implement the sinkhorn algorithm through simple iterative process:

S(R) = lim
l→∞
Tr(R) = R ⊘ (cid:0)R1n1⊤

(Tc (Tr (exp(R)))) ,
(cid:1) ,

Tc(R) = R ⊘ (cid:0)1n1⊤

n

n R(cid:1) ,

(11)

(12)

where Tr(R) and Tc(R) represent the row and column normalization operators, respectively; ⊘
indicates element-wise division; and 1n is a column vector of ones. As established by (Sinkhorn,

6

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

Under review as a conference paper at ICLR 2025

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

Figure 3: An overview of the learning framework. The P-Net is a small model incorporating optimal
transport (OT) algorithm, trained jointly with the LLM under the adversarial optimization paradigm.
Note that the permutation matrix operates on the input sequence’s embeddings (simplified here as
text sequences for clarity). After training, only the LLM is retained while the P-Net is discarded.

1966), the Sinkhorn operator S(R) strictly converges to a doubly stochastic matrix as the number of
iterations l approaches infinity.

As we need to sample a permutation matrix from the Sinkhorn distribution (doubly stochastic matrix)
(sample step in Figure 3) and build a differentiable process, Gumbel sampling (Jang et al., 2017) is
applied to the Sinkhorn operator:

Π = S

(cid:19)

(cid:18) R + U
τ

, Uij = − log (cid:0)− log (cid:0)U ′

ij

(cid:1)(cid:1) , U ′

ij ∼ Uniform(0, 1),

(13)

where U ∈ Rn×n is a matrix of Gumbel noise and τ is the temperature. As τ approaches zero,
S ((R + U )/τ ) approximates a permutation matrix Π ∈ Pn×n. The hyperparameters of the Sinkhorn
operator are studied in Appendix D.

By modeling permutation generation as an optimal transport problem and designing the P-Net
to implement it, we enable the transformation of the input permutation distribution into a target
permutation distribution. Next, we introduce how P-Net is co-optimized with the LLM to make the
target permutation distribution the most challenging for the current LLMs.

4.3 ADVERSARIAL OPTIMIZATION

As depicted in Figure 3, our framework employs an adversarial approach to co-optimize the LLMs
and the P-Net. Specifically, for each sample, the P-Net generates a challenging permutation designed
to maximize the LLM’s loss. In turn, the LLM seeks to minimize its loss despite the challenging
permutations introduced by the P-Net. Let θ denote the parameters of the LLM, and ϕ those of the
P-Net. We formalize the optimization process as follows.

We first optimize the P-Net, corresponding to the inner maximization step in Equation (6). For a
given example (p, x, y), we sample a permutation Π ∼ P-Net(ϕ; (p, x, y)) from P-Net. We then
compute the LLM’s loss on the permuted example (Π · p, x, y), denoted by ℓ(θ; ϕ; (Π · p, x, y)). The
objective is to optimize the P-Net parameters ϕ to maximize this loss:

L(ϕ; θ)lm = E

(p,x,y)∼ ˆP ,Π∼P-Net(ϕ;(p,x,y))[ℓ(θ; ϕ; (Π · p, x, y))]

(14)

Note that the Sinkhorn operator is implicitly included in Π ∼ P-Net(ϕ; (p, x, y)).

To prevent the P-Net from exploiting trivial solutions, such as outputting uniform matrices that dilute
the semantic content of the demonstrations, we introduce an element-wise entropy constraint term
that encourages Π to be as distinct as possible:

L(ϕ)ent = E

(p,x,y)∼ ˆP ,Π∼P-Net(ϕ;(p,x,y))

(cid:88)

i,j

Πij(1 − Πij).

This leads to the following combined optimization for the P-Net:

ˆϕ⋆ = arg max
ϕ∈Φ

(L(ϕ; θ)lm − βL(ϕ)ent) ,

(15)

(16)

7

Under review as a conference paper at ICLR 2025

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

where β represents the penalty coefficient for the entropy constraint, further studied in Appendix D.
Note when optimizing Equation (16), θ remains constant.

We then optimize LLM, corresponding to the inner minimization step in Equation (6). For an
example (p, x, y), we get a challenging permutation from the previously optmized P-Net ( ˆϕ⋆), Π ∼
P-Net( ˆϕ⋆; (p, x, y)). We compute the LLM’s loss on this permuted example (Π · p, x, y), denoted by
ℓ(θ; ˆϕ⋆; (Π · p, x, y)). The objective is to optimize the LLM parameters θ to minimize this loss:

ˆθ⋆ = arg max
θ∈Θ

L( ˆϕ⋆; θ)lm,

(17)

Note when optimizing LLM, we incorporate the previously optimized parameter ˆϕ⋆ from the
P-Net and keep it constant.
From Equation (16) to (17), we complete a loop of iteration. In the next iteration, we substitute ˆθ⋆
into Equation (16) for a new round of optimization until convergence. The comprehensive training
algorithm is outlined in Appendix A.

5

IN-CONTEXT LEARNING WITH LINEAR FUNCTIONS

5.1 DATASETS AND EVALUATION METRICS

We investigate in-context learning on linear functions f (x) = w⊤x, where w ∈ Rd, follow-
ing (Garg et al., 2022; Guo et al., 2024b). For each w, we construct each example pi =
(x1, f (x1), . . . , xi, f (xi), xi+1) containing i input-output demonstration pairs and a query input
xi+1. A language model LMθ is trained to minimize:

Ep

min
θ

(cid:20) 1
k + 1

(cid:88)k

i=0

ℓ(LMθ(pi), f (xi+1))

(cid:21)

,

(18)

where ℓ(·) is the MSE loss and k is the maximum number of demonstrations. We evaluate using
normalized squared error ((LMθ(p) − w⊤xquery)2/d). Detailed settings are in Appendix B.1.

5.2

IMPLEMENTATION DETAILS AND BASELINES

Architecture and Training. We implement Lθ using a GPT-2 base model (Radford et al., 2019)
and train it from scratch on a 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 as a BERT-base (Devlin et al., 2019a) and also trained from scratch.
Implementation details are in Appendix B.2.
Baselines. Consistent with (Garg et al., 2022), we adopt an empirical risk minimization method
with curriculum learning (Bengio et al., 2009; Wu et al., 2020) (ERM+CL) to train the model. The
training process gradually increase the number of demonstrations presented to the model, allowing
for progressive learning of more complex patterns and making the training more stable.

5.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.

As shown in Table 1, the performance gap between average and worst-case performance across
permutations for the baseline methods was significant, indicating substantial vulnerability to 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% at
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% at 5 shots.

Figure 4 depicts the proportion of successfully attacked samples in terms of (1) different attack
success thresholds and (2) number of demonstrations (shots). The former considers more pessimistic

8

Under review as a conference paper at ICLR 2025

Table 1: Normalized MSE across permutations.

Shot Method

3

4

5

ERM+CL
PEARL

ERM+CL
PEARL

ERM+CL
PEARL

Avg.

1.45

Worst.

2.67

0.86 (+40.7) 0.92 (+65.5)

1.20

3.34

0.79 (+34.1) 1.11 (+66.8)

1.28

5.03

0.87 (+32.0) 1.33 (+73.6)

Figure 4: Comparison of attack success rates.

scenarios (attacked samples drop a large margin), while the latter examines larger input spaces. We
observed that PEARL’s advantage increased as the threshold grew. At δ > 50%, the defence success
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 a large threshold).
Moreover, PEARL’s performance improved with an increasing number of shots, suggesting better
scalability compared to baseline methods.

6

INSTRUCTION FINE-TUNING OF LARGE LANGUAGE MODELS

6.1 EXPERIMENTAL SETUPS

Datasets. Our instruction tuning data were derived from Super-Natural Instructions (Wang et al.,
2022). We selected 17 representative tasks: 9 natural language generation (NLG) tasks and 8 natural
language understanding (NLU) tasks. Following (Wang et al., 2022), we randomly designated 4
datasets as held-out test sets and used the remaining 13 datasets for training. This resulted in a
training set of 1,950 examples and a test set of 400 examples (see Appendix C.1 for details). Each
example follows a format of (p, x, y), where p is an ICL prompt containing 2 to 5 demonstrations.
Evaluation Metrics. Following the practice in Super-Natural Instructions (Mishra et al., 2022; Wang
et 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 a single "average" metric
across all datasets, following the methodology in FLAN (Wei et al., 2022; 2023).
Baseline and Implementation Details. To evaluate our framework, we compare it with other
learning algorithms, including Empirical Risk Minimization (ERM) (Min et al., 2022), ERM with
Demonstration Shuffling (ERM+DS) (Zhang et al., 2018), ERM with Instance Mixup (ERM+IM)
(Zhang et al., 2018), InfoAC (Xiang et al., 2024), Batch-ICL (Zhang et al., 2024) and CurStable
(Chang & Jia, 2023). We use FLAN-large as the P-Net and experiment with five LLMs: Llama3-8B,
Llama2-7B/13B, Mistral-7B, Gemma-7B. More details on baselines and implementations are in
Appendix C.2. Hyperparameters are in Appendix C.3.

6.2 EVALUATION RESULTS

We evaluate PEARL from three perspectives: (1) comparison with training-stage methods (empirical
risk minimization [ERM] and its augmentations, InfoAE), (2) comparison and integration with
inference-stage techniques (Batch-ICL, CurStable), and (3) scalability to many-shot in-context
learning (ICL; Agarwal et al., 2024) with more demonstrations.

Table 2 presents the comparative performance of various 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 2 shots
to 29.4% at 4 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 LLaMA3-8B’s fine-tuning, where training loss plateaus within one
epoch. The rapid convergence suggests that focusing on challenging permutations during training is
more effective than using random ones—an observation consistent with WizardLM (Xu et al., 2024).

Our method demonstrates substantial improvements over strongest training-stage and inference-stage
baselines, achieving 9–21% worst-performance gains. On inference-stage methods, Batch-ICL boosts

9

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

020406080100Threshold (%)020406080100Attack Success Rate (%)ERM+CL shot 3ERM+CL shot 4ERM+CL shot 5PEARL shot 3PEARL shot 4PEARL shot 5Under review as a conference paper at ICLR 2025

Table 2: Average and Worst-Case Performance of Llama3-8B on 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.

# Shot Method

Avg.

Worst.

Avg. Worst. Avg. Worst. Avg. Worst. Avg. Worst.

Average

CSQA

CurDial

CoLA

TMW

2

3

4

ERM
ERM+DS
ERM+IM
INFOAC

CURSTABLE
BATCH-ICL

57.3
57.5 (-0.2)
53.5 (-6.6)
55.7 (-2.9)

61.6 (+7.5)
58.6 (+2.2)

49.4
48.6 (-1.6)
44.4 (-10.1)
47.6 (-3.7)

52.1 (+5.4)
-

PEARL
+ CURSTABLE
+ BATCH-ICL

62.9 (+9.8)
65.6 (+14.5)
-

56.4 (+14.2)
58.0 (+17.4)
-

ERM
ERM+DS
ERM+IM
INFOAC

CURSTABLE
BATCH-ICL

57.8
56.1 (-2.9)
55.3 (-4.3)
56.3 (-2.6)

61.0 (+5.4)
58.6 (+1.3)

38.3
39.7 (+3.7)
39.8 (+3.9)
39.5 (+3.1)

41.4 (+8.0)
-

PEARL
+ CURSTABLE
+ BATCH-ICL

63.1 (+9.2)
65.0 (+12.5)
-

46.9 (+22.5)
48.9 (+27.5)
-

ERM
ERM+DS
ERM+IM
INFOAC

CURSTABLE
BATCH-ICL

PEARL
+ CURSTABLE
+ BATCH-ICL

59.7
57.7 (-3.4)
56.0 (-6.2)
58.6 (-1.8)

60.8 (+1.8)
58.5 (-2.0)

63.1 (+5.7)
65.0 (+8.8)
-

30.6
31.8 (+3.9)
32.4 (+5.9)
33.0 (+7.8)

32.3 (+5.6)
-

39.6 (+29.4)
41.4 (+35.1)
-

58.0
62.0
63.0
57.5

64.0
63.0

65.0
68.0
65.5

57.7
60.0
59.0
59.3

65.0
62.0

68.4
70.0
68.7

61.3
63.3
63.2
63.7

63.0
62.0

68.4
70.6
69.0

54.0
54.0
54.0
56.0

56.0
-

62.0
63.0
-

47.0
46.0
46.0
49.0

52.0
-

62.0
64.0
-

38.0
40.0
42.0
44.0

40.0
-

52.0
54.0
-

57.9
54.1
44.7
53.4

61.7
56.3

60.3
64.6
-

61.4
54.1
54.6
55.2

62.5
59.6

66.7
67.6
-

62.9
57.3
53.7
58.7

64.5
61.5

69.2
72.3
-

43.4
37.8
28.1
36.4

46.2
-

50.7
52.8
-

25.9
25.4
28.0
24.3

26.7
-

34.8
35.8
-

21.3
17.6
17.8
19.0

22.8
-

31.3
34.2
-

62.0
61.0
57.0
63.0

68.4
65.5

71.0
74.0
72.0

61.9
60.0
57.6
62.1

64.0
64.0

64.7
68.4
65.6

63.3
60.1
57.6
63.9

64.1
63.3

64.7
66.3
65.0

58.0
60.0
56.3
61.5

62.0
-

68.0
70.0
-

52.0
56.0
53.1
55.8

54.0
-

56.0
58.0
-

45.8
52.0
48.5
51.0

48.0
-

52.0
54.0
-

51.1
51.5
49.4
48.7

52.3
49.6

55.1
55.9
-

50.3
50.3
50.0
48.4

52.3
48.7

52.4
54.1
-

51.1
49.9
49.6
48.1

51.5
47.2

50.1
50.6
-

42.0
42.7
39.2
37.3

44.1
-

44.8
46.2
-

29.4
31.5
31.9
28.8

32.7
-

34.7
37.6
-

17.5
17.8
21.3
17.0

18.4
-

23.0
23.2
-

both average and worst-case performance on classification tasks (CSQA, CoLA); however, it exhibits
limited or negative effects on generation tasks (CurDial, TMW), limiting applicability. In contrast,
CurStable performs well on both task types via demonstration selection. Notably, when combined
with inference-time methods, our approach yields additional performance improvements of 3–5%,
highlighting the complementary nature of our method.

We evaluate PEARL and ERM in the many-shot ICL set-
ting. As shown in Fig. 5, PEARL achieves surprising
worst-case performance gains from 24% to 40% when gen-
eralizing to larger shots, despite being trained with fewer
shots and shorter sequences. This suggests our method
helps LLMs learn robust features that generalize well to
many-shot ICL. Detailed results are in Appendix. F.

Analyses of hyperparameters and extended experiments
on LLaMA2, Mistral, and Gemma are provided in Appen-
dices D and E, respectively.

7 CONCLUSION

Figure 5: Scaling to many-shot ICL.

We introduced a novel permutation-resilient learning framework (PEARL) to enhance the robustness
of LLMs against different permutations. PEARL employs a hard Permutation mining Network (P-Net)
that utilizes the Sinkhorn algorithm to generate challenging permutations, combined with adversarial
training to systematically improve LLM performance. Through empirical evaluations in both the
synthetic ICL task and the instruction tuning task, our framework has proven effective in mitigating
attacks and enhancing the generalization of LLMs. This research addresses a significant vulnerability
in LLMs, setting a foundation for the development of more resilient future language models.

10

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

Under review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

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/abs/2404.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. URL https://api.semanticscholar.
org/CorpusID:761793.

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. URL

https://arxiv.org/abs/2212.10378.

Yongqiang Chen, Binghui Xie, Kaiwen Zhou, Bo Han, Yatao Bian, and James Cheng. Positional
information matters for invariant in-context learning: A case study of simple function classes,
2023. URL https://arxiv.org/abs/2311.18194.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam
Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh,
Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam
Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James
Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev-
skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin
Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret
Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick,
Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Er-
ica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang,
Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern,
Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language model-
Journal of Machine Learning Research, 24(240):1–113, 2023. URL
ing with pathways.
http://jmlr.org/papers/v24/22-1144.html.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi 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, Quoc V. Le, and Jason Wei. Scaling instruction-finetuned language models. Journal of
Machine Learning Research, 25(70):1–53, 2024. URL http://jmlr.org/papers/v25/
23-0870.html.

11

Under review 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, Reiichiro Nakano, Christopher Hesse, and John
Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv.org/
abs/2110.14168.

Marco Cuturi.

In
Sinkhorn distances: Lightspeed computation of optimal
C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger (eds.), Ad-
vances in Neural Information Processing Systems, volume 26. Curran Associates, Inc.,
URL https://proceedings.neurips.cc/paper_files/paper/2013/
2013.
file/af21d0c97db2e27e13572cbf59eb343d-Paper.pdf.

transport.

Jacob Devlin, Ming-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 1 (Long and
Short Papers), pp. 4171–4186, Minneapolis, Minnesota, June 2019a. Association for Computational
Linguistics. doi: 10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding, 2019b. URL https://arxiv.org/
abs/1810.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? a case study of simple function classes. In Advances in Neural Information Processing
Systems, volume 35, pp. 30583–30598. Curran Associates, Inc., 2022.

Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning generative models with sinkhorn
In Amos Storkey and Fernando Perez-Cruz (eds.), Proceedings of the Twenty-
divergences.
First International Conference on Artificial Intelligence and Statistics, volume 84 of Pro-
ceedings of Machine Learning Research, pp. 1608–1617. PMLR, 09–11 Apr 2018. URL
https://proceedings.mlr.press/v84/genevay18a.html.

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 Linguistics. doi: 10.18653/v1/2024.
findings-acl.884. URL https://aclanthology.org/2024.findings-acl.884.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How
do transformers learn in-context beyond simple functions? a case study on learning with repre-
sentations. In The Twelfth International Conference on Learning Representations, 2024b. URL
https://openreview.net/forum?id=ikwEDva1JZ.

Pengfei He, Han Xu, Yue Xing, Hui Liu, Makoto Yamada, and Jiliang Tang. Data poisoning for

in-context learning, 2024. URL https://arxiv.org/abs/2402.02160.

Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang,
In International
and Weizhu Chen. LoRA: Low-rank adaptation of large language models.
Conference on Learning Representations, 2022. URL https://openreview.net/forum?
id=nZeVKeeFYf9.

W. Hu, G. Niu, I. Sato, and M. Sugiyama. Does distributionally robust supervised learning give

robust classifiers? In International Conference on Machine Learning (ICML), 2018.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In
International Conference on Learning Representations, 2017. URL https://openreview.
net/forum?id=rkE3y85ee.

Leonid Kantorovich. On the transfer of masses. In Doklady Akademii Nauk, volume 37, pp. 227–229,

1942.

12

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

Under review as a conference paper at ICLR 2025

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

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
In Andreas Vlachos and Isabelle Augenstein (eds.), Findings of the
to in-context learning.
Association for Computational Linguistics: EACL 2023, pp. 1385–1397, Dubrovnik, Croatia, May
2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.findings-eacl.102. URL
https://aclanthology.org/2023.findings-eacl.102.

Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summarization
Branches Out, pp. 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics.
URL https://aclanthology.org/W04-1013.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https:

//arxiv.org/abs/1711.05101.

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. 2780–2791, Online and Punta Cana, Dominican Republic, November
2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.emnlp-main.220. URL
https://aclanthology.org/2021.emnlp-main.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.18653/v1/2022.acl-long.
556. URL https://aclanthology.org/2022.acl-long.556.

Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permutations
with gumbel-sinkhorn networks. In International Conference on Learning Representations, 2018.

Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetaICL: Learning to learn
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 Linguistics: Human Language Technologies, pp. 2791–2809, Seattle, United States,
July 2022. Association for Computational Linguistics.

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
Computational Linguistics (Volume 1: Long Papers), pp. 3470–3487, Dublin, Ireland, May
2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.acl-long.244. URL
https://aclanthology.org/2022.acl-long.244.

T. Miyato, S. Maeda, S. Ishii, and M. Koyama. Virtual adversarial training: a regularization method
for supervised and semi-supervised learning. IEEE Transactions on Pattern Analysis and Machine
Intelligence, 2018.

Gaspard Monge. Memoire sur la théorie des déblais et des remblais. Histoire de l’Académie Royale

des Sciences de Paris, 1781.

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, Valerie 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

13

Under review as a conference paper at ICLR 2025

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

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 Gontijo-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, Łukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish
Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Hendrik
Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Łukasz 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, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto,
Michael, Pokorny, Michelle Pokrass, Vitchyr Pong, Tolly Powell, Alethea Power, Boris Power,
Elizabeth Proehl, 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 Lun-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.18653/v1/2024.acl-long.492.
URL https://aclanthology.org/2024.acl-long.492.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language

models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.

Abhinav Rao, Sachin Vashistha, Atharva Naik, Somak Aditya, and Monojit Choudhury. Tricking
llms into disobedience: Formalizing, analyzing, and detecting jailbreaks, 2024. URL https:
//arxiv.org/abs/2305.14965.

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: 10.1145/3411763.3451760. URL https://doi.org/
10.1145/3411763.3451760.

14

Under review as a conference paper at ICLR 2025

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust
neural networks. In International Conference on Learning Representations, 2020. URL https:
//openreview.net/forum?id=ryxGuJrFvS.

Shai Shalev-Shwartz 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/CJM-1966-033-9.

Vladimir Vapnik. The Nature of Statistical Learning Theory. Springer: New York, 1999.

Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Anjana
Arunkumar, Arjun Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, Eshaan Pathak,
Giannis Karamanolakis, Haizhi Gary Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby
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-naturalinstructions: Generalization via declarative instructions on 1600+
nlp tasks, 2022. URL https://arxiv.org/abs/2204.07705.

Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du,
Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International
Conference on Learning Representations, 2022. URL https://openreview.net/forum?
id=gEZrGCozdqR.

Jerry Wei, Le Hou, Andrew Lampinen, Xiangning Chen, Da Huang, Yi Tay, Xinyun Chen, Yifeng Lu,
Denny Zhou, Tengyu Ma, and Quoc Le. Symbol tuning improves in-context learning in 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/v1/2023.emnlp-main.61. URL
https://aclanthology.org/2023.emnlp-main.61.

Xiaoxia Wu, Ethan Dyer, and Behnam Neyshabur. When do curricula work?

arXiv preprint

arXiv:2012.03107, 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.findings-acl.386. URL https://aclanthology.org/2024.findings-acl.
386.

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 https://openreview.net/forum?id=CfXh93NDgH.

Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical

risk minimization, 2018. URL https://arxiv.org/abs/1710.09412.

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: ACL 2024, pp. 10728–10739, Bangkok,
Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.
findings-acl.638. URL https://aclanthology.org/2024.findings-acl.638.

Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: 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 https://proceedings.
mlr.press/v139/zhao21c.html.

15

Under review as a conference paper at ICLR 2025

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

for t = 1, . . . , m do

Algorithm 1 Adversarial Training Algorithm
1: Input: Training data ˆP , LLM θ, P-Net ϕ, P-Net iteration step m, Entropy coefficient β.
2: repeat
3:
4:
5:
6:
7:
8:
9:
10:
11: until convergence
12: Output: Optimized parameters θ and ϕ.

Sample an example (p, x, y) from ˆP .
Generate a permutation Π using P-Net: Π ∼ P-Net(ϕ; p, x, y).
Compute the LLM loss on the permuted input (Π · p, x, y): L(ϕ; θ)lm.
Compute the entropy regularization term L(ϕ)ent
Update P-Net parameters ϕ by ascending the gradient ∇ϕL(ϕ; θ)lm − β∇ϕL(ϕ)ent.

end for
Update LLM parameters θ by descending its gradient ∇θL(ϕ; θ)lm.

Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J. Zico Kolter, and Matt Fredrikson. Universal

and transferable adversarial attacks on aligned language models, 2023.

APPENDIX

A ADVERSARIAL TRAINING ALGORITHM

We present the adversarial training algorithm in Table A.

B DETAILED SETUP OF ICL WITH LINEAR FUNCTIONS

B.1 DATASETS CONSTRUCTION

We investigate training a language model to perform in-context learning on linear functions, following
(Garg et al., 2022; Guo et al., 2024b). The function class is defined as F = {f | f (x) = w⊤x, w ∈
Rd}, where d is the input dimension. Each data sample is constructed as follows:

(a) Function sampling: A weight vector w ∼ N (0, Id) is sampled, defining a linear function
f (x) = w⊤x.

(b) Input sampling: Inputs x1, x2, . . . , xk+1 ∼ N (0, Id) are independently drawn.
(c) Output generation: For each input, the corresponding output is computed as yi = f (xi) = w⊤xi
for i = 1, 2, . . . , k + 1.
The input prompt pi consists of i demonstrations and the (i + 1)-th example as the query: pi =
(x1, f (x1), x2, f (x2), ..., xi, f (xi), xi+1). We trained a language model Lθ, parameterized by θ, to
minimize the expected loss over all input prompts:

Ep

min
θ

(cid:20) 1
k + 1

(cid:88)k

i=0

ℓ(LMθ(pi), f (xi+1))

(cid:21)

,

(19)

where l(·) is the mean squared error (MSE) loss. During testing, we evaluated performance using the
same MSE metric. We report the normalized squared error ((LM (p) − w⊤xquery)2/d), where d is
the problem dimension.

B.2

IMPLEMENT DETAILS

Architecture. Following (Garg et al., 2022), we implement Lθ using a GPT-2 architecture (Radford
et al., 2019) with 12 layers, 8 attention heads, and a hidden dimension of 256. The model takes as
input a 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 on a generated dataset of 40k linear functions using
the AdamW (Loshchilov & Hutter, 2019). We employ a batch size of 128 and trained for 500k steps,

16

Under review as a conference paper at ICLR 2025

Table 4: Details of datasets used in instruction tuning from natural instructions.

Task ID Task Name

Source

Category

1297
442
908
288
582
151
1714
379
639
209
1516
589
1285

Question Answering
Question Rewriting
Speaker Relation Classification
Title Generation

QASC

QASC Question Answering
COM_QA Paraphrase Question Generation COM_QA
DialogRE
DialogRE Identify Familial Relationships
Gigaword
Gigaword Summarization
Natural Questions Question Answering
Natural Questions Answer Generation
Question Answering
TOM_QA
TOMQA Find Location Easy Clean
Dialogue Generation
ClariQ
ConvAI3 Sentence Generation
Text Categorization
AG News
AGNews Topic Classification
Dialogue Generation
MultiWOZ 2.2
MultiWOZ User Utterance Generation
Stance Detection
StarCon
Stance Detection Classification
Textual Entailment
IMPPRES
IMPPRES Natural Language Inference
Summarization
Amazon Reviews
Amazon Food Summary Text Generation
Text Matching
ArgKP
KPA Keypoint Matching

selecting the best checkpoint based on validation set performance. In the PEARL framework, we
randomly initialize the P-Net with a BERT-base-sized transformer encoder, also pre-training it from
scratch. During testing, we sample novel functions to assess the model’s ability to infer new weights
w through in-context demonstrations.

C DETAILED SETUP OF INSTRUCTION FINE-TUNING

C.1 DETAILS OF DATASETS

Table 3: Summary of datasets used
in instruction tuning.

Our instruction tuning data are derived from Super-Natural In-
structions (Wang et al., 2022), which are part of the FLAN v2
benchmark (Chung et al., 2024). We selected 17 representative
tasks, comprising 9 natural language generation (NLG) tasks
and 8 natural language understanding (NLU) tasks. Following
the methodology of Wang et al. (2022), we randomly desig-
nated 4 datasets as held-out test sets and used the remaining 13
datasets for training. Each training dataset contains 150 exam-
ples, and each test dataset contains 100 examples, resulting in a
training set of 1,950 examples and a test set of 400 examples, as summarized in Table 3 (details are
provided in the Appendix C.1. The details of datasets used in instruction tuning is presented in Table
4.

Category # Tasks # Samples

NLG
NLU

NLG
NLU

1050
900

Training

200
200

Testing

Split

2
2

7
6

C.2 BASELINE AND IMPLEMENTATION DETAILS

To evaluate the performance of our trained model, we 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
et 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 of demonstrations during
training. It can be considered a form of epoch-level data augmentation.

ERM with Instance Mixup (ERM+IM)(Zhang et 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 to 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).

17

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

Under review as a conference paper at ICLR 2025

Category Hyperparameter

LLMs

LoRA

P-Net

Learning rate
Batch size
Max sequence length
Weight decay coefficient
Epoch

Rank
Alpha
Dropout
P-Net target modules

LLMs target modules

Temperature
Iteration coefficient
Entropy constraint
Noise
Learning rate
Batch size
Max sequence length

Value

3e-5
16
512
0.1
2

8
32
0.1
q, v
q_proj, k_proj, v_proj,
o_proj, gate_proj,
up_proj, down_proj

0.1
80
1.0
0.3
1e-4
16
512

Table 5: Hyperparameter settings used in our main experiment.

InfoAC: (Xiang et al., 2024) is a training-stage method that employs contrastive learning to enable
earlier tokens to access information from later tokens, amining to mitigate the order sensitivity of
ICL inherent in autoregressive LM.

Batch-ICL: (Zhang et al., 2024) is an inference-stage method that transforms an n-shot ICL prompt
into n individual one-shot prompts and then ensembles the results to improve robustness.

CurStable: (Chang & Jia, 2023) is another inference-stage method that enhances ICL performance
by selecting optimal demonstration samples. This selection process involves performing multiple
inferences with different prompts on a validation set, calculating the expected performance when
each demonstration is used, and assigning an importance score to each. The demonstrations with the
highest scores are then selected to form the demonstration pool.

By including these baselines—training-stage (ERM, ERM+DS, ERM+IM, and InfoAC) and inference-
stage (Batch-ICL and CurStable)—we provide a comprehensive evaluation of our proposed method.

As for the proposed PEARL framework, we select the LLaMA3-8B model as our LLM and the
FLAN-large encoder as the P-Net. Both models are fine-tuned using LoRA (Hu et al., 2022), with
the number of finetuned parameters of P-Net being 1/20 that of the LLM. We train the models on
the instruction dataset for two epochs using a single NVIDIA A40 GPU, with a batch size of 16,
resulting in a total of 246 training steps. The optimizer used was AdamW. The learning rates for the
P-Net and the LLM are set to 1 × 10−4 and 3 × 10−4, respectively. For the Sinkhorn algorithm, we
use 80 iterations, a temperature parameter of 0.1, and an entropy constraint coefficient β = 1.0.

C.3 DETAILS OF HYPERPARAMETER SETTINGS

In this section, we provide a comprehensive overview of the hyperparameter settings used in our
experiments (Table 5). The hyperparameters can be categorized into three groups: (1) basic LLM
training parameters, such as learning rate and batch size; (2) LoRA configuration parameters; and (3)
P-Net optimization parameters. These hyperparameters were selected based on average validation
performance and kept consistent across comparative experiments to ensure fair comparison.

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

Figure 6: Impact of number of iterations and temper-
ature on the average/worst-case performance.

# Iter.

Temperature

0.03

0.1

0.3

80
200

55.7 / 40.0
55.7 / 40.0

55.7 / 40.0
55.8 / 40.0

55.4 / 39.6
55.8 / 40.6

Figure 7: Impact of entropy coefficient.

D ANALYSIS OF HYPERPARAMETERS IN INSTRUCTION FINETUNING

We conduct analysis to 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 in OT, testing values of 0.3, 1.0, 3.0, and 10.0 (Figure 7). 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 P-Net peaked at 1.0, suggesting optimal
learning conditions. As coefficients increased to 3.0 and 10.0, P-Net’s gradient norm decreased
again, suggesting excessive restrictions. The range of 1.0-3.0 provided an ideal balance, encouraging
P-Net to extract meaningful information from the LLM without oversimplifying or overcomplicating
the task. In contrast, the LLM’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 are 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 at 1, we vary the number of iterations (80, 200) and temperature
(0.03, 0.1, 0.3). As presented in Table 6, surprisingly, these substantial parameter changes result in
minimal performance variation. This suggests that the Sinkhorn algorithm in our framework is less
sensitive to these parameters than initially hypothesized, potentially indicating a wider range of stable
configurations for practical applications.

E EXTENDED INSTRUCTION FINETUNING ACROSS DIVERSE LLMS

We expanded our evaluation to include additional models: Mistral-7B, Gemma-7B, and earlier
generations such as Llama2-7B and Llama2-13B, as detailed in the tables from Table (6) to Table (8).

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 percentage points.

Adaptation of the Proposed Method In scenarios with three or more examples, our method
consistently demonstrated substantial improvements, often enhancing worst-case performance by
more than 10%. These results confirm the robustness and effectiveness of our approach.

F SCALING TO MANY-SHOT IN-CONTEXT LEARNING

19

972
973

974
975
976
977
978
979

980
981
982
983
984
985

986
987
988
989
990

991
992
993
994
995
996

997
998
999
1000
1001
1002

1003
1004
1005
1006
1007
1008

1009
1010
1011
1012
1013

1014
1015
1016
1017
1018
1019

1020
1021
1022
1023
1024
1025

0.31.03.010.001234567Gradient NormGradient Norm vs P-NetLLMUnder review as a conference paper at ICLR 2025

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

Table 6: 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.

2

3

4

5

ERM
PEARL 67.0 (+4.5)

64.1

58.1
62.4 (+7.5)

ERM
PEARL 69.5 (+4.3)

66.6

56.1
62.8 (+12.0)

ERM
PEARL 68.3 (+2.5)

66.7

50.4
57.1 (+13.4)

ERM
PEARL 70.2 (+3.4)

67.9

50.7
58.1 (+14.5)

67.0
68.0

67.0
70.0

68.9
69.9

67.5
70.4

64.0
66.0

62.0
66.0

60.0
62.0

56.0
64.0

54.6
59.4

63.7
70.1

67.6
71.6

70.7
76.7

41.8
49.0

38.9
60.1

47.8
54.8

52.6
59.3

81.0
82.0

80.0
83.6

74.2
74.9

76.0
73.3

78.0
78.0

76.0
78.0

52.0
66.0

56.0
66.0

53.7
58.4

55.6
54.1

55.9
56.8

57.4
60.4

48.5
56.7

47.3
47.0

41.6
45.5

38.2
43.0

Table 7: Instruction fine-tuning results for Gemma-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.

2

3

4

5

ERM
PEARL

ERM
PEARL

ERM
PEARL

ERM
PEARL

66.2 (+0.0)
66.3

59.5 (+2.0)
60.7

64.7 (+5.8)
68.4

52.5 (+13.0)
59.3

65.0 (+3.4)
67.2

46.5 (+13.0)
52.5

64.3 (+3.1)
66.3

46.3 (+10.2)
51.0

71.0
74.0

70.7
74.7

65.0
71.4

65.9
70.3

70.0
68.0

64.0
68.0

54.0
60.0

54.0
60.0

59.1
47.3

67.1
59.2

71.4
60.7

73.4
63.4

46.1
39.2

45.2
42.5

41.1
38.9

48.3
43.6

77.0
82.0

70.3
78.7

72.5
75.9

65.6
71.3

70.0
78.0

60.0
76.0

58.0
66.0

50.0
60.0

57.8
61.7

50.5
61.0

51.1
60.8

52.3
60.2

52.0
57.6

40.7
50.6

32.9
45.2

32.9
40.4

We evaluate the scalability of PEARL by extending our analysis to many-shot scenarios, testing
performance with 8 to 64 in-context examples (Table 10). 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
across four held-out tasks. Our analysis reveals persistent performance advantages of PEARL over
the ERM baseline across all shot regimes.

G 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 a balanced perspective. The results are shown in the Table 11.

Surprisingly, the results show that across all datasets and in every shot condition, PEARL’s best
performance consistently exceeded that of ERM. This indicates that our method not only optimizes
performance in worst-case scenarios but also slightly enhances best-case performance.

H SHOT EFFICIENCY

In the analysis of shot efficiency, we observe divergent trends between worst-case and average
performance metrics as the number of shots increases. Specifically, while worst-case performance
may decrease, average performance demonstrates improvement. This analysis is crucial for evaluating
the practical efficacy of training approaches in more realistic, variable conditions.

Our comparative analysis involves models trained with and without PEARL method. The results,
as summarized in Table 3, indicate that the PEARL-trained models generally achieve comparable
average performance to non-PEARL models using approximately two to four times as many shots.

20

Under review as a conference paper at ICLR 2025

Table 8: Instruction fine-tuning results for Llama2-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.

2

3

4

5

ERM
PEARL

ERM
PEARL

ERM
PEARL

ERM
PEARL

56.6 (+1.5)
57.4

46.3 (+0.4)
46.5

58.2 (+2.3)
59.6

34.0 (+19.1)
40.4

58.9 (+2.7)
60.5

19.9 (+59.1)
31.6

61.9 (+1.6)
62.9

25.8 (+24.7)
32.1

56.0
58.0

52.7
56.3

60.0
61.2

59.0
62.4

50.0
48.0

34.0
40.0

26.0
40.0

32.0
38.0

61.3
55.2

64.0
66.2

68.1
69.4

74.2
73.3

50.2
44.7

36.4
46.2

24.4
40.1

43.9
43.4

58.2
62.0

66.0
67.0

60.2
62.4

65.7
64.8

42.0
48.0

36.0
42.0

14.0
24.0

10.0
24.0

50.7
54.4

50.1
48.7

47.3
48.9

48.6
51.0

43.1
45.4

29.4
33.5

15.1
22.4

17.1
23.0

Table 9: Instruction fine-tuning results for Llama2-13B 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.

2.0

3.0

4.0

ERM
PEARL

ERM
PEARL

ERM
PEARL

66.3 (+2.4)
67.9

65.7 (+4.2)
68.5

56.6 (+7.3)
60.7

46.2 (+8.7)
50.3

65.8 (+0.9)
66.4

33.2 (+21.1)
40.2

56.0
64.0

55.7
62.7

58.2
63.3

46.0
58.0

38.0
44.0

28.0
42.0

72.6
73.8

76.4
81.0

79.6
80.4

56.2
64.2

51.3
58.4

41.6
45.5

83.0
81.0

77.7
76.7

73.7
69.4

76.0
76.0

56.0
56.0

38.0
42.0

53.4
52.6

53.1
53.5

51.8
53.1

48.0
44.4

39.6
42.6

25.0
29.1

In some instances, the performance equivalence exceeds this range, suggesting substantial gains in
sample efficiency.

As shown in Table 12, the results demonstrate that a PEARL-trained model, on average, requires
between 50% and 75% fewer shots to achieve performance levels comparable to those of a non-
PEARL model. This reduction in the required number of shots translates into a significant decrease
in computational complexity, from O(N 2) to O((N/2)2) or O((N/4)2), enhancing computational
efficiency.

1080
1081

1082
1083
1084
1085
1086
1087

1088
1089
1090
1091
1092
1093

1094
1095
1096
1097
1098

1099
1100
1101
1102
1103
1104

1105
1106
1107
1108
1109
1110

1111
1112
1113
1114
1115
1116

1117
1118
1119
1120
1121

1122
1123
1124
1125
1126
1127

1128
1129
1130
1131
1132
1133

21

Under review as a conference paper at ICLR 2025

Table 10: 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.

8

16

32

64

ERM
PEARL

ERM
PEARL

ERM
PEARL

ERM
PEARL

61.8 (+7.6)
66.5

21.3 (+39.2)
29.7

66.9 (+5.3)
70.5

21.3 (+23.7)
26.3

67.4 (+3.8)
70.0

19.3 (+36.4)
26.4

68.1 (+3.5)
70.4

20.6 (+36.7)
28.2

61.4
67.7

67.3
70.9

67.5
70.0

68.1
69.5

36.0
44.0

36.0
46.0

32.0
44.0

38.0
46.0

68.3
77.1

76.5
83.9

77.8
82.6

76.9
82.9

22.7
28.7

31.4
37.5

30.7
40.3

27.7
38.9

62.7
65.0

67.2
70.1

68.2
70.6

72.2
74.2

16.0
32.0

8.0
12.0

6.0
12.0

8.7
19.6

54.8
56.2

56.5
56.9

56.1
56.6

55.0
55.1

10.6
14.0

9.7
9.8

8.6
9.1

8.0
8.1

Table 11: Best performance comparison between ERM and PEARL

#Shot Method Average

gain CSQA CurDial CoLA TMW

2

3

4

5

ERM
PEARL
ERM
PEARL
ERM
PEARL
ERM
PEARL

64.1
68.8
72.8
77.0
82.9
84.3
86.8
89.3

7.2%

5.7%

1.7%

2.9%

68.8
73.4
70.3
73.4
81.3
82.8
84.4
87.5

64.4
69.2
85.0
87.9
92.4
93.6
95.3
96.5

64.1
70.3
65.6
79.7
78.1
81.2
81.3
85.9

59.2
62.1
70.3
66.9
79.7
79.5
86.2
87.3

Table 12: Average performance with and without PEARL

# Shots
wo PEARL
w PEARL

2
57.3
62.9

4
59.7
63.1

8
61.8
66.5

16
66.9
70.5

32
67.4
70.0

64
68.1
70.4

22

1134
1135

1136
1137
1138
1139
1140
1141

1142
1143
1144
1145
1146
1147

1148
1149
1150
1151
1152

1153
1154
1155
1156
1157
1158

1159
1160
1161
1162
1163
1164

1165
1166
1167
1168
1169
1170

1171
1172
1173
1174
1175

1176
1177
1178
1179
1180
1181

1182
1183
1184
1185
1186
1187

