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

We thank the reviewer for their detailed feedback. We have revised our manuscript significantly to address you and the other reviewers’ concerns; we have summarized these changes in a global comment to all reviewers. Below we address your specific comments:

Novelty of Paper and Technical Significance of Results: The primary novelty of the work is in studying and understanding how in-context learning in transformers is non-robust. Specifically:

• For linear transformers we show they are vulnerable because they implement gradient descent - demonstrating how seemingly desirable algorithmic properties can lead to exploitable weaknesses.
• In GPT2 style transformers, we have shown (with addition of new results in Sec. 5.6 regarding transferability of adversarial attacks from ordinary least squares (OLS) to transformers and vice versa) that these transformers are non-robust in a way that is distinct from both gradient descent and ordinary least squares. This is a novel result as prior works had strongly suggested that these transformers are implementing gradient descent or OLS.
• Our discovery that attacks transfer readily between smaller transformers but poorly between larger ones provides the first evidence of non-universal learning mechanisms in architecturally identical transformers.
• The success of adversarial training is surprising given the historical difficulty of robust regression with adaptive adversaries - a problem that saw meaningful progress only in the last decade (citation)
• The generalization of adversarial training from K to K' > K examples suggests transformers can learn remarkably robust algorithms through simple training procedures, and also illustrates transformer’s ability to perform non-convex optimization in-context.

Please also see our global comment where we discuss the relevance of this work to various sub-communities within machine learning.

Reason for studying regression setting:

• Transformers trained to do in-context learning are increasingly being used for tabular data settings including regression (Requeima et al., 2024, Hollmann et al. 2022, Ashman et al. 2024). This is part of a general trend where in-context learning capabilities are being utilized to develop performant models across various domains, e.g., vision, reinforcement learning and robotics (Zhang et al. 2023,Elawady et al. 2024, Raparthy et al. 2023).
• Additionally, studying a more abstract setting like regression enables developing a more robust understanding of a transformer's behavior, e.g., it allowed us to theoretically show and diagnose the vulnerability of linear transformers to hijacking attacks. At the same time, we would like to stress that robust regression is a highly challenging problem with computationally efficient methods being developed for this problem only in the last decade. Thus, transformers' ability to perform robust regression is both surprising as well as likely to be of interest to the robust statistics community.
• Studying in-context learning capabilities of transformers using linear regression is also not unprecedented and is in line with the number of prior works (Akyürek et al. 2022, Garg et al. 2022, Oswald et al. 2023, Zhang et al. 2024, Ahn et al. 2023).

New defenses: Given that standard adversarial training already works well in our setting without significant performance degradation, developing novel defense methods does not seem necessary. We would again stress that given the historical difficulty of robust regression in the presence of an adaptive adversary, the success of adversarial training in this setup and especially the efficiency of adversarial training is a surprising finding.

Please let us know if you have any further questions. We would be happy to address any remaining concerns or expand on any of these points.

Respected reviewer, we have given a detailed response to your comments below (and in the global response). As there are only few days remaining in the discussion period, we would greatly appreciate if you could review our response and let us know if you have any further questions. If we have successfully addressed your concerns, we request that you please revise your score accordingly.

We thank the reviewer for their detailed feedback. We have revised our manuscript significantly to address you and the other reviewers’ concerns; we have summarized these changes in a global comment to all reviewers. Below we address your specific comments:

Relevance to the community: Please see our global response which answers this question in detail. As suggested by you, we have also revised our introduction and discussion sections to make the contributions of this work, and their implications, more salient.

Meaning of attacks in linear regression context These attacks can be thought of as robustness to adversarial outliers, either in the feature/data space ($x$), the label space ($y$), or both ($z=(x,y)$). Equivalently, these can be viewed as robustness to data-poisoning attacks.

Comparison with OLS Our new results in Section 5.6 show that alignment between OLS and transformers weakens significantly out-of-distribution, with transformers exhibiting additional adversarial vulnerabilities relative to OLS.

Meaning of adversarial training from linear regression perspective This is a great question. For GPT2-trained transformers, our experiments strongly suggest that they don’t implement any standard regression algorithm - for instance, ridge regression produces a solution which is linear in the labels, which would thus be susceptible to “y-attacks” (renamed label attacks) in the same way that the linear transformer is, but our Figure 3 demonstrates this is not the case. Indeed, there are very few known algorithms for linear regression which are robust to hijacking attacks, and such algorithms often require complex routines to mitigate the effects of (unknown) adversarial examples (Cherapanamjeri et al., 2020). We think it would be interesting to develop a mechanistic understanding of what algorithm the adversarially-trained transformers implement.

“I would have tried to inspect better why attacks don't work on GPT2 style model…” Prior work pointed to the possibility that GPT2 models implemented either gradient descent (Oswald et al. 2023, Zhang et al. 2024) or ordinary least squares for regression problems (Akyürek et al. 2022, Garg et al. 2022). Our results demonstrate this is not the case, since attacks derived from gradient descent (and our Theorem 4.1 shows linear transformers implement gradient descent) don’t transfer to GPT2, and attacks from GPT2 don’t transfer to OLS (our new Section 5.6). This shows which algorithms GPT2 does not implement, but we aren’t confident about which algorithm it does implement, and we think this is an interesting direction for future research.

Please let us know if you have any further questions. We would be happy to address any remaining concerns or expand on any of these points.

Respected reviewer, we have given a detailed response to your comments below (and in the global response). As there are only few days remaining in the discussion period, we would greatly appreciate if you could review our response and let us know if you have any further questions. If we have successfully addressed your concerns, we request that you please revise your score accordingly.

Testing on larger models like LLaMA While computational constraints prevent us from training 8B+ parameter models, our analysis of scaling effects (Section 5.3 and Appendix B.1) suggests that model scale alone does not improve robustness to these attacks.

TRADES vs standard adversarial training Given that standard adversarial training already works well without significantly impacting clean performance (Figure 5), we found more sophisticated approaches unnecessary for this setting.

Clarity on x/y/z terminology and experiment details: We are sorry for the lack of clarity regarding terminology and the experiment section. We have updated our paper to change the x/y/z attack labels to feature, label, and joint attack. We have also updated our figures to be more explicit about which architecture and training procedure we use.

Context-specific attacks: We are unsure what the reviewer means here and would appreciate further clarification. In our setting, every attack is context-specific, since the adversary looks at the sequence of tokens and based on that, then adds or perturbs one or more tokens.

We would be happy to clarify any remaining concerns. We hope these revisions and explanations address your questions about the motivation and clarity of our work.

We thank the reviewer for their detailed feedback. We have revised our manuscript significantly to address you and the other reviewers’ concerns; we have summarized these changes in the global comment to all reviewers. Below we address your specific comments:

Motivation for and relevance of hijacking attacks: In LLMs, hijacking attacks can be considered to be an instance of prompt-injection / prompt-manipulation attacks which are now widely recognized as serious security vulnerabilities (e.g., Anwar et al. 2024, Section 3.5). For sequence-based models like transformers/LLMs, prompt-manipulation is a natural attack vector through which an adversary can influence a model’s prediction by adding adversarial tokens into the context. In this setting, a small change of the input corresponds to a small number of tokens added/changed in a long prompt. Transformers are increasingly being used in critical settings which rely upon in-context learning to perform their task, e.g., clinical decisions where the patient history and data is given as context (Nori et al. 2023), robotic control algorithms which rely on ICL (Elawady et al. 2024), automated code completion (Patel et al. 2024), etc., and in many settings the end-user/adversary doesn’t have the ability to manipulate the internals of the transformer but does have the ability to append tokens to the input of e.g,. an API. As such, we believe understanding the robustness of in-context learning under hijacking attacks is a particularly relevant problem, since they could change the outcome of a clinical decision, manipulate the movement of a robot’s arms, insert code backdoors, and so on.

We would also like to emphasize that hijacking attacks on LLMs have already been demonstrated in prior work, e.g.,

• Qiang et al, 2023 give a practical attack on LLMs that closely mirrors our x-attack (now renamed featuredata attack) in construction and show that their attack results in LLM outputting the target token.
• Bailey et al. 2023 give a hijacking attack on vision-language models (VLM) where they manipulate images present in the context to control the behavior of VLM.

Note that applications of in-context learning in transformers (and hence dangers of hijacking attacks) are not just restricted to LLMs. They are seeing rapid adoption in many domains, e.g., vision (Zhang et al. 2023, Kirsch et al. 2022) and robotics/reinforcement learning (Elawady et al. 2024, Raparthy et al. 2023) with both continuous and discrete input representations.

Our global comment further goes into the details of the novelty and relevance of our contributions.

Why adversarial examples fail to transfer: As we mention in our global comment, we provide new experiments in Section 5.6 which further probe transferability of attacks between transformers and OLS. Together with prior results shown in Figure 1 and Figure 6, we have a fairly complex picture of transferability:

• Within GPT2-style transformers, attacks transfer well across small and medium transformers but transfer poorly from and to larger transformers – even between standard transformers with identical architectures.
• Attacks derived from the linear transformer (i.e., GD solution) do not transfer well to any GPT-2 style transformers
• attacks derived from OLS transfer to medium-scale transformers relatively well but not larger transformers (Fig. 7)
• attacks derived from all transformers transfer poorly to OLS. This suggests that GPT2's out-of-distribution behavior differs fundamentally from both gradient descent and OLS, with some shared vulnerabilities to attacks from OLS but vulnerable to a broader set of attacks than OLS is. This challenges prior theories about what algorithms these models implement.

Respected reviewer, we have given a detailed response to your comments below (and in the global response). As there are only few days remaining in the discussion period, we would greatly appreciate if you could review our response and let us know if you have any further questions. If we have successfully addressed your concerns, we request that you please revise your score accordingly.

I have taken account of the global comment, and I keep my current rating. If you are very into argument, I suggest you submitting to a more theoretical-driven conference.