Understanding In-Context Learning of Linear Models in Transformers Through an Adversarial Lens
Authors: Usman Anwar
Johannes von Oswald
Louis Kirsch
David Krueger
Spencer Frei
Venue: Accepted by TMLR ()

Abstract: In this work, we make two contributions towards understanding of in-context learning of linear models by transformers. First, we investigate the adversarial robustness of in-context learning in transformers to hijacking attacks — a type of adversarial attacks in which the adversary’s goal is to manipulate the prompt to force the transformer to generate a specific output. We show that both linear transformers and transformers with GPT-2 architectures are vulnerable to such hijacking attacks. However, adversarial robustness to such attacks can be significantly improved through adversarial training --- done either at the pretraining or finetuning stage --- and can generalize to stronger attack models. Our second main contribution is a comparative analysis of adversarial vulnerabilities across transformer models and other algorithms for learning linear models. This reveals two novel findings. First, adversarial attacks transfer poorly between larger transformer models trained from different seeds despite achieving similar in-distribution performance. This suggests that transformers of the same architecture trained according to the same recipe may implement different in-context learning algorithms for the same task. Second, we observe that attacks do not transfer well between classical learning algorithms for linear models (single-step gradient descent and ordinary least squares) and transformers. This suggests that there could be qualitative differences between the in-context learning algorithms that transformers implement and these traditional algorithms.

Published in Transactions on Machine Learning Research (08/2025)
Understanding In-Context Learning of Linear Models in Transformers Through an Adversarial Lens
Usman Anwar University of Cambridge
ua2370com.ac.
Johannes Oswald
joosuald@google..com
Google, Paradigms of Intelligence Team
Louis Kirsch
Ukirsch @ggoogle com
Google Deep
David Krueger
kruegerd@mälla quebec
Mila & Université de Montréal
Spencer Frei
sfrei@uedavis.cc edu
UC Dovis
Reviewed on OpenReview: https: // openreview. net/ forum? id= C1MXJO07SJ
Abstract
In this work, make two contributions towards understanding of in-context learning of linear models by transformers. First, we investigate the adversarial robustness of in-context learning in transformers to hijacking attacks a type of adversarial attacks in which the adversary's goal is to manipulate the prompt to force the transformer to generate specific output. We show that both linear transformers and transformers with GPT-2 architectures are vulnerable such hijacking attacks. However, adversarial robustness such attacks can be significantly improved through adversarial training done either at the pretraining finetuning stage and generalize to stronger attack models. Our second main contri- bution is a comparative analysis of adversarial vulnerabilitia across transformer models and other algorithms for learning lincar models. This revcals two novel findings. First, adversar- ial allacks transfer poorly between larger transformer models trained from different seeds despite achieving similar in-distribution performance. This suggests that transformers of the same architccture trained according the same recipe may implement different in-context learning algorithms for the same task. Second, we observe that attacks do not transfer well between classical learning algorithms for linear models (single-step gradient descent and ordinary least squares) and transformers. This suggests that there could be qualitative dif- ferences between the in-context learning algorithms that transformers implement and these traditional algorithms.
1 Introduction
Transformers exhibit sophisticated in-context learning capabilities across variety of settings such as lan-
guage [Brown et al. 2020). vision KKrsch ct al. 2022 Bar et al. 2022 Zhang et al. 2023), tabular data (Holl-
mann el al. 2022 Requeima et al. 2024 Ashman el al. 2024 reinforcemer learning and robolics (Chen
et al.) 2021 Raparthy et al. 2023 Team et al. 2023) Elawady et al. 2024). The mechanisms underlying this behavior, however, remain poorly understood. While recent works have made progress by studying transformer behavior on supervised learning tasks, fundamental questions about how these models learn and
implement algorithms in-context remain open (Anwar et al. 2024 Section 2.1).

---

Published in Transactions on Machine Learning Research (08/2025)
In this work, we investigate the mechanisms of in-context learning through the lens of adversarial robustness to hijacking attacks -a threat model where adversary manipulates examples in the in-context set to force the model to output specific target values (Qiang et al. 2023 Bailey et al. 2023). Beyond its direct practical relovance for deployed language models and emerging applications of in-context learning across various domains for sensitive applications like clinical decision-making (Noi et al. 2023 or robot control (Elawady ct al. 2024); studying hijacking attacks provides a powerful tool for probing and understanding the algorithms that transformers learn to implement in-context.
Our investigation focuses the setting of linear regression tasks, where we analyze two architecture classes: single-layer linear attention models and GPT-2 style transformers. Through combination of theoretical analysis and extensive experiments, we show that both the linear transformer, and GPT-2 style transformers, do not learn adversarially robust in-context learning algorithms by default. Specifically:
1. We prove that single-layer linear transformers, which prior work showed implement gradient descent on in-context data (von Oswald et al. 2022), can be hijacked through perturbation of just a single token (Theorem 4.1 We then empirically demonstrate that GPT-2 style transformers are also vulnerable to hijacking attacks, which we find through gradient-based methods.
2. Encouragingly, we show that adversarial training can effectively improve robustness of GPT-2 style trans- formers, with impressive generalization: training on perturbations of K examples yiclds robustness against manipulation of K' K tokens. This is particularly surprising given the historical difficulty of achieving robustness against adaptive adversaries in regression tasks (Diakoikkllas & Kane 2019).
We then provide comparative analysis of adversarial vulnerabilities of different transformers and traditional solvers. Specifically, analyze to what extent adversarial examples designed for particular model or solver transfer to other models or solvers. This comparative analysis yields multiple interesting insights:
1. We find that hijacking attacks transfer readily between smaller transformers but show poor transferability between larger transformers of identical architecture but trained from different random seeds, providing the first evidence that architecturally identical transformers trained with the same recipe may learn distinct in-context learning algorithms.
2. We find that successful attacks against linear transformers (which perform gradient descent) fail to transfer to GPT2 architectures. Similarly, we find that adversarial attacks transfer poorly between transformers and ordinary least squares (OLS). This shows that the ont-of-distribution behavior of transformers is mechanistically distinct from both gradient descent and OLS - calling into question prior explanations about what algorithms these models implement to learn in-context (Garg et al., 2022 Akyürek et al. po22p.
2 Related Works
In-Context Learning of Supervised Learning Tasks: Our work is most closely related to prior works
that have attempted to understand in-context learning of linear functions in transformers (Garg et al.] 2022
Akyürek et al., 2022 von Oswald et al., 2022 Zhang et al.) 2024 Fu et al.] 2023 Ahn et al. 2023 Vla-
dymyrov et al. 2024 von Oswald et al. 2022) provided a construction of weights of linear self-attention
layers (Schmidhuber 1992 Katharopoulos et al. 2020; Schlag al., (021)) that allow the transformer to
implement gradient descent over the in-context examples. They show that when optimized, the weights of the linear self-attention layer closely match their construction, indicating that linear transformers implicitly perform mesa-optimization. This finding is corroborated by the works of /Zhang et al. (2024) and Ahn et al. 2023). A number of works have argued that when GPT2 transformers are trained on linear regression, they learn to implement ordinary least squares (OLS) (Garg et al.) 2022 Akyürek et al. 2022 Fu et al. 2023 More recently, Vladymyrov et al. (2024) show that linear transformers also implement other iter- ative algorithms on noisy linear regression tasks with possibly different levels of noise. Bai al. (2024) show that transformers perform in-context algorithm selection: choosing different. learning algorithms
2

---

Published in Transactions on Machine Learning Research (08/2025)
to solve different in-context learning tasks. Other neural architectures such recurrent neural networks have also been shown to implement in-context learning algorithms (Hochreiter et al., 2001) such as bandit algorithms (Wang et al. 2016) or gradient descent (Kirsch & Schmidhuber 2021).
Hijacking Attacks: While a considerable amount of research has been conducted on the security aspects of LLMs, most of the prior research has focused on jailbreaking attacks. To the best of our knowledge, Qiang et al. (2023) is the only prior that considers hijacking attack on LLMs or transformers during in- context learning. They show that it is possible to hijack LLMs to generate unwanted target outputs during in-context learning by including adversarial tokens in the demos. He et al. (2024) also consider adversarial perturbations to in-context data, however, their goal is simply reduce the in-context learning performance of the model in general, and not in targeted way. Bailey et al. (2023) demonstrate that vision-language models can be hijacked through adversarial perturbations to the vision modality alone. Similar to our work, both Qiang et al. (2023) and Bailey et al. (2023) assume a white-box setup and use gradient-based methods for finding adversarial perturbations to hijack the models.
Robust Supervised Learning Algorithms: There are a number of frameworks for robustness in machine learning. The framework we focus on in this work is data contamination, potiiioonnoiooinnn....oo where an adversary can manipulate the data in order to force predictions. Surprisingly, designing efficient robust learning algorithms, even for the relatively simple setting of linear regression, has proved quite challenging, with significant progress only being made the last decade (Diakonikolas & Kane 2023). Different algorithms have been devised which work under & contamination model where only labels y can be corrupted (Bhatia et al.) 2015
2017: Suggala et al.] 2019) or when both features x and labels y can be corrupted (Klivans et al. 2018 Diakonikolas et al.)
2019 Cherapanamjeri et al., 2020). Note that all the aforementioned work focus hand-designing robust learning algorithms for each problem setting. In contrast, we are concerned with
understanding the propensity of the transformers to learn to implement robust learning algorithms.
There are a number of other related frameworks for robustness in machine learning, e.g., robustness with respect to imperceptible (adversarial) perturbations of the input (Goodfellow al., 2015 Madry et al. 2018). We do not focus these attack models in this work.
3 Preliminaries
In this work, we investigate whether the learning algorithms that transformers learn to implement in-context are adversarially robust. We focus on the setting of in-context learning of linear models, a setting studied significantly in recent. years (Garg et al., 2022
Akyürek et al. 2022 von Oswald et al. 2022 Zhang et al. 2024 Ahn et al. 2023). We assume pre- treinini data that are sampled as follows. Each linear regression task is indexed by 5 E [], with each task consisting of N labeled examples (I+,a Vr,a), query example gr,q
parameters Wy 1...d N(O, Ia), features 27,1; greqy Vr.query w. Fr.query-
l.d. N(O, Ia) (independent of wx), and labels V7, w Zr,i
The goal is to train a transformer on this data (by method to be described shortly) and examine if, after pre-training, when we sample a new linear regression task (by sampling a new, independent w ~ N(0, Ia) and features Ti, i == I,...,,), M), the transformer can formulate accurate predictions for new, independent query examples. Note that the number of examples M in a task at test time may differ from the number of examples N per task observed during training.
To feed data into the transformer, we need to decide on tokenization mechanism, which requires some care since transformers map sequences of vectors of fixed dimension into sequence of vectors of the same length and dimension, while the features x; are d-dimensional and outpuls yi are scalars. That is, from a prompt of N input-output pairs (I,, and a test example Tquery for which we want to make predictions, the question is how to embed
p= (#1, V1 ....... N, UN. Zquery):
into a matrix. We will consider two variants of tokenization: concatenation (denoted Concat), which con- catenates Iy and yi and stacks each sample into column of an embedding matrix, and then appends
3

---

Published in Transactions on Machine Learning Research (08/2025)
(aquery. 0)T E R the last column:
E(P)
T1 D2
NN
Jquery
0 E R (d(++)x((+1))
(Concat)
(1)
1/ 1/2
UN
The notation E(P) emphasizes that the embedding matrix is function of the prompt P, and shall sometimes denote this as E for case of notation. This tokenization has been used in number of prior works
on in-context learning of function classes (von Oswald et. al. 2022 Zhang et al. 2024 Wu et al. 2023).
+ 1, 1, M + 1) entry of the transformer output. This allows for last-token prediction formulation of the squared-loo objective function: if (E a transformer, the objective function for B batches of data consisting of N + samples (Ir, (r, (Fr,qquery query)))) each batch embedded into Er,, is
With the Concat tokenization the natural predicted value for 2M+1 appears in the (d
1 Σ. (U;())))aaa,, Vrrqqerr)) 2B
I(0)
(2)
We will also consider an alternative tokenization method, Interleave, where features a and are interleaved into separate tokens,
E(P) 0
12
JN
Equery
(((+)))(N+1))
(Interleave)
(3)
y1
0
YN
By using causal masking, i.c. forcing the prediction for the i-th column of E, to depend only on columns S i, this tokenization allows for the formulation of next-loken prediction averaged across all N pairs of examples,
[(00= = 2B
1 E ((MMask 0)) Vr,+11*
(4)
where we treat Vr,N+1 Vr.query- This formulation was used in the original work by Garg el al. (022
We consider in-context learning in two types of transformer models: single-layer linear transformers, where we can theorctically analyze the behavior of the transformer, and standard GPT-2 style transformers, where we use experiments to probe their behavior. In all experiments, we focus on the setting where d 20 and the number of examples per pre-training task is N 40.
3.1 Single-Layer Linear Transformer Setup
Linear transformers are a simplified transformer model in which the standard self-attention layers are replaced
by linear self-attention layers KKthhroppoulss et al.) 2020 von Oswald et al. 2022 Ahn et al. 2023 Zhang
et al.
2024 Vladynyrov el al. 2024). In this work, specifically consider single-laye layer linear sell-attention
(LSA) model,
(EsA(E;)) = fisa(E; WWKQ := = E+ E WKQE N
This is a modified version of attention where we remove the softmax nonlincarity, merge the projection and value matrices into single matrix WVV E Rd+1xd+1, and merge the query and key matrices into single matrix WKQ E [Rd+ix++1 For the linear transformer, we will assume the Concat tokenization.
Prior work by Zhang et al. 2024) developed an explicit formula for the predictions fLSA when is pre- trained on noiseless lincar regression tasks (under the Concat tokenization) by gradient flow with particular initialization scheme. This corresponds to gradient descent with infinitesimal learning rate d V (T)) in the infinite task limit B of the objective
L(0) lim Ż(0)= B-+00
[(I(E,;)))aIIN)) x J w) w) 2
(6)
Lid.
w. N(0.1) Z Br N(O.I)
3.2 Standard Transformer Setup
For studying the adversarial robustness of the in-context learning standard transformers, we use the same setup as described in Garg et al. (2022). Namely, we use a standard GPT2 architecture with the Interleave tokenization. We provide details the architecture and the training setup in Appendix C
4

---

Published in Transactions on Machine Learning Research (08/2025)
3.3 Hijacking Attacks
We focus on a particular adversarial attack where the adversary's goal to hijack the transformer. Specif- ically, the aim of the adversary is to force the transformer to predict specific output Ybad for zquery when given prompt P (21, 9 yM, F.uery). The adversary can choose one or more pairs (Iz, yi) to replace with an adversarial example (radv: Vadv).
We characterize hijacking attacks in this work along two axes: (i) the type of data being attacked (it) number of data-points or tokens being attacked. The adversary may perturb either the æ feature (r,, y4) H> (Tadv, We). which we call atttuk,, or label y, (I,, ys) + (Ig, Vadv), which we refer to as label-attack, or simulltnnously perturb the pair (xi, yi) + (Nadu; Vadu), which we refer to as joint-attack. We will primarily focus feature-attack and label-attack to feature-attack (sce Figurcs and 2). Furthermore, we allow for the adversary to perturb multiple tokens in the prompt P. A k-token attack means that the adversary can perturb at most k pairs (1,, Bt) in the prompt
the behavior of joint-attace is qualitatively quite similar
We note that hijacking attacks are different from jailbreaks. jailbreaking, the adversary's goal is bypass safety filters instilled within the LLM (Willison 2023 Kim et al., 2024). A jailbreak may be considered successful if it elicit any unsafe response from the LLM. While on the other hand, the goal of hijacking attack is to force the model to generate specific outputs desired by the adversary (Bailey et al. 2023), which could potentially be unsafe outputs, in which case the hijacking attack would be considered a jailbreak as well. good analogy for jailbreaks and hijack attacks is untargeted and targeted adversarial attacks as studiod in the context of image classification [Liu et al. 2016
4 Robustness of Transformers to Hijacking Attacks
In this sention, we first provide a theoretical result characterizing the lack of robustness of linear transformers to single-token hijacking attacks. We then empirically investigate the robustness GPTT2 style transformers trained to solve linear regression in-context, and the efficacy of adversarial training towards improving the robustness of these transformers.
4.1 Single-Layer Linear Transformers
We first consider robustness of a linear transformer trained to solve linear regression in-context. As reviewed previously in the Scction 3.1 this setup has been considered in several prior works (von Oswald et al.) 2022; Zhang et al. 2024 Ahn et al. 2023 who all show that linear transformers learn to solve linear regression problems in-context by implementing (preconditioned) step of a gradient descent. We build on this prior work to show that the solution learned by linear transformers is highly non-robust and that an adversary can hijack a lincar transformer with very minimal perturbations to the in-context training set. Specifically, we show that throughout the training Irajectory, an adversary can force the linear transformer to make any prediction it would like by simply adding a single (Iadv Badv) pair to the input sequence. We provide a constructive proof of this theorem in Appendix
Theorem 4.1. Let L 12 O and let fisa(- :((()) be the linear transformer trained by gradient flow on the population loss using the initialization of Zhang et al. (2024), and denote ((xo) as the infinite-tine time limit of gradient flow. For any time E R, U {oo} and prompt P (I1, Ui..... FM UM Equery) with Tquery ~ N(O, I), for any Doad E R, the following holds.
1. If Fadv Ia), there crists Vadv Uadu P. Equery Ubad Fadv) E R s.t. with probability 1 over the draws of Jadu; Squery; by replacing any single ezample (I), i M, with (Fadus Vadv), the oulpul the perturbed prompt Padv satisfies
((t)) = Vbad-
I Note that standard transformers with the Interleave tokenization, k-token attack corresponds 2k takens being manipulated (see
5

---

Published in Transactions on Machine Learning Research (08/2025)
2. If Yadv # O, there crists Tadv Zadv (t, P., quuery, Ubad, Vadv) R s.t. with probability 1 over the draw of Mquuer; by replacing any single ezample (Ii,, Vi), i S M, with (Nadv: Vadu), the oulput on the perturbed prompt Padv) satisfies (querr([((Paaa 0(t)) = Ybad-
Theorem 4.1 demonstrates that throughout the training trajectory, by adding single (Iadvs Madv) token adversary can force the transformer to make any prediction the adversary would like. Moreover, the (Fadv: Madv) pair can be chosen so that either Tadv is in-distribution (i.e., has the same distribution as the training data and other in-context examples) or Vadv is in-distribution. We provide explicit formulas for each of these attacks in the Appendix (see (17) and (18).
At a high level, the non-robustness of the linear transformer plementing a learning algorithm one step gradient step - that generalizes well but is inherently non-robust. At more mechanistic level, this non-robustness can be altributed to the learned in-context algorithm' 's inability to identify and remove outliers from the prompt. This property is shared by many learning algo- rithms for regression problems: for instance, ordinary least squares, as an algorithm which is linear in the labels y, can also be shown to suffer similar problems as the linear transformer outlined in Theorem While non-robustness of the transformers to hijacking attacks has been established in prior works (Qiang et al. 2023 Bailey et al. 2023), this is the first result that provides a mechanistic explanation as to why transformers are vulnerable to hijacking attacks.
consequence of the linear transformer im-
Tn Section 5.2 we empirically verify that our attack can successfully force linear transformer output the adversary's is desired output Ybad, and evaluate its transferability to GPT-2 style transformers as well.
4.2
GPT-2 Style Standard Transformers
In this section, we empirically investigate the robustness of GPT2-style standard transformers to gradient- based hijacking allacks, and whether adversarial training (during pre-training or by fine-tuning) can improve the robustness of the transformers.
Metrics: To evaluate the impact of our adversarial attacks, we use two metrics: ground truth error (GTE), and largeled allack error (TAE). Ground-truth error measures mean-squared error (MSE) between the trans- former's prediction on the corrupted prompt Padv and the ground-truth prediction, i.c., Velean = w Quuery Targeted attack error similarly measures mean-squared error (MSE) between the transformer's prediction on the corrupted prompt and Ybad. Let y be the transformer's prediction corresponding Equery: then:
B
B
Ground Truth Error
Σ
(0: -Velean) )
Targeted Attack Error =
(ū) -bbad)
(7)
B
B
1=1
Attack Methodology: We choose Ybad according to the following formula,
Ybad = (1 - a1 Equery + aw) Zquery
(8)
Here wy is the underlying weight vector corresponding to the clean prompt P and wi 1 w, and a e lid. [0, 1] is a parameter. We construct W1 by first sampling a w' IN N(O, and applying Gram-Schmidt orthogonalization process to w' to make it orthogonal to wy. The resultant vector then normalized to have T.2 norm of When a -> 0, the target label Pbad is more similar to the in-distribution ground truth, while a + represents a label which is more out-of-distribution.
To carry out gradient-based attack, we randomly select a Aeeat number of input examples where Kreet is specified beforchand-and and initialize their values to zero. We then optimize these kiest tokens by minimizing the targeted attack error, for target Ybad from ( for different values of a E (0, Both during training and testing, we set the sequence length of the transformer to be 40. See Appendix C.3 for further details on attack procedure.
Based on whether we are perturbing 'features' or 'labels' within the input prompt, we denote our attacks as ei- ther feature-attack label-attack. When simultaneously attack both the features and label, call

---

Published in Transactions on Machine Learning Research (08/2025)
that joint-attack. Specifically, given prompt P (I1, ".....,......................... M yas, Fquery): for feature-attack. we replace (21, v1) with (Fadv: v1), for label-attack we replace (I1, with (I1, Madv), and for joint-attack, we replace (I1,31) 21) with (æadv: Madv).
Results: We find that gradient based adversarial attacks are generally successful. Our main results appear in Figure | under the label Atrain = O, which show the targeted attack error for an layer transformer averaged over 1000 prompts and 3 random initialization seeds when a from ( We note that for feature-attack, an adversary can achieve very small targeted attack error with perturbing just single token. However, for label-attack achieving low targeted attack generally requires perturbing multiple y-tokens. Note that this is in contrast with lincar transformers, for which we have previously shown that hijacking is possible with perturbing just single y-token. Finally, joint-attach behave in qualitatively similar way to feature-attack but are slightly more effective (this is most notable for keest = 1). Additional experiments investigating different choices of a appear in Appendix B.3
4.3 Adversarial Training
A common tactic to promote adversarial robustness of neural networks is to subject them to adversarial training i.e., train them adversarially perturbed samples (Madry et al. 2018). In our setup, create adversarially perturbed samples by carrying out the gradient- gradient-based attack outlined in Section 4.2 on the model undergoing training. Namely, for the model f at time i, for each standard prompt P, we take a target adversarial label Ybad and use the gradient-based attacks from Section 4.2 to construct an adversarial prompt Padv-
We consider two types of setups for adversarial training. In the first setup, train the transformer model from scratch on adversarially perturbed prompts. We call this adversarial pretraining. In the second sctup, we first train the transformer model on standard (non-adversarial) prompls P for Ti number of steps; and
88
attack 10
rrgeet SErICE
16
28
20
28
Clean 1
4
5
Clean 3
2
3
Clean 1 2
takens attacked
(Krnn)
of tokens attacked (Ruee
of tokens attacked (butl)
(a) aeature-attack.
(b) label-attack.
(c) joint-attack.
Kuain
Krain=
Ktrain=3
A-PT
A-FT
Figure 1: For both adversarial pretraining (A-PT) and fine-tuning (A-FT) against label robustness
against label-attack improves significantly, especially when trained budget of ktrain 3 perturbed tokens. The results are shown for 8 layer transformers with GPT-2 architecture.
40
ottack 86
rrge Back errict
28
28
Clean 1
4
3
Clean 1
Clean 1
2
5
8
takens attacked
(Kor)
of tokens attacked (Roee)
of tokens
attacked Ubet
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Kurinn
Kuain=
Ktrain=3
A-PT
A-FT
Figure 2: For both adversarial pretraining (A-PT) and fine-tuning (A-FT) against feature-attack. robust- ness against feature-attack and joint-attack improves for 7+ token attacks when trained on hirain The results are shown for 8 layer transformers with GPT22 architecture.

---

Published in Transactions on Machine Learning Research (08/2025)
then further train the transformer model for T2 number of steps adversarial prompts. We call this setup adversarial fine-tuning. In our experiments, unless otherwise specified, we perform adversarial pretraining for 10 steps. For adversarial fine-tuning, we perform 5 10 steps of standard training and then 10 steps of adversarial training, i.c., T1 = 10 and T2 = 10
The adversarial target value Ybad is constructed by sampling weight vector w ~ N(O, I) independent of the parameters , which determine the labels for the task т and setting 1). To keep training efficient, for
0:10
Kown 1
Knan
Ubad w wT Zquery (i.e., each task we perform 5 gradient steps to construct the ad- versarial prompt. We denote the number of tokens attacked during training with ktriin and experiment with two values of Ataann and kuain = 3. Unless stated otherwise, we use an layer transformer.
0.05
0.00
y--A
No
Figure 3: While there a moderate trade- off between robustness and (clean) accuracy when training against label-attack, the tradeoff is very small for feature-attace and joint-attack training.
Adversarial training improves robustness even with only fine-tuning. In Figures [ and 2 show the ro- bustness of transformers under k-token hijacking attacks when they are adversarially trained on either Keature-attack or label-attack. We see that adversarial training against at- tacks of a fixed type (e.g. feature-attack or label-attack) improves robustness to hijacking attacks of the same type, with robustness under feature-attack seeing particular improvement. Notably, there is little difference between adversarial fine-tuning and pretraining, showing little benefit from the increased compute requirement of adversarial pretraining.
Adversarial training against one attack model moderately improves robustness against another. Following adversarial training against label-attack, we see modest improvement in the robustness against feature-attack and joint-attack, while adversarial training against feature-attack results in significant improvement. against joint-attack (as expected, given that 20 of the 21 dimensions joint-attack uses
is shared by feature-attack) and modest improvement against label-attace We show in Fig. 12 the
results for adversarial training against joint-
Adversarial training against k-token attacks can lead to robustness against K' > token attacks. In both
Target attack error
40
Target attack error
40
Secuence Length
Sequence Twighh 101
Secuence Length 261
Fig. and 2 (as well as Fig. 12 we see
29
-20
that training against k = 3 token at-
tacks can lead to significant robustness against k = 7 token attacks, especially in the case of models trained against. feature-attack and joint-attack.
Clean2
15
$ of tokens attacked (Hue)
& of tokens attacked (Rima)
Figure 4: Larger context lengths improve robustness for fixed number of tokens attacked, but not for a fixed propor- lion. The number of layers is kept fixed at 8 while varying the context length. The results are averaged across three inde- pendently trained transformers (with identical hyperparameters). The shaded area indicates standard error.
Minimal accuracy vs. robustness tradeoff. In many supervised learning problems, there is an inherent tradcoff between the robustness of model and its (non-robust) accuracy (Zhang et al. 2019). In Fig. 3 we compare the perfor- mance of models which undergo adver- sarial training vs. those which do not, and we find that while there is a moderate tradeoff when undergoing label-attack training, there is little tradeoff when undergoing feature-attack and joint-attack train- ing.
On the whole, given the challenging nature of robust regression problem Diakonikolas & Kane 2019 the success of adversarial training is both surprising and remarkable, and hints at the ability of transformers to solve highly challenging non-convex optimization problems in context.
8

---

Published in Transactions on Machine Learning Research (08/2025)
Hodel
Source Model
16
0.0
25.8
34.7
25.1
43.4
32.8
43.1
3.1 0.0 2.0
3.3
0.0
1.4
12
1.2
18.5
21.6
8.8
4.3
Target Model
Target Model
Source Model.
8 633.3
15.9
0.0
5.9
4.3
5.0
14.3
(a) layers.
(b) 6 layers.
633
25.4
11.8
8.8
6.4
11.6
13.9
10.3
19.5
Source Model
1.2 7.3 12.1
Source Model
0.4
39.8
31.5
61.7
28.2
3.2
3.3
2.4
8.8
6.9
23.3
8.3
8.9
11.4 5.3 8.1
15.3
28.8
0.8
16
2
Target
Target Model
Target
(c) 12 layers.
(d) 16 layers.
(e) Source: L layers, Target: L' layers.
Figure 5: Targeted attack error when transferring an attack from a source model to a target models. Attacks transfer better between smaller-scale models, but not to larger-scale models (right) even across random seeds (left). Adversarial samples were generated using cnature-attack with k 3.
4.4 Effect of Scaling Depth and Context Length
Some recent works indicate that larger neural networks are naturally more robust adversarial attacks (Bartoldson et al. 2024 Howe et al. 2024). However, we did not observe any consistent improvement in adversarial robustness of in-context learning in transformers in setup with scaling of the number of layers, can be seen in Figure g in the appendix.
We also studied the effect sequence length, which scales the size of the in- context training set. We show in Figure g that for a fixed number of tokens attacked, longer context lengths can improve the robustness to hijacking attacks. However, for a fixed proportion of the context length altacked, the robustness to hijacking attacks approximately the same across context lengths. Note that in this setup the transformers are evaluated the same sequence length with which they were trained with (i.c., M M N). We explore this in more detail in the appendix (sce Appendix B.2
5 Transferability of Hijacking Attacks
Intuitively, two predictive models are similar, they should be vulnerable to similar targeted adversarial allacks. Thus, evaluating whether hijacking attacks transfer between two predictive models learned from distinct learning algorithms can help evaluate what extent the two models (and their underlying learning algorithms) could be considered similar. Note that we are specifically interested in targeted transfer; i.e., want adversarial samples generated by attacking source model to predict Yoad to also cause vietim model to predict Ybad- Transfer of targeted attacks is generally less common than the transfer of untargeted attacks (Liu et al. 2016).
In this section, we first analyze the transferability of adversarial attacks between different transformers, and surprisingly find that adversarial attacks can fail to transfer between even transformers of the same architecture and trained with the same rocipe-their sole difference being the random scods used for training. Next, we investigate the transfer from linear transformers to GPT-2 style transformers, and find that these attacks do not transfer. Lastly, we investigate the transfer between ordinary least squares (OLS) and GPT-2 style transformers, and find that the hijacking attacks transfer poorly between them.
5.1 Transfer Between Different Transformers
Lastly, we look the transferability of the adversarial attacks between different GPT-2 style transformers trained to solve linear regression in-context.
9

---

Published in Transactions on Machine Learning Research (08/2025)
Ground truth errot
66
Target attack error
66
66
Target attack error
8
Ground truth
+6
20
20
20
(a) feature-attack
(b) feature-attack
(c) label-attack
(d) label-attack
Linear Transformer
Transformer (layers=2) Transformer (layers=8)
Transformer (layers=16)
Figure 6: Robustness of different SGD-trained transformers when using attacks constructed from the infinite- time gradient flow solution via Theorem 4.1 for different, larget values 1/aa = (1 - a)w" Vquery + arw] (query; where w1 1 w. While these attacks reduce ground truth error across all model classes, the targeted attack error is only small for the lincar transformer. Shaded arca is standard crror.
aa)""xqury
For brevity, we restrict our focus to feature-attack here; transferability of label-attack follows a simi- lar pattern and discussed in Appendix [.4] We first consider within-class transfer, i.e., transfer from one transformer another transformer with identical architecture but trained from a different random initializa- tion. We present these results in Fig. (aad), averaged over 1000 prompts. We generate adversarial samples according to the methodology described in Sec. 4.2 While sampling Ybad, we use a = 1. Diagonal entries correspond to telf-transfer, i.e., TAE of the generated adversarial samples for the source model. These results show that for transformers with smaller capacities (3 and layers) attacks transfer quite well, but transfers become progressively worse as the models become larger. This suggests that higher- transformers could implement different in-context learning algorithms when trained from different sceds.
We next consider across-class Imnsser, i.e., transfer between transformers with different number of layers. We present these results in Fig, F(e), for one randomly chosen source and target model with given number of layers. The results are also averaged over 1000 prompts. These results show similar trend as within-class transfer: atlacks from small-to-medium capacity models transfer better to her small-to-medium capacily models, while larger capacity models transfer poorly to all other capacity models.
5.2 Transfer Between Linear Transformers and Standard Transformers
We implement separate feature-attack and label-attace based on formulas given in equa- tions I1 and 18). Recall that these formulas presume that the network being attacked is a linear trans- former which implements one step of gradient descent when trained to solve linear regression in context. As done previously, Mbad is chosen according to
In Figure@wee show the robustness of tGD-trai single-layer linear transformers and standard transformers of different depths as function of (ef. (8) These results averaged over 1000 different samples and 3 random initialization seeds for every model type (see Appendix C for further details on training). We find that the gradient flow-derived attacks transfer to the SGD-trained single-layer linear transformers, as the targeted attack error is near zero for all values of a. Moreover, while standard (GPT2) transformers trained to solve lincar regression in-context incur significant ground-truth error when the prompts are perturbed using the altacks from Theorem 4.1 these allacks are not successful as largeled altacks, since the targeted error is large. This behavior persists across GPT2 architectures of different depths. As linear transformers implement step of gradient doscent, poor transfcrability of those attacks suggests that when trained linear regression tasks, GPT2 architectures do not implement one step of gradient descent, as has been suggested in some prior works fvon Oswald et al. 2022
5.3 Transfer Between Transformers and Least Squares Solver
It has been argued that transformers trained to solve linear regression in-context implement ordinary least squares (OLS) Garg et 2022 Akyürek al.) (202)). so, adversarial (hijacking) attacks ought to transfer between transformers and OL.S. In Figure we show mean squared error (MSE) between predictions of OL.S
10

---

Published in Transactions on Machine Learning Research (08/2025)
Between Prediction
40
2.0
MSE Between Prodict Prediction
40
1.3
30
20
20
CLean
0.1
0.3
0.5
8.7
10
0
Clean
0.1 0.3
0.5
0.7
1.0
Clean
0.1
0.3
0.5
0.7
1.0
a
a
(a) OLS + Transformers.
(b) Transformers OLS.
Layers - 16
Layers- 12
Layers-8
Layers-6
Layers
As the targeted prediction Ybad becomes more out-of-distribution (a + 1; cf. S we see (b) attacks derived from transformers result in significantly different predictions than that of OLS, while (a) attacks derived from OLS result similar predictions for some classes transformers but not all.
Figure 7:
and transformers on adversarial samples created by performing feature-attack on OLS and transformers respectively. can be clearly observed that as the targeted prediction Ybad becomes more out-of-distribution (a + 1), MSE between predictions made by OLS and transformers also increases. Furthermore, MSE is considerably larger when adversarial samples are created by altacking transformers. This collectively indicates that the alignment between OLS and transformers is weaker out-of-distribution and that the transformers likely have additional adversarial vulnerabilities relative to OLS. We provide additional results and expanded discussion in Appendix B.5
6 Discussion
This study makes several contributions to understanding how transformers perform in-context learning, particularly concerning adversarial robustness and the nature of the learned algorithmss
We conclusively establish that neither linear transformers nor standard GPT-2 style transformers inher- ently learn adversarially robust in-context learning algorithms. Specifically, within lincar transformers, this vulnerability arises because they implement standard, non-robust learning algorithm. This finding pro- vides mechanistic insight into the adversarial non-robustness of transformers previously observed in other works (Qiang et al. 2023 Bailey et al. 2023 On a positive note, we empirically demonstrate that adver- sarial training can improve robustness to hijacking attacks, and this improved robustness shows degree of generalization. This an encouraging result, especially as robust. regression against. adaptive adversaries is known be challenging (Diakonikolas Kane 2019).
Furthermore, analysis of the transferability of adversarial attacks reveals two novel phenomena. Firstly, hijacking attacks may fail to transfer effectively across larger transformer models trained with identical procedures but different random seeds, suggesting a non-universality of in-context learning mechanisms even within a fixed architecture and data distribution. Secondly, adversarial samples sourced from traditional linear regression algorithms (like ordinary least squares and gradient descent) are largely ineffective against GPT-2 style transformers, and vice versa. This provides indirect evidence that the in-context lcarning algorithms acquired by GPT-2 might difler qualitatively from these traditional learning algorithms.
6.1 Limitations
While the improvement in robustness through adversarial training promising, the relative simplicity of the lincar regression task used our study means these results should be interpreted with caution. Further investigations on more complex tasks are required fully assess the capability of transformer architectures to learn adversarially robust learning algorithms. Understanding and 'reverse-engineering the specific al-
11

---

Published in Transactions on Machine Learning Research (08/2025)
gorithms transformers implement in these robust settings could also provide novel insights for algorithm design.
An important limitation of our transferability analysis is its non-mechanistic nature. While the differing behavior on adversarial (out-of-distribution sample strongly hints at qualitative differences between learned algorithms, it leaves open questions about the precise mechanistic nature of these differences. Specifically, it remains unclear whether these differences arise from different parameterizations of same underlying al- gorithm, transformers learning variety of different types of learning algorithms, due to some other reason.
6.2 Future Work
Prior work has shown that gradient descent on neural network paramcters can exhibit an implicit bias towards solutions that generalize well but lack adversarial robustness (Frei al. 2023). Our work shows that a similar bias exists for linear transformers trained solve linear regression in-context. Future work could investigate the extent to which this bias affects the in-context learning algorithms discovered by transformers in general.
Our work, on the whole, highlights that developing a thorough understanding of in-context learning within transformers may be more challenging than previously thought, and emphasises the need of developing mechanistic understanding of these transformers. Further mechanistic analyses of these transformers is needed to better understand what in-context lcarning algorithms these transformers learn. Such reverse engineering efforts could also help inspire better algorithm design for various robust learning problems.
References
Kwangjun Ahm, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement precon- ditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? invostigations with lincar models. arXiv preprint arXiv:2211.1560 2022.
Usman Anwar, Abulhair Saparov, Javier Rando, Daniel Palcka, Miles Turpin, Peter Hase, Ekdeop Singh Lubana, Erik Jenner, Stephen Casper, Oliver Sourbut, et al. Foundational challenges assuring alignment and safety of large language models. arXiv preprint or Xiv:24404 00932, 2024.
Matthew Ashman, Cristiana Diaconu, Adrian Weller, and Richard E Turner. In-context in-context lcarning with transformer neural processes. arXio preprinl or Xiv:2406.. 13493, 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Provable in- context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
Luke Bailey, Euan Ong, Stuart Russell, and Scott Emmons. Image hijacks: Adversarial images can control generative models at runtime. arXiv preprint arXim:2309. 00000063, 2023.
Amir Bar, Yossi Gandelsman, Trevor Darrell, Amir Globerson, and Alexei Efros. Visual prompting via image inpainting. Advances in Neural Information Processing Systems, 35:25005- 25:25005-25017, 2022.
Brian R Bartoldson, James Diffenderfer, Konstantinos Parasyris, and Bhavya Kailkhura. Adversarial ro-
bustness limits via scaling-law and human-alignment studies. arXin preprint arXiw:2404.09340 2024.
K. Bhatia, P. Jain, and Kar. Robust regression via hard thresholding. In Advances Neural Information Processing Systems 28, pp. 721-729, 2015.
K. Bhatia, P. Jain, Kamalaruban, and P. Kar. Consistent robust regression. In Advances in Neural Information. Processing Systems 30, pp. 2110-2119, 2017.
12

---

Published in Transactions on Machine Learning Research (08/2025)
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877 1901,,991,111 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Loe, Aditya Grover, Misha Laskin, Picter Abbeel, Ar- avind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15099 2021.
Yeshwanth Cherapananjeri, Efe Aras, Nilesh Tripuraneni, Michael I Jordan, Nicolas Flammarion, and Pe- ter L Bartlett. Optimal robust linear regression in nearly linear time. arXiv preprint orXiv:2007.08133 081177 2020.
Ilias Diakonikolas and Daniel M Kane, Recent advances in algorithmic high-dimensional robust statistics. arXiv preprint arXiv: 1.05911, 2019.
Tlass Diakonikolas and Daniel M. Kane. Algorithmic Higl--iieensional Robust Statislics. Cambridge Uni- versity Press, 2023.
Tlis Diakonikolas, Weihao Kong, and Alistair Stewart. Ellicient algorithms and lower bounds for robust linear regression. In Proceedings of the Thirtieth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 2745 2754. SIAM, 2019.
Ahmad Elawady, Gunjan Chhablani, Ram Ramrakhya, Karmesh Yadav, Dhruv Batra, Zsolt Kira, and Andrew Szot. Relic: A recipe for 64k steps of in-context. reinforcement learning for embodied ai. arXio preprint arXit:2410.02755 2024.
Spencer Frei, Gal Vardi, Peter L. Bartlett, and Nathan Srebro. The double-edged sword of implicit bias: Generalization vs. robustness in relu networks. In Advances in Neural Information Processing Systems (NeurIPS). 2023.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17008 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in- context? a case study of simple function classes, Advances Neural Information Processing Systems, 35: 30583-30598, 2022.
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In International Conference on Learning Representations, 2015.
Pengfei He, Han Xu, Yue Xing, Hui Liu, Makoto Yamada, and Jiliang Tang. Data poisoning for in-context learning. arXiv preprint arXiv:2402.02100 021600 2024.
Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. Artificial Neural Networks ICANN 2001: International Conference Vienna, Austria, August 21 -25, 2001 Proceedings 11, pp. 87 94. Springer, 2001.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: A transformer that solves small tabular classification problems in a second. arXiv preprint ar XxXx:2207.018866 2022.
Nikolhaus Howe, Michal Zajac, Ian McKenzie, Oskar Hollinsworth, Tom Tseng, Pierre-Lur Bacon, and Adam Gleave. Exploring scaling trends 10m robustness. arXio preprint orXi:2207.118 00..4077.11,,, 2024.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autorogressive transformers with lincar attention. In International conference on machine learning, pp. 5156 5165. PMLR, 2020.
Tacyoun Kim, Suhas Kotha, and Aditi Raghmattaan Jailbreaking is best solved by definition. arXiv preprint arXiv:2403. 14725, 2024.
13

---

Published in Transactions on Machine Learning Research (08/2025)
Louis Kirsch and Jürgen Schmidhuber. Meta learning backpropagation and improving it. Advonces in Neural Information. Processing Systems, 14:14122-141344 2021.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprind 0rXiv:2212.04455 2022.
Adam Klivans, Pravesh K Kothari, and Raghu Meka. Efficient algorithms for outlier-robust regression. In Conference On Learning Theory, pp.
PMLR, 2018.
Yanpei Liu, Xinyun Chen, Chang Liu, and Dawn Song. Delving into transferable adversarial examples and black-box attacks. arXiv preprint
2016.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu, Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Represendations); 2018.
Harsha Nori, Yin Tat Lee, Sheng Zhang, Dean Carignan, Richard Edgar, Nicolo Fusi, Nicholas King, Jonathan Larson, Yuanzhi Li, Weishung Liu, et al. Can generalist foundation models outcompete special- purpose tuning? case study in medicine. arXiv preprint ar arXft:2811. 2023.
Yao Qiang, Xiangyu Zhou, and Dongxiao Zhu. Hijacking large language models via adversarial in-context learning. arXiv preprint arXiv:2311.09943 2023.
Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikacl Henaff, and Roberta Railcamu. Gencralization new sequential decision making tasks with in-context learning. arXiv preprinl arXiv:2312.0380 038001 2023.
James Requeima, John Bronskill, Dami Choi, Richard Turner, and David Duvenaud. Llm processes: Numerical predictive distributions conditioned on natural language. arXiv preprint arXiv:2405. 12856, 2024.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight program- mers. In International Conference on Machine Learning, pp. 9355-9366. PMTR, 2021.
Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative recurrent nets. Neural Computation, 1992.
A. S. Suggala, K. Bhatia, P. Ravikumar, and P. Jain, Adaptive hard thresholding for near-optimal consistent robust regression. In Proceedings of the Sherty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 2892 2897. PMLR, 2019.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoope- hand, Nathalic Bradley-Schmice Michael Chang, Natalic Clay, Adrian Collister, et al. timesae adaptation in an open-ended task space. arXio preprinl arXin:2301. 200.0000 2023.
Max Vladymyrov, Johannes Von Oswald, Mark Sandler, and Rong Ge. Lincar transformers are versatile in-context learners. arXiv preprint arXiv:2408.14180 2024.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, An- drey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677, 2022.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel 7 Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXio preprint arXiv:1611.05700 2016.
Simon Willison. Multi-modal prompt injection, 2023. http:://shoovilllss.ee//////////
mul ti paaimmmmoll-rooot--njettiiooooot Accessed on: August 20, 2024.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? Preprint, 0rXiv:2310.08391, 2023.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric P. Xing, Laurent El Ghaoui, and Michael I. Jordan. The- oretically principled trade-off between robustness and accuracy. In International Conference on Machine Learning (ICML), 2019.
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Yuanhan Zhang, Kaiyang Zhou, and Ziwei Liu. What makes good examples for visual in-context learning? Advances in Neural Information Processing Systems, 36:17773-17794 2023.
15

---

Published in Transactions on Machine Learning Research (08/2025)
Appendix
A Proofs
Notation: We denote [n] (1, We write the inner product of two matrices A, B E R B) tr(ABT). We use O, and 0mxn to denote the zero vector and zero matrix of size and x n, respectively. We denote the matrix operator norm and Frobenius norm as 1l-112 and 11-1 We use Ia to denote the d-dimensional identity matrix and sometimes we also use when the dimension is clear from the context.
Setup: As described in the main text, we consider the setting linear transformers trained on in-context examples of linear models, a sctting considered in number of prior theoretical works on transformers (von Oswald al. 2022 Akyürek al. 2022 Zhang et. al., 2024 Ahn et al. 2023 Wu el. al. 2023). Let 24 E Rd and Y4 E I. For prompt P = (I1, "...,,.... UN, 2N+1), say its length is N. For this prompt, use
embedding which stacks (Ii, vi)
E Rd+1 into the first N columns with (x)+1,0) E RR++1 the last
column:
= B(P)
F1
2N43 IN+1 UN 0 0+)) E R(d+1)((41))
EN
(9)
1/2
We consider single-layer linear self-attention (LSA) model, which is modified version of attention where remove the softmax nonlincarity, merge the projection and value matrices into single matrix WPV e Rd+1,d+1, and merge the query and key matrices into single matrix WKQ E RR+11d+1 Denote the set of parameters as 0 = (W kQ, WPV) and let
E BKR
(ssA(E;)) = E ++ WWVV E.
(10)
N
The network's prediction for the query example IN+1 is the bottom-right entry of matrix output by fisA. (Euery(E: 0) =
We may occasionally use an abuse of notation by writing (query(E; as (query(P) or Jquery with the under- standing that the transformer always forms predictions by embedding the prompt into the matrix E and always depends upon the parameters 0.
We assume training prompts are sampled as follows. Let A be a positive definite covariance matrix. Each training prompt, indexed by t E N, takes the form of P. = (ær,1 Wr.I...,,TT,,, Br.N), Fx,,), where task weights wy iid. N(O, Ta), inputs Fr,s sid. N(O, and labels y=. = (wr, The empirical risk over B independent prompts is defined
[(0)
1
Σ
B
(T.N+1(Er:0)) - (w,, (wN)))
(11)
2B
We consider the behavior of gradient flow-trained networks over the population loss in the infinite task limit
B + 00:
1 E -. w...N( 2r, Rid N(OAA) [ [0r,N+1(EF: -(rr (r,N+1))
L(0) lim 2(0)
(12)
B
Note that we consider the infinite task limit, but each task has finite set of N i.i.d. (æi, v:) pairs. We consider the setting where fLSA is trained by gradient flow on the population loss above. Gradient flow captures the behavior of gradient descent with infinitesimal step size and has dynamics d 0 -V
We repeat Theorem 4.1 from the main section for convenience.
Theorem 4.1. Let t 2 O and let fLsA(- ((()) be the linear transformer trained by gradient flow on the population loss using the initialization of Zhang al. (2024), and denote ((00) as the infinite-time limit of gradient flow. For any time E R, U {oo} and prompt (I1. V1. ......... . FM, Mdd Iquery) with Tquery N(O, I), for any Yoad R, the following holds.
16

---

Published in Transactions on Machine Learning Research (08/2025)
1. If aadv N(O,Ia) there crists Vadv Yadv (t, P. ..uueyy Ybad. Tadv) E R s.t. with probability 1 over the draws of Jadv; squery; by replacing any single example (Ii Us), i is M, with (Tadvs Vadv),, the output the perturbed prompt Padv satisfies Bquery (M(eey(((Paav):::) ((t)) = Vbad-
2. If Yadv # 0, there ezists Zadv Zadult, P, Equery, ybad, Vadv) Rd s.t. with probability 1 over the draw of Nquery; by replacing any single ezample (I), i I M, with (Fadv: Vadu), the output on the perturbed prompt Padv) satisfies Bquery! (E((adv)) A(t)) = Vbad-
Proof. Let us partition the Wkq and WPV matrices as,
W (w21))
KQ W12 KQ
W.P 11
PV PV W22
wko =
WPV
KQJ
(w221 PV )T
W22
KQ where W WHY E Rdxd, W12 w wi21 embedding matrix E with M + columns,
KQ
WP E Rd, and W22
KQ
way
E R. Then by definition, for an
EE M
WKQ KQ (w22
Jquery (E;0) ((
)
Equery-
(13)
Due to the lack of positional encoders, the prediction is the same when replacing (2aR, yx) with (Iadvs Madu) for any k, so for notational simplicity of the proof we will consider the case of replacing (21,31) with (Tadus Vadu ) So, let consider the embedding corresponding to (Tadus Madv: 22, 1/....,, IM MM, :query), so that
[aadv adv
Ei= T400 + "'query query Vadv adv + >i== y1.0
yadv + + 9134
EE
=
M
Vadv + -i=2 Vi
M
Expanding, we have
(w
PV) M
M
Jquary (E:0)
Iadv adv
T
Σ
x + L quary
WKO QQ Quuery
2=
(wFy))
VadvWadd
Σ
yiit
(w
(O) Iquery
M
W22
W KQ
Yadv adv
+
Equery
M
1=2
PV
M
W22
Vadv
+
(w
KQ T
Equery
M
2=2
When training by gradient flow over the population using the initialization of (Zhang et al. 2024 Assumption 3.3), by Temmas C.I, C.5, and C.6 of (Zhang et. al. (2024) we know that for all times R4 U (x)), it KQ
PV (t) + 0. In particular, the prediction holds that w (t) = wP (t) = w21 (t) = 0 and W (t) + 0 and w formula above simplifies to
PV Waz
(t) M
Equery (E; 0())
yadvil odv +
Σ 2m
yorr T
(1)aquary-
(14)
For notational simplicity let us denote W(t) = w PV (t)W
so that
M
1
((E = M
Vadv adv +
Σ yoor 2m
W (t) (q)ury
The goal to take Mead E R and find (æadv: Vodv) such that (EE; (t)) Ybad- Rewriting the above equation that this is equivalent to finding (Fadv; Madv) such that.
17

---

Published in Transactions on Machine Learning Research (08/2025)
Vadv adv W (t)xEquryy M
Ybad
1
Σ yar W (t)xquer)
(15)
M
From here we see that if W (u)rryy o 0 then by setting
MW (quury
T addv
Ybad
Σ y: W qt)arr
(16)
[[ (t) ?
we guarantee that (( Ybad- By Zhang el al. (2024 Lemmas A.3 and A.4), we know W(I) o 0 for all t. Since ((t) * and Zquery N(O, independent of Wt),, we know (ft) a.s. Therefore the
identity (16) suffices for constructing adversarial tokens, and indeed for any choice of Yadv # 0 this directly
allows for constructing x- a-based adversarial tokens,
MW (WT(())auua Qqury
aavv
Vadv W (t) 1|2
Ybad
M Σ ya W(t)xrqury)
(17)
On the other hand, if we want to construct adversarial token by solely changing the label y, we can
return to 15). Clearly, as long as aadv W Qt)xeu * O, then dividing both sides by this quantity allows for
solving Yadv- we assume Tadv another in-distribution independent N(0, I) sample, then since W(t) * 0
guaranlees that II ((1 qu)rry * 0 and so we can construct.
M
Ybad
M Ei yeut W (t) (f)rquerry
Vadv
(18)
2 adv
W Wuu)xqqeryy
18

---

Published in Transactions on Machine Learning Research (08/2025)
B
Additional Results
In this section, we provide additional results that expand on some of the results given in the main text.
B.1 Effect of Scale
We conducted experiments with transformers with different number of layers to evaluate whether scale has any elfect on adversarial robustness of the transformer not. We observed no meaningful improvement in the adversarial robustness of the transformers with increase in the number of layers. This is shown in the figure below for Ybad chosen with a = See Section 4.4 in the main text for relevant discussion.
40
Ground truth
20
28
Clean
of tokens attacked (Kaza)
tokess attacked (Kim.)
(a) feature-attack.
(b) label-attack.
40
48
sttack
20
28
Clean
5
Clean
3
6
attacked
(c) feature-attack.
(d) label-attack.
Layers=82
Layers=04
Layers=08
Layers=16
Layers=03
Layers=06
Layers=12
Figure 8: Increasing the scale of the transformer does not improve the adversarial robustness of in-context learning in transformers.
B.2 Effect of Sequence Length
We show here the complete set results, for both feature-attack and label-attack, on how an increase in sequence length positively impacts adversarial robustness if adversary can manipulate the same number of tokens (for all sequence lengths), but if the adversary can manipulate the same proportion of tokens (which would amount to different number of tokens for different sequence lengths), increase in sequence length has negligble effect on the adversarial robustness. See Section 4.4 in the main text for relevant discussion.
19

---

Published in Transactions on Machine Learning Research (08/2025)
49
Ground truch
20
Clean
Clean 2
7
10
12 (M
(a) feature-attack
(b) feature-attack
10
48
Target attack
20
28
Clean
Clean 2
12 of tokens attacked (Kaw)
of tokons attacked (Rowe)
(c) label-attack.
(d) label-attack.
Sequence Length 41
Sequence Length
101
Sequence Length 201
Sequence Length 481
Figure 9: Effect of increase in sequence length.
B.3 Gradient- Based Adversarial Attacks & Adversarial Training
In the main text (in Sections 4.2 and 4.3
gave results for atlacks performed with ybad chosen by setting a = 1 in equation 8 Here, we present results for a = 0.5 and a 0.1. These results are qualitatively similar to the case of a and are presented only for completeness. Furthermore, the main text, we showed only target attack error for our attacks due to space constraints, while here we present results for both ground truth error and target attack error.
B.3.1
a 1.0
Ground truth errar
40
Ground truth error
48
Ground truth
20
20
20
Clean 1
3
5
Clean
2
4
5
6
Clean
1
2
3
5
6
# tokens attacked (Krett
of tokens
attacked
(Ktest
tokens
attakked (Keest
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attaakk error
attack errror
Target attack error
48
40
40
20
Clean 1
3
5
Clean
2
4
6
Clean
1
3
5
tokens attacked (krma)
of tokens attacked (Kres))
tokens attacked (Kres)
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
0
=1rain=1
Ktrain
A-PT
A-FT
Figure 10: Adversarial training against label-attack. A-PT denotes adversarial pretraining and A-FT denotes adversarial finctuning. Atrai denotes the number of tokens attacked during training and Ktrain corresponds to model that has not undergone adversarial training all.
20

---

Published in Transactions on Machine Learning Research (08/2025)
Ground truth error
48
Ground truth arror
40
Ground truth
20
Clean 1
3
5
Clean
2
4
5
6
7
Clean 1
3
4
5
7
# of tokens attacked (Kuet)
# of tokens
attacked (Klea)
# of tokens attacked
(Kuet)
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack errer
attack errar
Target attack errer
28
28
28
Clean 1
3
5
Clean
2
4
Clean 1
3
5
# of tokens attacked (Keest)
of tokens attacked (Kiezt)
# of tokens attacked (Keest
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
Ktrain=0 0
1train=1
Ktrain =3
A-PT
- --- A-FT
Figure 11: Adversarial training against
Ground truth errer
40
Ground truth crror
40
Ground truth
20
20
20
Clean 1
3
5
Clean
2
4
5
6
7
3 Target attack crror
Clean 1
# of tokens attacked (Keett
# of tokens attacked (ktest
# of tokens attacked (Kies)
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack error
48
Target attack errar
48
48
28
20
20
Clean 1
3
5
Clean
2
4
5
6
Clean 1
3
5
# of tokens attacked (kuet)
of tokens attacked (Kue))
# of tokens attacked (kuss)
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
0
=1
Ktrain=3
A-PT --- A-FT
Figure 12: Adversarial training against joint-attack.
21

---

Published in Transactions on Machine Learning Research (08/2025)
B.3.2
0.5
Ground truth error
16
Ground truth error
10
Ground truth
Clean 1
3
5
Clean
2
4
5
6
Clean 1
2 3
5
6 (Keett
# tokens attacked
(Krest)
of tokens
attacked
(Krest)
# tokens
attacked
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack error
Target attack error
Target attack error
10
10
10
Clean 1
3
5
Clean
2
4
5
6
Clean
1
3
5
tokens attacked (Kuett
of tokens attacked (Klest)
tokens attacked (Kues)
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
ktrain=0 0
=train=1
ktrain=3
A-PT
A-FT
Figure 13: Adversarial training against label-attack.
Ground truth error
10
10
Ground truth error
10
Ground truth
Clean
3
5
Clean
2
Clean
3
5
# of tokens attacked
(Kiest)
# of tokens attacked (kue))
# of tokens attacked (Kiest
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack errar
Target ottack error
Target attack crrar
16
10
10
Clean 1
3
5
Clean
2
4
Clean 1
3
5
# of tokens attacked (Kuea)
# of tokens attakee (kuez)
# of tokens attacked (Kuea)
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
0train= 0
A-PT Figure 14: Adversarial training against feature-attack.
=1rain=1
Ktrain=3
--- A-FT
22

---

Published in Transactions on Machine Learning Research (08/2025)
Ground truth errer
16
Ground truth errer
10
Ground truth
Clean
3
Clean
Clean
3
5
tokens attacked
(Krest)
of tokens attacked (Krest)
tokens attacked
(Kies)
(a) leature-attack.
(b) label-atlack.
(c) joint-altack.
Target attack error
16
Target attack error
Target sttack error
16
Clean tokens attacked (Kiest)
Clean
2
Clean
3
of tokens attacked (Krest)
tokens attacked (Kies))
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
00
ktrain=1
ktrain=3
A-PT
A-FT
Figure 15: Adversarial training against joint-attack.
B.3.3
0.1
Ground truth errer
0.4
Ground truth errror
0.4
Ground truth error
0.4
0.2
0.2
0.0
0.0
Clean
1
2
4
5
Clean
2
3
4
6
Clean 1
2
4
5
7
tokens
attacked (kwa)
of tokens attacked
(kmar)
tokens
attacked (Kua)
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack error
0.6
Target attack error
8.6
Target attack error
0.6
0.4
0.2
0.2
0.2
0.0
0.0
0.0
Clean
1
2
4
5
Clean
3
6
Clean 1
2
4
5
# of tokens
attacked (Krest)
of tokens attacked (Krest)
# of tokens
attacked (Krett
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
0
=1rain=1
ktrain=3
A-PT
A-FT
Figure 16: Adversarial training against label-attack A-PT denotes adversarial pretraining and A-FTT
dcnotes adversarial finctuning. ktain denotes the number of tokens attacked during training and Brran = 0 corresponds to a model that has not undergone adversarial training at all.
23

---

A-PT Published in Transactions on Machine Learning Research (08/2025)
Ground truth error
0.4
Ground truth error
0.4
0.4
Ground truth
0.2
0.2
Clean 1
2
4
5
Clean
2
3
4
6
Clean 1
2
4
5
# of tokens
attacked (Kuett
of tokens attacked
(Kest)
# of tokens
attacked (Kust)
(a) feature-attack.
(b) label-attack.
(c) joint-attack.
Target attack error
Target attack error
Target attack error
0.4
0.4
0.4
0.2
8.2
0.2
0.0
0.0
Clean 1
2
4
5
Clean
3
6
Clean 1 2
4
5
# of tokens attacked (Kiest)
# of tokens attacked
(Kiet)
# of tokens
attacked (Keet)
(d) feature-attack.
(e) label-attack.
(f) joint-attack.
0
11
k3rain=3
A-PT
A-FT
Figure 17: Adversarial training against Stature-attack.
Ground truth crror
0.4
Ground truth error
Ground truth error
0.4
8.4
0.2
0.2
8.2
Clean 1
2
4
5
Clean
3
4
6
Clean 1
2
4 attacked (Kust)
5
# of tokens
attacked (Kwst)
of tokens attacked
(Kuat)
# of tokens
(a) feature-attack.
(b) label-atlack.
(c) joint-altack.
Target attack errar
Target attack error
Target attack error
0.4
0.4
0.4
8.2
8.2
0.2
0.0
0.0
0.0
Clean 1
2
4
5
Clean
3
6
Clean 1
2
4
5
# of tokens
attacked (Krest
of tokens attacked
(Krest)
# of tokens
attacked (Kiest
(d) feature-attack.
(e) label-attack.
(I) joint-attack.
00
=1rain=1
ktrain=3
A-FT
Figure 18: Adversarial training against joint-attack
B.4 Transferability
In Section 5.1 we briefly presented some results around transfer of adversarial examples generated using one transformer to other transformers either with the same architecture or different architecture. We present complete results here, for both feature-attack and label-attack. As in the main text, we first present results for transfer across same class of transformers, i.c., transformers with same number of layers and then present results for transfer across different classes of transformers.
24

---

Published in Transactions on Machine Learning Research (08/2025)
Source Model.
0.0 1.7
3.7
Source Model
0.0
1.5
1.5
Source Model
0.0
2.7
4.5
Source Model
0.0
2.8
2.5
Source Modet
0.0
2.2
2.9
Source Model
1.2
7.3
12.1
Source Model
0.4
39.8
31.5
1.3 0.8 3.3
3.1
0.0
2.0
3.1
8.1
2.1
3.3
0.0
1.4
3.7
0.0
3.2
13.7
8.1
11.4
21.8
0.0
26.6
3.5 4.3 8.0
2.2
1.7 8.0
3.7
1.7 0.1
4.3
3.0
0.0
3.6
4.1
8.1
11.4
5.3
0.1
15.3
28.8 0.0
Target
Target Model
Target Model
Target Model
Target Model
Target
Nodel
Target Model
(a) layers.
(b) 3 layers.
(c) layers.
(d) 6 layers.
(e) layers.
(f) 12 layers.
(g) 16 layers.
Figure 19: Target Attack Error for different target models on adversarial samples generated using source model with the same number of layers. Adversarial samples were generated using feature-attack with k = 3. Transfer of adversarial samples across transformers progressively becomes poorer as number of layers increases.
Source Modet
Source Modet
Source Model
Source Model
Source Modet
Source Model
Source Model
29.7
3.3
1.2 0.5
8.1
8.3
1.0
1.9
0.2
1.3
6.9
8.9
24.3
3.5 3.4 2.3
1.4 1.5
0.4
0.9
0.8
6.3
1.2
1.6 0.2
4.8
3.9
1.1
7.8
4.7
1.2
17.8
38.5
0.5
Target Model
Target Model
Target Model
Target Model.
Target Model
Target Nodel
Target Model
(a) layers.
(b) 3 layers.
(c) 4 layers.
(d) 6 layers.
(e)
layers.
(I)
12 layers.
(g) 16 layers.
Figure 20: Same as above figure 19 but adversarial samples were generated using label-attack with
k
As with feature-attack, transfer of adversarial samples samples across transformors progressively becomes poorer as number of layers increases.
16
16
40
12
12
12
12
8
8
6
6
Source Model
Source Model
8
40
20
10
2
16 16316661688 12
8
6
4 4
3
3
16 16666177
12
6
2 2
Target Model
Target Model
(a)
(b)
Figure 21: Target Attack Error for different target models on adversarial samples possibly generated ing source model with different number of layers. In (a) adversarial samples were generated using feature-attack with k = 3. In (b) adversarial samples were generated using label-attack with k = 7. Transfer is generally worse when
25

---

Published in Transactions on Machine Learning Research (08/2025)
B.5 Hijacking Attacks Ordinary Least Square
Linear regression can be solved using ordinary least square. This solution can be written in closed-form follow:
y f(x, Fquery) = (x xx X querry
(19)
where X [zj : :.;:;::]]ll :>X] and Y= .y...... yNl. We implement a gradient based adversarial attack on this solver by using Jax autograd to calculate the gradients V Vxf(X,,; (queery and f Y, "query) Similar to our gradient-based attack on the transformer, X and Y. In OLS, X and Y are not lokenized, however, for consistency of language, we will continue to refer to the individual entries of these matrices, ie., I,, Yi as tokens. We perform 1000 iterations and use a lcarning rate of 0.01 for both feature-attack and label-attack.
only update randomly chosen subset of entries withing
Figure 22 shows results for feature-attack and y-attack respectively on OLS for Ybad chosen by using a 1.0. The adversarial robustness of OLS is qualitatively similar to that of the transformer; for fixed compute budget, single-token label-attass are much less successful compared to single-token beaurr-attack, and larget altack error is lower when greater number of lokens are allacked.
We further look at the transfer of adversarial attacks between transformers and OLS. Specifically, by attack- ing OLS we create a set of adversarial samples and then mcasure the mcan squared error (MSE) between the predictions OLS and different transformers on these adversarial samples, and vice versa. Figure 23 shows the transfer for adversarial samples for different values of a for sampling Ybad- For feature-attack, we attack 3 indices and for y-attack, we attack 7 indices. We can make following observations from this increases. This indicates lack of alignment. between the predictions made by OLS and transformers OOD. (ii) For feature-attack, MSE between predictions is significantly lower when adversarial samples are sourced by attacking OLS rel- alive to when adversarial samples are sourced by attacking the transformers. In other words, adversarial samples transfer better from OLS to transformers but not vice versa. This hints at the fact that adversarial robustness of the transformers worse than that of OLS. (iii) For y-attack, the aforementioned asymmetry in transfer above does not exist except for transformers with layers 16 and 12. (iv) Finally, we note that transformer with 16 layers clearly always behaves in anomalous fashion, with transformers with layers 12 and also sometimes behaving anomalously, which is in line with the discussion in previous section intra-transforma transfer of adversarial samples.
figure: (i) the predictions made by OLS and transformers tend to diverge
In Figure 24 we present complementary results showing MSE between predictions of OLS and transformers on adversarial samples when different number of tokens are attacked for a = 1.0. These results further support the observations made the previous paragraph.
26

---

Published in Transactions on Machine Learning Research (08/2025)
0.04
Target Attack Error
12
18
0.03
Target Attack
0.02
8.01
8.00
Cleen
Clean
of tokens attacked (Klett
of takens attacked (Ktest)
(a) feature-attack.
(b) label-attack.
Figure 22: The adversarial robustness ordinary least squares to gradient-base based hijacking attacks is quali- tatively similar that of the transformers.
MSE Between Prodiction
40
48
38
2.0
28
0.0
9.5
18
Clean
0.1
0.3
0.5
0.7
1.0
(a) feature-attack: formers.
OLS Trans-
(b) feature-attack:
Transformers
4 OLS.
HSE Betveon Predictisn
38
Clean
0.1
0.3
0.5
0.7
1.0
(c) label-attack:
OLS Trans-
(d) label-attack:
Transformers
formers.
OLS.
Layers 16
Layers-12
L
-
Layers-2
Figure 23: The mean squared error between the predictions being made by the transformer and OLS on adversarial samples tends to increase as the 'OOD-ness' of the Ybad increases. Furthermore, the difference in prediction is generally higher when the hijacking altacks are derived using the transformer (notice the differences in scale). For feature-attack we attack 3 tokens and for y-attack we attack 7 tokens when creating adversarial samples.
27

---

Published in Transactions on Machine Learning Research (08/2025)
MSE Betwoon Prestetton
49
49
20
Clean
2
3
5
Clean 2
of tokens attacked (ktest
# of tokens attacked (Kes))
(a) feature-attack: OLS Trans-
(b)
feature-atlack:
Transformers
formers.
>
OLS.
VSE Betwenn Prediction
Trediction
10
10
0
Clean
2
6
Clean
1
2
of tokens
attacked (kta))
tokens
attacked (Krest)
(c) lttack:
OLS Trans-
(d) label-attack:
Transformers
formers.
OLS.
Layers- 16
Layers- 12
Layers-8 8
-
Layers -2
Figure 24: The mcan squared error between the prodictions being made by the transformer and OLS on adversarial samples lends to be higher when the adversarial samples are sourced by attacking transformers. In the above plot, we use 1.0 for sampling Yoad-
28

---

Published in Transactions on Machine Learning Research (08/2025)
C Training Details and Hyperparameters
C.1 Linear Transformer
To match the setup considered in Theorem 4.1 we implement linear transformer single-layer attention- We train the linear transformer for 2M steps with
only linear transformer described in equation 10 batchsize of 1024 and learning rate of 10-6
C.2 Standard Transformer
Our training setup closely mirrors that of Garg et al. (2022). Similar to their setup, we use a curriculum where we begin training on data with dimensionality incrementing it by every 5k steps. Details of our architecture are given in Table | We guve the number of parameters present in various transformer models with different number of layers in Table 2 Important
and gradually increase the dimensionality to 20 by
training hyperparameters mentioned in Table 3
Parameter
Value
Embedding Size 256
Number of heads
8
Positional Embedding
Learned
Number of Layers
8 (unless mentioned otherwise)
Causal Masking
Yes
Table 1: Architecture for the transformer model.
Number of Layers Parameter Count
1, 67,, 601
2, 46, 553
3, 253, 505
6
4, 833, 409
8
6, 413, 313
12
9, 573, 121
16
12, 732, 929
Table 2: Hyperparameters used for training transformer models with GPT-2 architecture.
Hyperparameter
Value
Learning Rate
x 10
Warmup Steps
20,000
Total Training Steps
500,000
Batch Size
64
Optimizer
Adam
Table 3: Hyperparameters used for training transformer models with GPT-2 architecture.
C.3 Adversarial Attack and Adversarial Training Details
We implement our adversarial attacks as simple gradient descent on the (selected) inputs with the target attack error as the optimization objective. We briefly experimented with variations gradient descent, e.g., gradient descent with momentum but found those to perform at par with simple gradient descent.
29

---

Published in Transactions on Machine Learning Research (08/2025)
When performing feature-attacia used learning rate of 1 and when performing label-attack,, we used a learning rate of 100. When performing joint-attack, we used learning rate of when perturbing x-tokens and a learning rate of 100 when perturbing y-tokens. We chose the learning rates based on best performance within 100 gradient steps. Using lower values learning rates resulted in proportionally slower convergence, and hence were avoided.
In all our plots, we show results across three different models and 1000 samples for each model.
Differences Between Adversarial Attacks and Adversarial Training: The two major differences in our adversarial traning setup, compared with adversarial attacks setup are:
During adversarial attacks (done on trained models at test time), sample ybad according to the expression
g but during adversarial training we sample Ubad by sampling weight vector N(0 Ia) independent of the task parameters W, and setting Ybad w Iquery
During adversarial attacks, we perform 100 steps of gradient descent, but in adversarial training, we only perform steps of gradient descent.
Both the above changes were done to help improve the efficiency of adversarial training.
30


Reviews:
[{"review_id": "ChRd8ufuFX", "reviewer": "Reviewer_E14B", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe paper raises very interesting and pertinent questions related to in-context learning and adversarial attacks. The paper provides extremely clear and rigorous technical details, conducts several experiments, and presents methodology and results clearly.\n\nWeaknesses:\n\nThe key weaknesses relate to the clarity of the written work rather than the technical rigor.\n\nIt seems like the paper has many competing goals:\n* Adversarial training increases robustness to hijacking attacks in in-context learning prompts.\n* Can controlled studies of hijacking attacks in prompts shed light on how in-context learning really works?\n* How does in-context learning really work?\nAll in the context of transformers for linear regression.\nWhile the technical rigor and experiments are extremely thorough, as the reader I am not sure of the scope of the paper or what the goal really is, so I cannot quite see how those experiments contribute to the larger picture/goal/motivation. The authors would need to clearly distinguish between which question(s) are fully answered by the experiments in this paper vs. which questions are speculative/left for future work.\n\nTMLR is a venue that brings together work from many different subfields related to, involving, or directly studying machine learning. As a result, I found the title to be confusing. I believe the title should clearly indicate (1) the research question being explored, and (2) the type of model being used. From the paper, the type of model seems to be very narrow and specific to linear regression transformers (transformers pre-trained on the numerical task of linear regression). If this is true, the title/initial framing of the work is too broad compared to what is studied, and the title should be modified to clarify what the paper is about. Examples include \"Probing In-Context Learning in Transformers for Linear Regression with Adversarial Attacks,\" or \"Robustness of Transformers in Linear Regression Tasks: In-Context Learning, Hijack Attacks, and Adversarial Training\" depending on the scope of the paper (which remains a little unclear).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper looks at how transformers do in-context learning by testing their robustness on simple linear regression tasks. The authors use hijacking attacks to see what kind of learning algorithms the models are really using. They show that otherwise similar transformers can learn different behaviors and that adversarial training can help make them more robust to hijack attacks. Overall, it can shed light on how these models actually process and generalize from context.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is vital to address the questions and concerns in the weaknesses above.\n\nIn addition, the abstract on openreview (not in the pdf) lacks spaces between each sentence, so this needs to be corrected.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Tg7YZhFfjV", "reviewer": "Reviewer_xyqN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- Interdisciplinary: connects ICL of linear functions, adversarial attacks, and robust algorithm design\n- Clear mathematical setup: provides an easy-to-follow formulation of task embeddings and attack variants\n- Novel approach: suggests using adversarial vulnerability and transferability as tools towards probing ICL mechanisms\n- Empirical thoroughness: extensive numerical experiments across multiple architectures and attack types\n- Demonstrates effective adversarial training strategies for attack mitigation with an initial analysis of performance trade-offs\n\n**Weaknesses**\n- Interpretation issues: the claim that poor transfer indicates different algorithms may not be fully supported by the evidence\n- Organizational issues: figures presented out of logical order with terms used before introduction (see requested changes)\n- Limited parameter exploration: all experiments restricted to N=40 and d=20, even with very large GPT-2 type models\n- Notational issues: several undefined terms and notation shifts complicate understanding (see questions)\n- Attack design scope (minor): focuses primarily on gradient-based adversarial vector constructions without exploring alternatives (see questions)\"", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper makes two contributions to the study of in-context learning (ICL) capabilities of transformers in simplified scenarios (linear models).\n\n**#1 - Vulnerability to hijacking attacks and mitigation strategy:**\n- Provides a formal proof that linear transformers can be hijacked via a single token perturbation\n- Empirically demonstrates that conventional GPT-2 style transformers are likewise vulnerable\n- Shows that adversarial training improves robustness in GPT-2 models\n\n**#2 - Comparative analysis of attack transfer properties:**\n- Finds attacks transfer well between smaller transformers but poorly between large transformers trained with different seeds\n- Shows attacks transfer poorly between classical algorithms, small linear transformers, and large transformers\n- Interprets these observations as evidence for qualitatively different underlying ICL mechanisms (see 'Requested changes' below)\"", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Critical Adjustments** \n\n*A. Issue with claim #2 : \"Failure to transfer implies different mechanisms\"*  \nIn Section 5, the authors conclude that poor transfer of adversarial examples between transformers trained with different seeds (or betwween architectures/solvers) indicates that they implement qualitatively different in-context learning algorithms. However, this conclusion is not sufficiently supported by the evidence presented. \n\nNaively, poor transfer of adversarial examples could result from factors besides qualitative algorithmic differences, such as different parameter configurations implementing the same algorithm or different internal representations of the task embedding. Since the attacks rely on gradient-based optimization in input space, and these gradients are highly sensitive to model weights, it is expected that adversarial examples may not transfer even if the underlying learning procedure is functionally similar.\n\nAdditional evidence for their interpretation should be provided, or alternatively present this idea much more tentatively as a hypothesis rather than a conclusion.\n\n*B. The text is clear and straightforward up until section 4; the remaining sections and figures should be organized for readability*  \n- Present figures in logical order consistent with the text narrative\n- Ensure all terms are introduced before appearing in figures\n- Example issues:\n  - Section 4.1 (p6) redirects to Fig. 6 before Fig. 1 is introduced\n  - Fig. 6 references concepts (ground truth error, \\alpha, w_{perp}) before they are defined\n  - Fig. 1 (p7) uses unexplained terms (A-PT, A-FT, k_train)\n\n*C. Many small clarifications*  \n1. Context length variations (Section 4.3): Is only the test context length (M) being varied, or also training length (N)? \n2. Fig. 4: What do the error bars denote?\n3. Fig. 5: Why are the diagonals not always zero? Figure 2 suggests feature attack with k=3 gives ~0 error for 8 layers. Why are diagonals for 16-layer models lower than for 12-layer models? It would be helpful to understand how these scale on average. \n4. Clarify adversarial training parameters:\n- Specify in Section 4.4 whether adversarial training/tuning uses \\alpha=1 or other values\n- Consistently document experimental parameters across main text and appendices\n5. In Appendix A (p15), you introduce h(x) without definition. Similarly, in Eq. 13 (p16), the notation w21, etc. appears without prior definition.\n6. In Appendix A (p16) stating \"Due to the linear attention structure...\", would the same property not hold for softmax attention without positional embeddings? Is this property specific to linear attention or to the lack of positional information?\n7. Section 5.2, I am not sure that the following is claimed in these three works: \"..., GPT2 architectures do not implement one step of gradient descent, as has been suggested in some prior works (von Oswald et al., 2022; Ahn et al., 2023; Zhang et al., 2024)\"\n8. Fig. 6 caption: mention that it is inf-time solution of gradient flow (if that is the case). \n\n\n**Non-critical Improvements**\n1. Provide code to reproduce experiments.\n2. Explore alternative attack constructions: from the proof of Theorem 4.1, one can consider alternative constructions of x_adv beyond those parallel to W(t) x_query\n3. Add comparison to random feature attacks: e.g. for Figure 6, could compare GTE with random vector attacks of the same magnitude; this would serve as a valuable bridge to connect to black-box attack scenarios\n4. Explanation of training curriculum (mentioned in C.2, p28)\n5. Minor points and typos:\n- typos (e.g. p2 , line 3; p29 y_bad = w^T y_bad; p9 y_{b}ad)\n- section 3 - some redundancy in the text, e.g. the same notion appears twice \"Since transformers output a sequence of tokens of the same length and dimension as their input, ...\"\n- section 3 - maybe the two embedding styles should be denoted differently somehow, E_{...} (P)?\n6. The baseline Target Attack Error (TAE) appears to be approximately 40 in many figures. How should it scale / what determines it in the experiments -- is it expected to be ~ 2 d \\alpha^2? Relatedly, Ground Truth Error (GTE) and Target Attack Error (TAE) often appear 'mirrored' in the supplementary figures. Is there a simple way to understand this approiximate tradeoff? How should we interpret this relationship, particularly as \\alpha changes?", "broader_impact_concerns": "The authors note (p3) that research on hijacking attacks of LLMs is still limited. Given the growing importance and public availability of these tools, it is worth considering the impact of research reporting potential susceptibilities of existing architectures. See e.g. the Impact section of Bailey et al. 2023. Despite the white-box setup and mathematical focus of the present work, it might be valuable for the authors to discuss this briefly in a Broader Impact section.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kizr36matS", "reviewer": "Reviewer_hMTR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The scope of this paper is wide, encompassing both theoretical and empirical studies and discussion of different aspects related to adversarial attacks (e.g., adversarial training, transferrability).\n2. The empirical studies are fairly comprehensive, and many different aspects are studied.\n3. The topics of in-context learning and hijacking attacks are of broad interest, and combining them is a good direction of research.\n\nWeaknesses:\n1. There are some serious presentation issues with the theoretical statements. See the requested changes section.\n2. There seems like a gap between the theoretical part and the empirical part of the paper, which is not carefully bridged. It would be interesting to see how (or if) the theory is useful in practice.\n3. While the linear regression setting makes the theoretical analysis feasible, it is not clear if it is important for empirical analyses and if the same empirical observations transfer to a more complicated setting.\n4. The paper contains too many ideas and distinct pieces of study. It is hard to grab the gist of it. A clear summary would be very helpful.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper analyzes the effects of adversarial attacks in the context of in-context learning using transformers. It considers a linear regression setting, where each datapoint consists of many input-output pairs sampled from a linear model with its own weight. An adversarial attack is introduced to drive the model's prediction to a particular output, by changing the inputs and/or the outputs of one or more conditioned vectors. A theoretical result on linear attention models is first presented to show the theoretical feasibility of such an adversarial attack. Then, an empirical study is done on GPT-2, followed by a discussion and an empirical study showing that adversarial training improves the robustness of the model. Finally, the transferability of hijacking attacks is studied and results suggest a lack of transferability, leading to a rethinking of prior characterization of in-context learning capabilities of transformers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. While the current abstract is fine and technically correct, I suggest compressing it and highlighting the important aspects of the paper to make it more engaging.\n2. Eq. (1) and (3) both use $E(P)$. Does it make sense to make the notations distinct so that it is clear which tokenization is used every time it appears in a formula?\n3. Eq. (4) does not look correct to me. My guess is that the author(s) meant $$\\frac{1}{N}\\sum\\_{i=1}^N ([f^{\\text{Mask}}(E\\_\\tau; \\theta)]\\_{d+1,2i+1} - y\\_{\\tau, i+1})^2.$$ Can you confirm and correct it if there is indeed a typo?\n4. The statement made in Theorem 4.1 has fundamental issues. It sounds like the $y\\_{\\text{adv}}$ in the first statement depends only on $t$ and nothing else. This sounds too strong and hard to believe. In fact, checking out eq. (17), we obviously have that $y\\_{\\text{adv}}$ also depends on $x\\_{\\text{adv}}$, the replacement index, and other unchanged inputs and outputs. This should be made very clear in the theorem. The same applies to the second statement.\n5. In eq. (8), there are infinitely many choices for $w\\_\\perp$. Which one did you use and why? Also, what is the norm of $w\\_\\perp$ and does it align with $w$?\n6. For experiments with adversarial training (e.g., Figure 1), can you also show how the convergence rate is affected by the noisy training data?\n7. For adversarial training, what if I do not know $P\\_{\\text{adv}}$ a priori? Is there a generally applicable method to still improve the model's robustness?\n8. There are several typos: (i) In section 1, some En Dashes (--) need to be changed into Em Dashes (---). (ii) In eq. (7), the $1/B$ on the left is larger than the $1/B$ on the right. (iii) In section 5.1, I cannot tell what $y\\_{b}ad$ is but I assume this is a typo.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)