Adversarial Robustness of In-Context Learning in Transformers for Linear Regression

Authors: Usman Anwar, Johannes von Oswald, Louis Kirsch, David Krueger, Spencer Frei
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: in-context learning, transformers, hijacking attacks, linear regression, linear transformers, transfer of adversarial attacks

Abstract:
Transformers have demonstrated remarkable in-context learning capabilities across various domains, including statistical learning tasks. While previous work has shown that transformers can implement common learning algorithms, the adversarial robustness of these learned algorithms remains unexplored. This work investigates the vulnerability of in-context learning in transformers to _hijacking attacks_ focusing on the setting of linear regression tasks. Hijacking attacks are prompt-manipulation attacks in which the adversary's goal is to manipulate the prompt to force the transformer to generate a specific output. We first prove that single-layer linear transformers, known to implement gradient descent in-context, are non-robust and can be manipulated to output arbitrary predictions by perturbing
a single example in the in-context training set. While our experiments show these attacks succeed on linear transformers, we find they do not transfer to more complex transformers with GPT-2 architectures. Nonetheless, we show that these transformers can be hijacked using gradient-based adversarial attacks. We then demonstrate that adversarial training enhances transformers' robustness against hijacking attacks, even when just applied during finetuning.  Additionally, we find that in some settings, adversarial training against a weaker attack model can lead to robustness to a stronger attack model.  Lastly, we investigate the transferability of hijacking attacks across transformers of varying scales and initialization seeds, as well as between transformers and ordinary least squares (OLS). We find that while attacks transfer effectively between small-scale transformers, they show poor transferability in other scenarios (small-to-large scale, large-to-large scale, and between transformers and OLS).

Decision: reject
Reviewer scores: 3, 6, 5

Reviews and discussion:
<thread>
creation date:
1730193429305

thread_type:
official_review

<thread_object>
creator: reviewer_kcot

<content>
summary:
This paper focuses on studying the adversarial robustness of in-context learning for transformers. The paper first proves that a single-layer transformer can be manipulated to output arbitrary predictions, and also uses gradient-based attacks to demonstrate the vulnerability of GPT-2.

strengths:
1. The paper has a solid theoretical analysis to show the vulnerability of single-layer transformers.

2. Beyond single-layer transformers, the paper also considers GPT-2, and the experimental results appear to support the conclusions well.

3. The paper is easy to follow and well-organized.

weaknesses:
1. The novelty of this paper is limited because the phenomenon that single-layer transformers and GPT-2 are vulnerable to adversarial attacks is not surprising, and the paper does not have new and strong technical contribution compared with existing gradient-based attacks.

2. The task that uses GPT-2 for linear regression is not well-motivated. We usually do not use GPT-2 to solve linear regression.

3. The paper does not propose any new defense method beyond adversarial training, which is a standard baseline.

questions:
No question, and please refer to the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730471831680

thread_type:
official_review

<thread_object>
creator: reviewer_g6bb

<content>
summary:
The paper shows that linear transformers trained on linear regression aren't robust to hijacking attacks. Moreover, the authors show that these attacks don't transfer from small models to more complex GPT2-style ones (and even among them). The latter can be hijacked with gradient-based optimization. Finally, adversarial training (or fine-tuning) is shown to be promising to prevent these attacks.

strengths:
- Methods and experiments are clearly explained and carried out. The work is original, but the scope could have been better stated.

- Results are well presented and commented on. However, they're not very significant.

- I personally like the part about adversarial training, as it is a promising method to make transformers more robust to attacks.

weaknesses:
- I don't get the point of this work and why it has to be considered relevant for the community. I would say this better, possibly on the first page.

- I would have tried to inspect better why attacks don't work on GPT2 style models. This would have been helpful in understanding better how these models perform linear regression.

questions:
1. What's the meaning of these attacks in the context of linear regression? Is there a more practical interpretation?

2. How do these results compare to when using OLS? Can the latter be used as a baseline?

3. What does adversarial training on attacks mean from a linear regression perspective? (eg. can it be related to some form of regularization?)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730878616492

thread_type:
official_review

<thread_object>
creator: reviewer_8xug

<content>
summary:
This paper explores the adversarial robustness of in-context learning in transformers, focusing on hijacking attacks in linear regression tasks. Single-layer linear transformers are shown to be easily manipulated, while complex models like GPT-2 resist simple attacks but remain vulnerable to gradient-based ones. Adversarial training enhances model robustness, especially during fine-tuning, and hijacking attacks transfer only among smaller models.

strengths:
1. The paper is overal clear and well-written with sufficient experiments.

weaknesses:
1. I find the use of x-attack, y-attack, z-attack to be confusing. The joint attack (z-attack) does not really corespond to a different dimension to attack. It is better to call it data, label, and joint attack, or something similar. 
2. From the motivation level, I am uncertain if hijacking attack is the main concern in the applications of LLM for transformers. These line of work along Garg et al. are used to demonstrate the in-context capability of transformers which is a simplified setting but sufficient. The typical adversarial learning considers the setting where the input is modified only slightly, but the adversary can have large impact on the output. I don't see how this classical analysis can easily transfer to the text domain. How to define the allowable small perturbation. If we naively apply on token-level, people can easily find it. If the text contains number, and we only make small modification on the number, will it actually lead to different output. So, I am not sure if the paper can provide many insights into understanding the hijack attack for LLM. I will raise my score if the author sufficiently addresses this issue.
3. Since the authors mainly investigate the in-context learning setup, they should consider some context-specific attacks, for example within context or across-context. I don't think uniform attack on all possible entries are the only interesting scenario.
4. I appreciate the empirical effort by the authors. In terms of novelty, I would like to see more novel algorithm design and potential implications, since incontext is still a different setting from the standard ML setting. 
5. I find the experiment section to be hard to follow. Even though the model choice is discussed in the setup section, it is unclear what architecture and training procedure that each figure corresponds to e.g. figure 2.3.4.

questions:
1. Can the author provide why the adversarial examples fail to transfer? It is possible that the attack examples cannot be transferred from linear attention model to GPT2, but it can transfer from GPT2-small to GPT2-large. 
2. Has the author experiments with larger model than GPT2, such as llama 8B? I am wondering if the hijack attacks can be mitigated by the emergent ability of transfoermers. 
3. Did the author experiements with TRADES, which usually serves as a better defense than the original PGD.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732036976832

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Authors' Global Response -- New Results and Clarification on Contributions

comment:
We are thankful for all of the reviewers’ constructive and useful feedback, which has helped us to strengthen our work. Based on your comments, we have made several substantial improvements to clarify our key contributions and their broader significance. We have colored the major changes in blue for ease of reviewers. Two key updates are: 
- New experimental results comparing transformer behavior to ordinary least squares (Section 5.6 and Appendix B.5).
- Revamped introduction to better motivate our paper, and expanded discussion of implications for different research communities.  

## New Experimental Results
We have added new results investigating the transfer of adversarial examples from ordinary least squares (OLS) to transformers and vice versa. These new results show that 
- Attacks derived from OLS transfer relatively well (but not perfectly) to medium-scale transformers but do not transfer at all to larger transformers.
- Attacks derived from all transformers transfer poorly to OLS.
These new results imply that transformer’s out of distribution behavior does not match with OLS, challenging prior results which suggest transformers may be implementing OLS. 

## Relevance and Motivation of This Work
As transformers and large language models increasingly serve as the foundation for real-world applications, understanding how they perform in-context learning becomes crucial for both reliability and security. Our study of hijacking attacks - where an adversary can manipulate model behavior through carefully crafted context examples - reveals fundamental vulnerabilities in this learning mechanism. These findings are particularly relevant as ICL-based inference becomes prevalent across domains including embodied AI ([DeepMind ‘23](https://arxiv.org/abs/2301.07608), [Elawady et al. 2024](https://arxiv.org/abs/2410.02751)), and structured data (tabular data) processing ([Requeima et al., 2024](https://arxiv.org/abs/2405.12856), [Hollmann et al. 2022](https://arxiv.org/abs/2207.01848), [Ashman et al. 2024](https://arxiv.org/abs/2406.13493)), and is used for sensitive applications like clinical decision making ([Nori et al. 2023](https://arxiv.org/abs/2311.16452)) and robot control [Elawady et al. 2024](https://arxiv.org/abs/2410.02751).

We believe our work reveals several fundamental insights about transformer behavior that is likely to be of interest to multiple research communities:

**For researchers studying in-context learning:** Our use of hijacking attacks to examine alignment between transformers and traditional solvers is novel and resulted in novel findings.
- Our transfer analysis reveals that transformers' out-of-distribution behavior is mechanistically distinct from both gradient descent (which is implemented by linear transformers) and OLS. This  challenges prior explanations of how these models perform in-context learning ([Akyürek et al. 2022](https://arxiv.org/abs/2211.15661), [Garg et al. 2022](https://arxiv.org/abs/2208.01066), [Oswald et al. 2023](https://proceedings.mlr.press/v202/von-oswald23a.html), [Zhang et al. 2024](https://arxiv.org/abs/2306.09927), [Ahn et al. 2023](https://arxiv.org/abs/2306.00297)).
- Our discovery that adversarial attacks transfer readily between smaller transformers but poorly between larger ones provides the first evidence of non-universal learning mechanisms in architecturally identical transformers. 

Combined, these results indicate that developing an understanding of in-context learning even in a highly structured setting like linear regression is an open problem and may be more challenging than previously thought to be.

**For the robust statistics community:** 
Achieving adversarial robustness against adaptive adversaries in regression is historically very difficult, with meaningful progress only being made in the last decade ([Diakonikolas and Kane, 2019](https://arxiv.org/abs/1911.05911)) or so. Our finding that transformers can learn robust algorithms through simple adversarial training that generalizes beyond the training setting is thus particularly noteworthy and likely of interest to the robust statistics community.

**For security and adversarial robustness community**: 
Hijacking attacks on LLMs through perturbation of in-context dataset have already been demonstrated in prior work, but not well-understood ([Qiang et al. 2023](https://arxiv.org/abs/2311.09948), [Bailey et al. 2023](https://arxiv.org/abs/2309.00236)). On an abstract task, we have studied how these vulnerabilities arise resulting in two notable findings: 
- These vulnerabilities arise due to the underlying algorithmic mechanisms of ICL
- Even adversarially trained networks are vulnerable to deeper attacks (that target more tokens).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732555379107

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion Period Ending Soon

comment:
Dear reviewers, there is only one day left in the discussion period. If you have any unaddressed concerns, we would be thankful if you could share them and give us a chance to respond. If not, we hope that you may consider updating your scores.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission "Adversarial Robustness of In-Context Learning in Transformers for Linear Regression" investigates the question of robustness of in-context-learning in transformer models against small-scale modifications of the input context. The main vehicle of their analysis is a proof that that 1-layer transformers, which approximate gradient descent with ICL (as based on previous work in connecting ICL and optimization strategies), are non-robust because of this.
The authors then run a series of transfer experiments, between linear and normal tranformer models, OLS models, models of different seeds and sizes, and models trained with adversarial training.

Ultimately, the submission is unable to convince the reviewers and me that this is meaningful. The proof is an interesting niche, but hard to draw conclusions to more practical systems, where other formulations, such as Wolf et al, "Fundamental Limitations of Alignment in Large Language Models" already prove stronger theoretical statements about the existence of prompt injection attacks. There are some applications of the results as analysis tools to understand ICL algorithms, but the submission does not find a substantial conclusion.

Further, the submission notes that these findings call into question previous results about the algorithms approximated by ICL, but previous work makes no statement whether the instantiation of the algorithm is adversarially robust. 


I think the authors could improve this work by consolidating their findings into a more consistent presentation. Based on this, I do not recommend acceptance.