PEARL: Towards Permutation-Resilient LLMs

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

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 3

Reviews and discussion:
<thread>
creation date:
1730132057630

thread_type:
official_review

<thread_object>
creator: reviewer_n5jf

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

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.
* Aside from a few minor typos, the paper is well-written and easy to follow.

weaknesses:
**Threat model**

* "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.
* "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.

**Experiments.**
* "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?
* 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.
* 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.
* 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)$.  
* 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.
* The paper would benefit from a clearer explanation regarding how the hyperparameters are chosen.
* It would be worth describing how the Rogue-L metric is calculated and why it's a good metric to use here.

**Proposed method.**
* 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)?
* The [CLS] token seems to appear from nowhere. What does CLS stand for?
* 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.
* 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.
* 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. 
* 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.
* Unless I missed something, the variable $G$ is undefined in Algorithm 1. Is that the P-Net?
* "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?
* I didn't understand the steps between 9 and 10. Consider that if we plug in the losses to

$$\min_{\theta,\phi} L(\theta){\text{LM}} + L(\phi)_{P-Net}$$

then objective becomes:

$$ \mathbb{E}[\ell(\theta; \phi; (\Pi \cdot p, x, y)) - \ell(\theta, \phi; (\pi\cdot p, x, y)) + \beta L(\theta, \phi)_\text{ent}] $$

which is equal to 

$$ \mathbb{E}[ \beta L(\theta, \phi)_\text{ent}] $$

since the first two terms cancel.  Could the authors clarify this, because it seems like the loss reduces to an entropy penalty.

**Discussion of related work.**
* 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.
* 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.

**Miscellaneous.**
* "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?
* Figure 1 -- the colors are too difficult to tell apart. It would be worth considering colors outside of the blue family.

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. 
* Another missing (easy) baseline is the impact of picking a *random* permutation in Figure 1. Or, is this how the average performance was calculated?
* 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730222632033

thread_type:
official_review

<thread_object>
creator: reviewer_cdtr

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

strengths:
1. The paper offers a simple, straightforward, intuitive and, most importantly, well-performing solution to the problem of permutation-sensitivity.
2. The paper is clearly written and well-presented.
3. 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).
2. 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.

questions:
1. The separator line between 2-shot and 3-shot in Table 3 is offsetted.
2. 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.
3. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730584676150

thread_type:
official_review

<thread_object>
creator: reviewer_zwpa

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

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

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

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

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.

1. 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).

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

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

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

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

1. Minor points about formatting:

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

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

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

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?

1. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681181773

thread_type:
official_review

<thread_object>
creator: reviewer_sluu

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

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. 
+ The attacker model for finding worst-case perturbation looks to have a favorable computational complexity. 
+ 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. 

+ 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). 
+ 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. 
+ 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. 

Chang, Ting-Yun, and Robin Jia. "Data curation alone can stabilize in-context learning." arXiv preprint arXiv:2212.10378 (2022).
Chen, Yongqiang, et al. "Positional Information Matters for Invariant In-Context Learning: A Case Study of Simple Function Classes." arXiv preprint arXiv:2311.18194 (2023).
Zhang, Kaiyi, et al. "Batch-icl: Effective, efficient, and order-agnostic in-context learning." arXiv preprint arXiv:2401.06469 (2024).
Xiang, Yanzheng, et al. "Addressing Order Sensitivity of In-Context Demonstration Examples in Causal Language Models." arXiv preprint arXiv:2402.15637 (2024).

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. 
+ 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. 
+ 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732462773161

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear AC and Reviewers,

We sincerely appreciate your time and expertise in reviewing our paper. We are pleased that you found our work to be unique and elegant (ZWPA, cDtR, SLuu), particularly in our design of the P-Net (threat model) which employs Optimal Transport (OT) and the Sinkhorn algorithm to learn challenging permutations for LLMs. The approach is recognized as both technically elegant and practical for practitioners in their fine-tuning pipelines (ZWPA), with favorable computational complexity and no additional inference-time overhead (ZWPA, SLuu). We are glad that our Distributionally Robust Optimization (DRO)-based training algorithm, working in concert with P-Net, effectively improves LLMs' robustness to ICL permutations (ZWPA, cDtR, n5jf), as demonstrated through our comprehensive experimental evaluation (n5jf). Additionally, we are pleased that the clarity and presentation of our manuscript were well-received (ZWPA, cDtR, n5jf).


We believe that all of the reviewers' concerns can be addressed. Below, we provide brief responses to the main concerns and suggestions raised in the reviews:

**Comparison With Baselines** (SLuu, cDtR, n5jf): We have extended our discussions and comparisons to include four additional baselines: three inference-time baselines and one training-phase baseline. Our method demonstrates substantial improvements over these baselines, achieving 9–21% performance gains. Furthermore, when combined with inference-time methods, our approach yields additional performance improvements of 3–5%, highlighting the complementary nature of our method. We have also updated the related work section to better situate our contributions within the existing literature (*see Section 2 and Table 2 in Section 6*).

**Experiments on Diverse LLMs** (ZWPA, n5jf): To validate the general applicability of our method, we have expanded our experiments to include four additional LLMs: Mistral 7B, Gemma 7B, Llama 2 7B, and Llama 2 13B. The results, provided in the rebuttal and updated manuscript, demonstrate consistent and significant performance gains ranging from 7% to 59% when the number of shots is greater than 2. These improvements across diverse model families indicate the broad applicability of our approach across different LLMs (*see Tables 6–9 in Appendix G*).

**Scaling to Many-Shot Settings** (ZWPA, cDtR): We have extended our evaluations to 8-, 16-, 32-, and 64-shot settings. The results show that our method generalizes well to large-shot and long-sequence scenarios, achieving an impressive 24–40% performance gain, even when trained with smaller settings. This demonstrates the excellent generalization capability of our approach (*see Figure 5 in Section 6 and Table 10 in Appendix F*).

**Clarification on Threat Model and Training Process** (n5jf): We acknowledge that our initial wording may have led to misunderstandings regarding the attack setting and training process of P-Net and the LLM. To address this, we have carefully revised the formulas and explanatory text in Sections 3 and 4 of the manuscript to enhance clarity and make it easier to understand. Additionally, we have provided explanations of these points in our responses (*see Sections 3 and 4, revisions highlighted in blue*).

All these revisions have been made in the manuscript and are highlighted for your convenience.

Please let us know if there are any outstanding concerns, and we are happy to discuss them further. We would appreciate it if you could take our responses into consideration when making the final evaluation of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission proposes a method for improving the robustness of in context learning to permutations in the demonstrations.

+ The paper is well-written.
+ The method is technically sound.

- There were some concerns raised regarding limited empirical evaluation.