On the Robustness of Transformers against Context Hijacking for Linear Classification
Authors: Tianle Li
Chenyang Zhang
Xingwu Chen
Yuan Cao
Difan Zou
Venue: NeurIPS 2025 poster ()
Keywords: in-context learning
transformers
robustness
deep learning theory
learning theory
Primary area: theory

Abstract: Transformer-based Large Language Models (LLMs) have demonstrated powerful in-context learning capabilities. However, their predictions can be disrupted by factually correct context, a phenomenon known as context hijacking, revealing a significant robustness issue. To understand this phenomenon theoretically, we explore an in-context linear classification problem based on recent advances in linear transformers. In our setup, context tokens are designed as factually correct query-answer pairs, where the queries are similar to the final query but have opposite labels. Then, we develop a general theoretical analysis on the robustness of the linear transformers, which is formulated as a function of the model depth, training context lengths, and number of hijacking context tokens. A key finding is that a well-trained deeper transformer can achieve higher robustness, which aligns with empirical observations. We show that this improvement arises because deeper layers enable more fine-grained optimization steps, effectively mitigating interference from context hijacking. This is also well supported by our numerical and real-world experiments. Our findings provide theoretical insights into the benefits of deeper architectures and contribute to enhancing the understanding of transformer architectures.

On the Robustness of Transformers against Context Hijacking for Linear Classification
Tianle Li'; Chenyang Zhang"; Xingwu Chen', Yuan Cao², Difan Zou 1.2
Xingwu
Institute of Data Science, The University of Hong Kong
School of Computing & Data Science, The University of Hong Kong (tianleli, chyzhang, xingau){cconectt. hkue. hk {yuancao, dzou) hk
Abstract
Transformer based Large Language Models (LLMs) have demonstrated powerful in-context learning capabilities. However, their predictions can be disrupted by factually correct context, a phenomenon known as context hijacking, revealing a significant robustness issue. To understand this phenomenon theoretically, we explore in-context linear classification problem based recent advances in linear transformers. In our setup, context tokens are designed as factually correct query-answer pairs, where the queries are similar to the final query but have opposite labels. Then, we develop a general theoretical analysis on the robustness of the linear transformers, which is formulated training context lengths, and number of hijacking context tokens. A key finding is that a well-trained deeper transformer can achieve higher robustness, which aligns with empirical observations. We show that this improvement arises because deeper layers enable more fine-grained optimization steps, effectively mitigating interference from context hijacking. This is also well supported by our numerical and wall-wrldd experiments. Our findings provide theoretical insights into the benefits of deeper architectures and contribute to enhancing the understanding of transformer architectures.
function of the model depth,
Introduction
Transformers 667 have demonstrated remarkable capabilities in various fields of deep learning, such as natural language processing [60] [I] [68, 66 [52] 26]. A common view of the superior performance that is, transformers can
of transformers lies in its remarkable in-context learning ability 14
18 44
flexibly adjust predictions based on additional data given in context contained in the input sequence itself, without updating parameters. This impressive ability has triggered a series of theoretical studies attempting to understand the in-context learning mechanism of transformers [51] 30 77. 31 76).
These studies suggest that transformers can behave as meta learners [18] implementing certain meta algorithms (such as gradient descent [69 2 85 based on context examples, and then applying these
algorithms to the queried input.
Despite the benefits of in-context learning abilities in transformers, this feature can also lead to certain negative impacts. Specifically, while well-designed in-context prompts can help generate desired responses, they can also mislcad the transformer into producing incorrect or even harmful outputs, raising significant concerns about the robustness of transformers
21
45
86 For instance,
significant body of work focuses on jailbreaking attacks [5
50
62
25
81
which aim to design
"qqual contribution. Corresponniig author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Model
Context Input
Output
lesk_Lang-ap
Tesk
All models
Rafael Nadal is best sport Is
tennis
TEGK, Sports
Task Tesk_Cky
GFT2-SMALL (12 Layers)
2x [Rafael Madal good playing basketball). Rafael Nadall is best sport t5
basketbell
14 of
30
GPT2-MEDIUM [24 Layers)
BK [Rafael Nadal is not (food at playing backetball.
Rafael Nadal best sport 1s
basketbell
GPT2-LARGE (36 Layers)
[Rs(Rafeel Nadal good playing basketball) Rafael Nadal a best sport 18
basketball
GPTZ-SMALL GPTZ-MEDIUM GATZE
XFT2XU
Model
Figure 1: Context hijacking phenomenon in LLMs of different depths. Left: there are no or only a few factually correct prepends, LLMs of different depths can correctly predict the next token. When the number of prepends increases, the outputs of models are disrupted. Right: Four different types of tasks are introduced, each with fixed template, and tested on LLMs of different depths. The horizontal axis is the model with depth from small to large, and the vertical axis is the average number of prepends required to successfully interfere with the model output. Experiments show that deeper models perform more robustly. (Experimental setup given Appendix H.I
specific context prompts that bypass the defense mechanisms of large language models (LLMs) to produce answers to dangerous or harmful questions (c.g.. "how to build a bomb?"). It has been demonstrated that, as long as the context prompt is sufficiently long and flexible to be adjusted, almost all LLMs can be successfully attacked (6). These studies can be categorized under adversarial robustness, where attacker is allowed to perturb the contextual inputs arbitrarily induce the
transformer model to generate targeted erroneous outputs 663 53 22 80
However, in addition to the adversarial attack that may use harmful or incorrect context examples, it has been shown that the predictions of LLMs can also be disrupted by harmless and factually correct context. Such a phenomenon is referred to context hijacking [37 (36), which is primarily discovered on fact retrieval tasks, i.c. the output of the LLMs can be simply manipulated by modifying the context with additional factual information. For example, shown in Figure [ the GPT2 model can correctly answer the question "Rafael Nadal's 's best sport is" with "tennis" when giving context examples. However, if factually correct context examples such "Rafael Nadal is not good at playing basketball" are provided before the question, the GPT-2 model may incorrectly respond with "basketball". Then, it is interesting investigate whether such phenomenon depends on different tasks and transformer architectures. To this end, we developed a class of context hijacking tasks and counted the number of context examples that led to incorrect outputs (see Figure D]. Our findings indicate that increasing the number of prepended context examples amplifies the effect on the transformer's prediction, making more likely to generate incorrect outputs. Additionally, we observed that deeper transformer models exhibit higher robustness to context hijacking, requiring more prepended context examples alter the model's 's output. Therefore. conducting a precise robustness analysis regarding context hijacking could provide valuable insights in understanding the architecture of the transformer model.
In this paper, we aim to develop a comprehensive theoretical analysis on the robustness of transformer against context hijacking. In particular, we follow the general design of previous theoretical works 2 (2) on the in-context learning of transformers. by considering the multi-layer linear transformer models for linear classification tasks, where the hijacking examples are designed as the data on the boundary but with an opposite label to the queried input. Starting from the view that the I--laee transformer models can implement L-step gradient descent on the context examples, with an arbitrary initialization, we formulate the transformer training as finding the optimal multi-step gradient descent methods with respect to the learning rates and initialization. Then, we prove the optimal multi-step gradient strategy, and formulate the optimal learning rate and initialization as the function of the iteration number (i.e., model depth) and the context length. Furthermore, we deliver the theoretical analysis on the robustness based on the proved optimal gradient descent strategy, which shows that as the transformer become deeper, the corresponding more fine-grained optimization steps can be less affected by the hijacking examples, thus leading to higher robustness. This is well aligned with the empirical findings and validated by our numerical and real-world experiments. We summarize the main contributions of this paper follows:
2

---

We develop the first theoretical framework to study the robustness of multi-layer transformer model against context hijacking, where the hijacked context example is designed as the data with the factually correct label but close to the prediction boundary. This is different from a very recent rel related work on the robustness of transformer perturbed, which could be factually incorrect.
that allows the context data to be arbitrarily
Based on the developed theoretical framework, we formulate the test robust accuracy of the transformer function with respect to the training context length, number of hijacked context examples, and the depth of the transformer model. The key of our analysis is that we model the in-context learning mechanism of a well- trained multi-layer transformer an optimized multi-step gradient descent, where the corresponding optimal initialization and learning rates can be precisely characterized. This could of independent interest to other problems that involve the gradient descent methods on linear problems.
Based on the developed theoretical results, we demonstrate that deeper transformers are more robust because they are able to perform more fine-grained optimization steps on the context samples, which can potentially explain the practical observations of LLMs in the real world (see Figure []. The theoretical results are well supported by synthetic numerical experiments and real-world LLMs experiments in various settings.
Related Works. Our work is related to recent works on in-context learning via transformers. mechanism interpretability of transformers, and robustness of transformers. Due to space limit, defer them to Appendix[A] A
2 Preliminaries
2.1 Data model
To
understande understand the mechanism of context hijacking phenomenon, we model it as a binary classification task, where the query-answa pair is modeled as the input-response pair ((x, v) E R x (=1}). In particular, we present the definition of the data model as follows:
(Data distribution). Let w* E Rd be a vector drawn from a prior distribution on the d dimensional unit sphere S4-1, denoted by pa- (-), where f E Sd-1 denotes the expected direction of w*. Then given the generated w", the data pair (x, y) is generated as follows: the feature vector is X (O,,,, ( La) and the corresponding label is y = sign( ((w", x)).
Definition 2.1
Modeling natural language problems as linear problems is common setting in theoretical research, because linear problems have sufficient representation power, supported by many previous works [69 The innovation of our construction is that we precisely exploit the characteristics of the context hijacking phenomenon and model it as a semantic binary classification problem. Compared with previous works that directly study linear problems, such as linear regression, our modeling is more practical. Moreover, to extend our theoretical results to more complex situations, we also consider nonlinear models and nonlinear tasks, which is given in Appendix [.3]]
2
85 84
33
47
Based on the data distribution of each instance, we then introduce the detailed setup of the in-context learning task in our work. In particular, we consider the setting that the transformer is trained on the data with clean context examples and evaluated on the data with hijacked context.
Training phase. During the training phase,
given n clean context examples {(x1, ".......................... (xn, 3n)} and a query Xquery with its label Vquery- In particular, here we mean the clean examples the {(x), //),,.... (Xn, }n)) are drawn from the same data distribution Dw- as (Xquery: Mquery). Then, the input data matrix for in-context learning is designed as follows:
Z
=
X 91
Xn
Xquery
E g(d+1)x(x+1).
(2.1)
Vo
Here, to ensure that the dimension of Xquery aligns with those of other input pair (x; y: ) we concatenate it with 0 as a placeholder for unknown label yquery- Ideally, we anticipate that given the input Z, the output of the transformer model, denoted by Vquery can match the ground truth one. Moreover, we also emphasize that within each data matrix Z. the context examples and the queried data should be generated based on the same ground truth vector w*, while for different input matrices, e.g., Z and Z', we allow their corresponding ground truth vectors could be different, which are i.i.d. drawn from the prior p3- (.).
3

---

The training data distribution simulates the pre-training data of the large language model. Unlike existing works ST where the prior of w is assumed to have zero mean, we consider setting where w has a non-7ero mean (i.e., (3*). This approach is inspired by empirical observations (see Figure I that transformer models can perform accurate zero-shot predictions. Consequently, our model can encapsulate both memorization and in-context learning, where the former corresponds to recovering the mean of the prior distribution, i.c., 3*, and the latter aims to manipulate the {(x,, VI).... (Xn, }n)} effeettivlyy. In contrast, existing works primarily focus on the latter, thereby failing to fully explain the interplay between memorization and in-context learning.
Test phase. During the test phase, context examples are designed based on the query input Xquery to effectively execute the attack. Inspired by empirical observations (Figure and prior experience with jailbreaking attacks [6]. we choose to use repeated hijacking context examples during the test phase. Specifically, since the hijacked context should be factually correct, we consider data similar to the queried input but with correct and opposite label of low confidence. Mathematically, this involves projecting Xquery onto the classification boundary. To this end, given the target query data (Xquery, Vquery), ). we formalize the design of the hijacked context example as follows.
Definition 2.2 (Hijacked context data). Let (x, y) be input pair and w' be the corresponding ground truth vector. Additionally, denote x as the projection of x on the boundary of classifier, i.e.
XI (Ia - (w*())))))))) ) Then, the query pair (Xquery, Mquery) is generated as Xquery X 1 + Fow and yquery = sign ((w*, Xquery)) )) = sign(a) with being a random variable, and the hijacked context example is designed as Xhc = x_ and yhc = Yquery-
Note that we pick (Xhe: w") = 0 to enforce hijacked context lies on the boundary of the classifier. A more rigorous design is set Xhe x1 7) yquery w for some positive quantity n) where it can be clearly shown that yhe by enforcing 1) + 0+.
(Xte, w*)) = Vquery. Definition 2.2 concerns the limiting regime
Then, based on the above design, the input data matrix in the test phase is constructed as follows:
Zhe
Xho
Xhe
Xquery
R (x++1)
(2.2)
The
The 0
Here we use N to denote the number of hijacked context examples. The example (Xhes Uhe) can also be interpreted to the closest data to Xquery but with a different label -Vquery, which principally has the ability to perturb the prediction of Xquery. An innovation compared with previous work is that our data model design is semantic. For example, in Figure I we can assume that Xhe = "Rafael Nadal not good at playing", the = "basketball", Xquery = "Rafael Nadal's best sport is"" and Vquery "tennis". Additionally, because the prediction is highly likely to be correct in the zero-shot regime (i.e., N ), the prediction in the test phase can be viewed as a competition between model memorization and adversarial in-context learning. This dynamic is primarily influenced by the number of hijacked context examples.
2.2 Transformer model
Following the extensive prior theoretical works for transformer 84 85 16
27
2].
we consider
linear attention-only transformers, prevalent simplified strueture to investigate the behavior of transformer models. In particular, We define an L-layer linear transformer TF as a stack of L single- head linear attention-only layers. For the input matrix Z E [Rd++1)((++1,, the i-th single-head linear attention-only layer TF, updates the input as follows:
Z (Z.1) Z-1 + P. P.Z.MZZT,,,.... Q.Z,,11,
(2.3)
where M := In 9) ((++1)x(--1) is the mask matrix. We design this architecture to constrain
the model's focus to the first n in-context examples. Moreover, the matrix P == W. E (d+1)x(d++1) serves as the value matrix in the standard self-attention layer, while the matrix Q = W W, E I(d+1)((++1) consolidates the key matrix and query matrix. This mild re- parameterization has been widely considered in numerous recent theoretical works [34] 72 65 35 To adapt the transformer for solving the linear classification problem, we introduce an additional linear embedding layer We g E (d(+1))(d-11, Then the output of the transformer TF is defined as
Iquery TZo: WE, (Pr, Qr)

---

=- ..... IF1 We (Z(oo)((a)))6
- (Zrl((+1).(0+1))
(2.4)
Le. the negative of the (d + n + 1) entry of Zz, and this position replaced by 0 in the input Zo. The reason for taking the minus sign here is to align with previous work 669 62]. which will be explained in Proposition
2.3 Evaluation metrics
Based on the illustration regarding the transformer architecture, we first define the in-context learning risk of L-layer model TF in the training phase. In particular, let Der be the distribution of the input data matrix in (2.T) and the target Uquery, which covers the randomness of both (x, y) and w", then the risk function in the training phase is defined as:
R (TF) = Ëz Mquurry Drr (ZF 0) "qquery )²]
(2.5)
where 0 (Wa e: (Pr, ==) denotes the collection all trainable parameters of TF. This risk function will be leveraged for training the transformer models (where we use the stochastic gradient in the experiments).
Additionally, in the test phase, let Dio be the distribution of the input data matrix he in 2.2 and the target Vquery, we consider the following population prediction error:
(TF)
Pzue Buwery Dre
TF : !) Vqueery 0].
(2.6)
3
Main theory
In this section, we present how we establish our theoretical analysis framework regarding the robust- ness of transformers against context hijacking. In summary, can briefly sketch our framework into the following several steps:
Step 1. We establish the equivalence between the L-layer transformers and L steps gradient descent, converting the original problem of identifying well-trained transformers to the problem of finding the optimal parameters of gradient descent (i.e., initialization and learning rates).
Step 2. We derive the optimal learning rates and initialization of gradient descent, revealing its relationship with the number of layers 7. and training context length
Step 3. By formulating the classification error of a linear model obtained by L steps gradient descent with optimal parameters on hijacking distribution Dte. characterize how the number of layers L, the training context length n and test context length N affect the robustness.
3.1 Optimizing over in-context examples
which connects the in-context learning of transformer with the gradient descent algorithm, we follow similar approach by showing that, in the following proposition, multi-layer transformer can implement multi-step gradient descent, starting from any initialization, on the context examples.
Inspired by line of recent works 85
9
16 2
51
Proposition 3.1. For any L-layer hingle-head linear transformer, let Yquery be the output of the /-th layer of the transformer, i.c. the (d + 1, 1)-th entry of Ze. Then, there exists single-head linear (1)
<(1)
(l) transformer with layers such that Equery = (w gd x Xquery). )- Here, w (f),
"d are the parameter vectors
obtained by the following gradient descent iterative rule and the initialization w w!t can be arbitrary:
w (+++) (!)) T.Vī(w gd)
gd gd
where I(w) = 2
Σ ((w, xx) - -2:)2
(3.1)
1=1
Here Tj can be any d x matrix.
5

---

As H2 could be any d x d matrix, Proposition 3.1 demonstrates that the output of the L-layer transformers is equivalent to that of a linear model trained via L-steps of full-batch preconditioned gradient descent on the context examples, with (Ti [- being the learning rates. This suggests that each L-layer transformer defined in (2.4). with different parameters, can be viewed as an optimization process of a linear model characterized by a distinct set of initialization and learning rates {w ...... ---1}. Therefore, it suffice to directly find the optimal parameters of the gradient descent process, without needing to infer the specific parameters of the well-trained transformers.
Among all related works presenting similar conclusions that transformers can implement gradient descent, our result is general as we prove that transformers can implement multi-step gradient descent from any initialization. In comparison, for example, [85] shows that single-layer transformer with MLP can implement one-step gradient descent from non-zere initialization. [2] demonstrate that linear transformers can implement gradient descent, but only from 0 initialization.
3.2 Optimal multi-step gradient descent
Based on the discussion in the previous section, Proposition 3.1 problem of identifying the parameters of well-trained transformers into the task of finding the optimal learning rates and initialization for the gradient descent process (3.1). In this section, we present our conclusions regarding these optimal parameters. As we consider optimizing over the general training distribution Der. where the tokens x;'s follow the isotropic distribution, it follows that the updating step size should be equal in each direction from the perspective of expectation. Therefore we consider the case F Ia to simplify the problem, with og being scalar for all I (0..... L - 1}. In the following, we focus on the optimal set of parameters (w (0) gd 00; OL-1). Specifically, we
successfully transforms the original
consider the population loss for w(()
as gd
(L) R(w (L) Er,w' [(w [d Xquery) Vquery) )²}.
where T = ((x1,,1) (Xn; Vn); (Xquery, (query)) is the set of all classification pairs This definition resembles RTFF) defined in (2.5). We attempt to find the w(()) that minimizes this (0) population loss, along with the corresponding learning rates {or (eo and initialization w gd which can generate this w via gradient descent. We first present the following proposition demonstrating (L)
that there exists commutative invariance among the learning rates {ar)f for producing w "ri
Proposition 3.2. Let {ao, (i,,..,, ar.-1) ".,,...,a..111 } be set of learning rates, and {a6, ai,...... .2 be another set of learning rates that is a permutation of {oo, ai...., az-1}, meaning both sets contain the same elements, with the only difference being the order of these elements. With w gd R"
(L)
denoting as the parameters achieved by learning rates {ag, ar....,,...11 and w-( Rd as the gd'
parameters achieved by learning rates {aá: a',....,''_1) from the same initialization w (0) gd
it holds
that w (L.) w gd'
gd
Proposition 3.2 implies that the learning rates at different steps contribute equally to the overall optimization process. Consequently, we will consider a consistent learning rate a through the entire gradient descent procedure, which significantly reduces the difficulty of analysis and does not incur any loss of generality. Additionally, given the correspondence between each step of gradient descent and each layer of transformers, Proposition immlles the possibility of analogous behavior among layers of deep transformers during training. Such a result provides a promising approach to the optimization of deep transformers. Now we are ready to present our main results regarding the derivation the optimal parameters a and w gd
Theorem 3.3. For training distribution Dtr in Definition 2.1 suppose that the training context length n sufficiently large such that n 2 Mmmax (dd,, dd.)). Additionally, suppose that the perturbation of w around its expectation 3 is smaller than 4, i.e. (w",g*) > 0. Based on these assumptions,
è Herr slightly abuse the notation of Dt: denote both the distribution of Z, Wquery and I', w*.

---

F-10
Cptmolut
R(w ((w(L)), (L)), gd take the T.l AcCurao
0.75
20007
Ouarma
110013
Opondd
9.70
0.50
60304
9.05
0.05
2,001
100000
3.60
0,003
602044
2.00100
0.70
010044
0.05
6
Figure 2: Gradient descent experiments using single-layer neural network. We use grid search to obtain the optimal learning rate for different training context lengths n and different steps of gradient descent L. Then we use the corresponding optimal learning rate to perform multi-step gradient descent optimization on the test dataset. The results show that longer training context lengths and more gradient descent steps lead to smaller optimal learning rate and better optimization.
the optimal learning rate a and initialization (0) gd
i.e a, w = arg min (0)
gd
value as follows:
a
W gd
cB"
where c is an absolute constant.
Theorem 3.3 clearly identifies the optimal learning rate a and initialization w. (0)
gd Specifically,
shows that the optimal initialization w gd aligns the direction of the expectation G", with its length independent of the number steps L, and the context length n. Such a conclusion complies with our intuitions as the initialization w!d represents the memory of large language models, which is not dependent on the task-specific context examples. In contrast, the optimal learning rate a is inversely related both n and This suggests that in both cases: with more in-context examples; and (ii) with more layers, the output of pre-trained transformers will equal to that of a more fine- ine-ggained gradient descent process using a smaller learning rate. Generally, a small-step strategy ensures the convergence of the objective, highlighting the potential benefits of deeper architectures and training inputs with longer context.
(0)
3.3 Robustness against context hijacking
The previous two subsections illustrate that for any input with context examples, we can obtain the corresponding prediction for that input from the well-trained transformers by applying gradient descent with the optimal parameters we derived in Theorem 3.3 As we model Dt. the distribution of hijacking examples, to examine the robustness of I.-layer transformers against hijacking, we only need to check whether the linear model achieved by L-step gradient on (Xhes Vue) can still conduct successful classification on Xquery- Specifically, we consider the classification error of the parameter
vector - (L. gd
((w(L)) PD. (Vquery (w (L) gdl
gd Xquery) < <0),
where T = ((Xhe: Dhe), (x queryy Myquery)) ). and (L)
gd is obtained by implementing gradient descent
(L)
(0) Similar to the previous result, E(w (I.) is identical to ((TF) defined in (2.6). Based on these preliminaries, we are ready to present our results regarding the robustness against context hijacking. We first introduce the following lemma illustrating that when the context length of hijacking examples is small, we hardly observe the label flipping phenomenons the prediction from well-trained transformers.
on (Xhe, 3he) with L steps and the optimal a and w
Lemma 3.4. Assume that all assumptions in Theorem 3.3 still hold. Additionally, assume that the
length of hijacking examples N is small such that N < Ó 23/2 and follows any continuous distribution. Based on these assumptions, holds that
E(w (L)
E(w
gdd o(1)
(0)
7

---

Laye
Lepe
Layer
ugur
044
0.9
0.5
Layer
307
Lay11
0.5
0.
2.5
Contort Lengt?
5 CarTent Longth
CenceNt Le kg3
Contort Leog17
Figure 3: Linear transformers experiments with different depths and different training context lengths. By testing the trained linear transformers on the test set, we can find that as the number of interference samples increases. the model prediction accuracy becomes worse. However, deeper models have higher accuracy, indicating stronger robustness. As the training context length increases, the model robustness will also increase because the accuracy converges significantly more slowly.
n=30
n=40
1.000
1,000
0.915
0.975
0.950
20.950
0.925
0.925
Layer
0.000
Layer
Layor
30,000
Layer?
Layer
0.675
Layar
Layer
Layer
layer
Layer 3
0.850
Layer
25
30
75
100
125
150
25
20
75
303
125
150
Longth
9tth
Figure 4: Linear transformers experiments on training dataset. By testing trained linear transformers on the training set, the initial accuracy of the model is high and can be improved with the increase of context length, indicating that the model can use in-context learning fine-tune 3 w And deeper models have stronger optimization capabilities.
Lemma 3.4 demonstrates that when the context length of hijacking examples small, the classification error of the linear model obtained through gradient descent on these hijacking examples is very close that of the optimal initialization. The reasoning is straightforward: When N is relatively smaller compared to training context length and since the optimal learning rate a is on the order of the reciprocal n, the contributions from the hijacking examples become almost negligible in gradient descent iterations, allowing the model to remain close to its initialization. Consequently, we consider the case that N is comparable with n in the following theorem.
Theorem 3.5. Assume that assumptions Theorem 3.3 still hold. Additionally, assume that N ( /-/z n 2 (Nd.), and follows some uniform distribution. Based on these assumptions, holds that
E E(W E(W(L) (b) gd Sc (
Nd
1
nL
where C1. C2 are two positive scalar solely depending on the distribution of a and w*.
Based on a general assumption that a follows the uniform distribution, Theorem 3.5 formulates the upper bound of the classification error as a function of the training context length n, the number of hijacking examples N, and the number of layers L. Specifically, this upper bound contains a term
monotonically increasing function for N and a monotonically decreasing function for n and I,, Theorem 3.5 successfully demonstrates two facts: (i) well-trained transformers with deeper architectures,
proportional to -(1 θ Nd) )². As -1-6 θ AL )²
those pre-trained longer context examples, will exhibit more robustness against context hijacking; and (ii) for given well- trained transformer, the context hijacking phenomenon is easier to observe when provided with more hijacking examples. These conclusions align well with our experimental observations (Figure 3).
Experiments
4.1
Optimal gradient descent with different steps
In our theoretical framework, the optimal gradient descent with more steps (I.) or longer training context length (r2) will have smaller learning rate per step (Theorem 3.3 and this combination

---

Cit
926
Tapk: Courtry
Sport
9224
3700
700
70%
7094
-254
622
5076
6071
5075
8 4.50%
10%
2214
0
1O5
ON6
0.4 Lovel
Lovd
Lovd
Lovel
Figure 5: Context hijacking in real-world LLMs. We investigate the label flipping rates of Qwen2.5- Base and Owen2.5-Innsruce, models of varying sizes, facing varying levels of context hijacking. Hijacking levels ranging from I to 8 represent context hijacking prefixes ranging in length from 10 to 80 sentences. Label flipping occurs when the model predicts the next token to be the topic described by the context prefix, which differs from the correct answer to the question. Our experimental results on real-world LLMs are still completely consistent with our theoretical conclusions: deeper models more robust to context hijacking.
of more steps with small learning rates will perform better on the optimization process over context samples (Theorem (.5). Our theory shows that trained transformer will learn the optimal multi-step gradient descent, which will make more robust during testing. Therefore, we directly verified the consistency between practice and theory in the multi-step gradient descent experiment.
We construct a single-layer neural network to conduct optimal multi-step gradient descent experiments. Each training sample (x,,)) is drawn i.i.d. from the distribution Drr defined in Section 2.1 We consider the learning rate that minimizes the loss of the test sample when the single-layer neural network is trained using to 8 steps of gradient descent, that is, the optimal learning rate ag corresponding to L-step gradient descent, which can be obtained by grid search. Figure Sshows that az decreases as L and increases, which is aligned with our theoretical results (Theorem 3.3).
Next, we discuss the second part of the theoretical framework, i.e., gradient descent with more steps and small step size performs more fine-graines optimization (Theorem 3.5). We apply the optimal learning rate searched in the training phase to the test phase, and perform gradient descent optimization on the test samples drawn from Dro with the optimal learning rate and its corresponding number of steps. We can find that with the increase in the number of gradient descent steps and the decrease in the learning rate, the performance of the model will be significantly improved.
4.2 Robustness of linear transformers with different number of layers
Applying our theoretical framework to the context hijacking task on transformers can explain it well, indicating that theory has practical significance. We train linear transformers with different depths and context lengths on the training dataset (Dtz). We mainly investigate the impact of training context length n, and model depth L and the testing context length N on model classification accuracy.
We first test the trained transformers on the training dataset to verify that the model can fine-tune the memorized 3 to w* According to the Figure 4 we can find that the model has a high classification accuracy when there are very few samples at the beginning. This means that the model successfully memorizes the shared signal f As the context length increases, the accuracy of the model gradually increases and converges, meaning that the model can fine-tune the pre-trained 3 by using the context samples. In addition, deeper models can converge to larger values faster, corresponding to the theoretical view that deeper models can perform more sophisticated optimization.
Then we conduct experiments on the test set. Observing the experiment results (Figure 3), we can see that as the context length increases, the accuracy of the model decreases significantly and converges to 50%, showing that the model is randomly classifying the final query Xquery. This is consistent with the context hijacking phenomenon that the model's robustness will deteriorate as the number of interference prompt words increases. When the number of layers increases, the models with different depths show the same trend as the context length increases, but the accuracy of the model will increase significantly, which is consistent with the phenomenon that deeper models show stronger robustness in practical applications. In addition, the model becomes significantly more robust as the training context length increases, because the accuracy converges more slowly the length increases.

---

4.3 Context hijacking in real-world LLMs
To further validate our theoretical results in more realistic scenarios, we investigate the label flipping models of varying sizes, facing different levels of
rates of Qwen2.5-Baal and Owen2.5-Instrun hijacking. The 0.5B, 1.5B, and 3B models have 24, 28, and 36 layers, respectively. We first construct four different datasets. Each question in the dataset consists of two parts: factually correct prefix of a certain length and a fact retrieval question. Each sentence of the context prefix will describe a topic from a different perspective and with different words. Therefore, rather than using repeated sentences, we use diverse sentences to describe a fact, which is more realistic. We divide the context hijacking into eight different levels according to the length of the context prefix, from level to level 8, which means the context has 10 to 80 sentences. We then observe whether the model predicts the correct next token. If the next token changes to the topic described by the context prefix (different from the answer to the fact retrieval question), we call this phenomenon "label flipping". For more detailed experimental details and examples, please refer to Appendix H..3
Figure shows how the label flipping rate of different models changes with the hijacking level. Although models exhibit varying resistance to context hijacking across different tasks, we can find that in practical L.MMs, longer hijacking context will significantly increase the label flipping rate (leading to lower accuracy), while increasing the model depth can well alleviate this problem. The experiment results are consistent with our theoretical conclusions, indicating that our theoretical results can be generalized to deeper and larger LLMs in practice. Additionally, we find that instruction fine-tuning can improve the model's robustness to context hijacking in most cases, but the effect is not significant, which provides new insights for future work, such as adversarial optimization.
5 Conclusion and limitations
In this paper, we explore the robustness of transformers from the perspective of context hijacking 337 We build a solid theoretical framework by modeling context hijacking phenomenon as a linear classi- fication problem. We first demonstrate the context hijacking phenomenon by conducting experiments on LLMs with different depths, i.e., the output of the LLMs can be simply manipulated by modifying the context with factually correct information. This reflects intuition: deeper models may be more robust. Then we develop a comprehensive theoretical analysis of the robustness of transformer, showing that the well-trained transformers achieve the optimal gradient descent strategy. More specifically, we show that as the number of model layers or the length of training context increase, the model will be able to perform more fine-grained optimization steps over context samples, which can be less affected by the hijacking examples, leading to stronger robustness. Specifically considering the context hijacking task, our theory can fully explain the various phenomena, supported by series of numerical experiments. We also conduct nonlinear experiments to extend our theoretical results.
The limitation of our work is that our model is specifically designed for context hijacking, a recent phenomenon from the real world, We only focus on the hijacking task, so new constructions may be required for other problems, and different results may be derived based our technical framework. Besides, our current analysis follows from the recent works on modeling transformer model as a gradient descent optimizer, it is also possible that the corresponding meta optimization algorithms more complicated ones. This will also lead to different theoretical results based on our techniques.
Our work provides new perspective for the robustness explanation of transformers and the under- standing of in-context learning ability, which offer new insights to understand the benefit of deeper architecture. Besides, our analysis on the optimal multi-step gradient descent may also be leveraged to other problems that involve the numerical optimization for linear problems.
Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers and area chairs for their helpful comments. This work is supported by NSFC 62306252, Hong Kong ECS award 27309624, Guangdong NSF MAAA1115512244,4 and the central fund from HKU IDS.
10

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidr, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arX1v:2303.00877, 2023.
[2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers leam imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614-45650, 2023.
[3] Yaroslav Aksenov, Nikita Balagansky, Sofia Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gorbatovski, and Daniil Gavrilov. Linear transformers with learnable kernel functions are better in-context models. In Proceedings of the 62nd Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 9584 199479997.. 2024.
[4] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022.
[5] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022.
[6] Cem Anil, Esin Durmus, Nina Rimsky, Mrinank Sharma, Joe Benton, Sandipan Kundu, Joshua Batson, Meg Tong, Jesse Mu, Daniel Ford, et al. Many-shot jailbreaking. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
[7] Usman Anwar, Johannes Von Oswald, Louis Kirsch, David Krueger, and Spencer Frei. Adver- sarial robustness of in-context learning in transformers for linear regression. arXiv preprint arXiv:2411.05189, 2024.
[8] Giovanni Apruzzese, Hyrum S Anderson, Savino Dambra, David Freeman, Fabio Pierazzi, and Kevin Roundy. "real attackers don't 't compute gradients": bridging the gap between adversarial ml research and practice. In 2023 IEEE Conference on Secure and Trustworthy Machine Learning (Sa7'LL, pages 339-364. IEEE, 2023.
[9] 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.
[10] Luke Bailey, Euan Ong, Starrt Russell, and Scott Emmons. Image hijacks: Adversarial images can control generative models runtime. In Forty-first International Conference on Machine Learning, 2023.
[11] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117748):30033-3077, 2020.
[12] Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning transformers and Ilms by learning learn discrete functions. In The Twelfth International Conference on Learning Representations, 2023.
[13] Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons in language models. URL https:/opeenippblic blob. core. windows. net/neuron- explainer/paper/ineex html. accessed: 14.05. 2023), 2023.
[14] Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
[15] Patrick Chao, Alexander Robey, Edgar Dobriban, Hamed Hassani, George Pappas, and Eric Wong. Jailbreaking black box large language models in twenty queries. arXiv preprint arXiv:2310.08419, 2023.

---

[16] Xingwu Chen, Lei Zhao, and Difan Zou. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. arXiv preprint 241V:2408.04552., 2024.
[17] Xingwu Chen and Difan Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404. 00000.00..1..... 2024.
[18] Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via lan- guage model in-context tuning. In 60th Annual Meeting of the Association for Computational Linguistics, ACL 2022, pages 719-730. Association for Computational Linguistics (ACL), 2022.
[19] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. In Forty-first International Conference on Machine Learning, 2023.
[20] Yixin Cheng, Markos Georgopoulos, Volkan Cevher, and Grigorios G Chrysos. Leverag- ing the context through multi-round interactions for jailbreaking attacks. arXiv preprint urXiv:2402.09177, 2024.
[21] Arijit Ghosh Chowdhury, Md Mofijul Islam, Vaibhav Kumar, Faysal Hossain Shezan, Vinija Jain, and Aman Chadha. Breaking down the defenses: A comparative survey of attacks on large language models. arXiv preprint arXiv:2403.044700 2024.
[22] Antonia Creswell, Murray Shanahan, and Irina Higgins. selection-inferent Exploiting large language models for interpretable logical reasoning. In The Eleventh International Conference on Learning Representations, 2022.
[23] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, pages 4005-4019, 2023. 2023.
[24] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. In International Conference on Leurning Representations, 2019.
[25] Gelei Deng, Yi Liu, Yuekang Li, Kailong Wang, Ying Zhang, Zefeng Li, Haoyu Wang, Tianwei Zhang, and Yang Liu. Jailbreaker: Automated jailbreak across multiple large language model chatbots. arXiv preprint arXiv:2307.0877 2023.
[26] Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint / 18/0.04805, 2018
[27] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting contextext. arXiv preprint arXiv:2410.01774. 2024.
[28] Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. Advances in Neural Information Processing Systems, 36, 2024.
[29] Jingwen Fu, Tao Yang, Yuwang Wang, Yan Lu, and Nanning Zheng. How does representation impact in-context learning: A exploration synthetic task. arXiv preprint 45555 2023.
[30] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
[31] Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. In The Twelfth International Conference on Learning Representations, 2023.
[32] Pengfei He, Han Xu, Yue Xing, Hui Liu, Makoto Yamada, and Jiliang Tang. Data poisoning for in-context learning, arXiv preprint ar arXiv:2402.02166. 2024.
[33] Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gradient descent with chain of thought. In The Thirteenth International Conference on Learning Representations, 2025.
12

---

[34] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Forty-first International Conference on Machine Learning, 2023.
[35] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37883 2022.
[36] Joonhyun Jeong. Hijacking context in large multi-modal models. arXiv:2312.07553, 2023.
arXiv preprint
[37] Yibo Jiang. Goutham Rajendran, Pradeep Kumar Ravikumar, and Bryon Aragam. Do llms dream of elephants (when told not to)? latent concept association and associative memory in transformers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[38] Aounon Kumar, Chirag Agarwal, Suraj Srinivas, Aaron Jiaxun Li, Soheil Feizi, and Himabindu Lakkaraju. Certifying Ilm safety against adversarial prompting. arXiv preprint DXXix:309.002055 2023.
[39] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
[40] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 19565-19594. PMLR, 2023a.
[41] Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis in-context linear estimation: Data, architecture, and beyond. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[42] Licong Lin, Yu Bai, and Song Mei. Transformers decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations, 2023.
[43] David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. Tracr: Compiled transformers as a laboratory for interpretability. Advances in Neural Information Processing Systems, 36, 2024.
[44] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang. Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9): 2023.
[45] Yang Liu, Yuanshun Yao, Sean-Francoi, Ton, Xiaoying Zhang, Ruocheng Guo Hao Cheng, Yegor Klochkov, Muhammad Faaiz Taufiq. and Hang Li. Trustworthy 1lms: A survey and guideline for evaluating large language models' alignment. arXiv preprint 438.522308.05374 2023.
[46] Yi Liu, Gelei Deng, Yuckang Li, Kailong Wang. Zihao Wang. Xiaofeng Wang, Tianwei Zhang, Yepang Liu, Haoyu Wang, Yan Zheng, et al. Prompt injection attack against lIm-integrated applications. arXiv preprint arXiv:2306.05499 2023.
[47] Arvind V Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, 2023.
[48] Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 15:17359-17372 2022.
[49] Sewon Min, Mike Lewis, Luke Zeutleooeer and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, pages 2791-2809, 2022.
13

---

[50] Zhenxing Niu, Haodong Ren, Xinbo Gao, Gang Hua, and Rong Jin. Jailbreaking attack against multimodal large language model. arXiv preprint arXiv:2402.02339, 2024.
[51] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[52] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 15:27730-27744, 2022.
[53] Lalchand Pandia and Allyson Ettinger. Sorting through the noise: Testing robustmess of information processing in pre-trained language models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 1583-1596, 2021.
[54] Onkar Pandit and Yufang Hou. Probing for bridging inference in transformer language models. In NAACL 2021-Armuu Conference of the North American Chupter of the Association for Computational Linguistics, 2021.
[55] Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2023.
[56] Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. In The Twelfth International Conference on Learning Representa- tions, 2023.
[57] Fábio Perez and lan Ribeiro. Ignore previous prompt: Attack techniques for language models. In NeurlPS ML Safety Workshop, 2022.
[58] Jorge Pérez, Pablo Barceló, and Javier Marinkovie. Attention is turng--ommlete. Journal of Machine Learning Research, 22(75):1-35, 2021.
[59] Yao Qiang, Xiangyu Zhou, and Dongxiao Zhu. Hijacking large language models via adversarial in-context learning, arXiv preprint arXiv:2311.09944 2023.
[60] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188)99
[61] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
[62] Xinyue Shen, Zeyuan Chen, Michael Backes, Yun Shen, and Yang Zhang. do anything now": Characterizing and evaluating in-the-wild jailbreak prompts on large language models. In Proceedings of the 2024 on ACM SIGSAC Conference on Computer and Communications Security, pages 1671-1685, 2024.
[63] Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed H Chi, Nathanael Schärli, and Denny Zhou. Large language models can be casily distracted by irrelevant context. In International Conference on Machine Learning, pages 31210-31227. PMLR, 2023.
[64] Chen Siyu, Sheen Heejune, Wang Tianhao, and Yang Zhuoran. Training dynamics of multi head softmax attention for in-context learning: Emergence, convergence, and optimality. In The Thirty Seventh Annual Conference on Learning Theory, pages 4573-4573 PMLR, 2024.
[65] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems, 36:71911-719477 2023.
[66] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
14

---

[67] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[68] Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 63-76, 2019.
[69] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[70] Jindong Wang, HU Xixu, Wenxin Hou, Hao Chen, Runkai Zheng, Yidong Wang. Linyi Yang, Wei Ye, Haojun Huang, Xiubo Geng, et al. On the robustness of chatgpt: An adversarial and out-of-distribution perspective. In ICLR 2023 Workshop on Trustworthy and Reliable Scage-Scaae Machine Learning Models, 2023.
[71] Jiongxiao Wang, Zichen Liu, Keun Hee Park, Zhuojun Jiang, Zhaoheng Zheng, Zhuofeng Wu, Muhao Chen, and Chaowei Xiao. Adversarial demonstration attacks on large language models. arXiv preprint arXi::2305. 14950, 2023.
[72] Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference on Machine Learning. 2024.
[73] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35: 1001-11083,,000
[74] Zeming Wei, Yifei Wang, Ang Li, Yichuan Mo, and Yisen Wang. Jailbreak and guard aligned language models with only few in-context demonstrations. arXiv preprint arXiv:2310.06387, 2023.
[75] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11090. PMLR, 2021.
[76] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference on Learning Representations, 2023.
[77] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2021.
[78] Xilie Xu, Keyi Kong, Ning Liu, Lizhen Cui, Di Wang, Jingfeng Zhang, and Mohan Kankanhalli. An Ilm can fool itself: A prompt-based adversarial attack. In The Twelfth International Conference on Learning Representations, 2023.
[79] Yifan Yao, Jinhao Duan, Kaidi Xu, Yuanfang Cai, Zhibo Sun, and Yue Zhang. A survey on large language model (Ilm) security and privacy: The good, the bad, and the ugly. Cigf-CConiiencce Computing, page 100211, 2024.
[80) Ori Yoran, Tomer Wolfson, Ori Ram, and Jonathan Berant. Making aanmenel- language models robust irrelevant context. In The Twelfth International Conference on Learning Representations, 2023.
[81] Jiahao Yu, Xingwei Lin, Zheng Yu, and Xinyu Xing. Gptfuzzer: Red teaming large language models with auto-generated jailbreak prompts. arXiv preprint arXiv:2309. 00022309..022.... 2023.
[82] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequee functions? International Conference Learning Representations, 2020.
15

---

[83] Chenyang Zhang, Xuran Meng, and Yuan Cao. Transformer learns optimal variable selec- tion in group-sparse classification. The Thirteenth International Conference on Learning Representations, 2025.
[84] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
[85] Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of linear transformer block: benefits of the mip component and one-step gd initialization. arXiv preprint arXiv:2402.14951, 2024.
[86] Yunqing Zhao, Tianyu Pang, Chao Du, Xiao Yang, Chongxuan Li, Ngai-Man Man Cheung, and Min Lin. On evaluating adversarial robustness of large vision-language models. Advances in Neural Information Processing Systems, 36, 2024.
[87] Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference on Learning Representations, 2023.
[88] Kaijie Zhu, Jindong Wang, Jiaheng Zhou, Zichen Wang, Hao Chen, Yidong Wang, Linyi Yang, Wei Ye, Yue Zhang, Neil Zhenqiang Gong, et al. Promptbench: Towards evaluating the robustness of large language models on adversarial prompts. arXiv e-prints, pages arXiv-2306, 2023.
[89] Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15043 2023.
16

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
17

---

Justification:
The paper discuss the limitations of the work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof'
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper provides the full set of assumptions and a complete (and correct) proof for each theoretical result in appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide experimental settings and training details in appendix.
Guidelines:
18

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
The code is being organized and we will provide it later.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
wtweeesubbmssiiooocooccc for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https c/iips.cc/public//uu
CodeeSumiiii wwttubbmissiiioi ssionoolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
19

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide experimental settings and training details in appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: We provide error bars and some other appropriate information about the statistical significance for most experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide experimental settings and training details in appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
20

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses such risks.
Guidelines:
21

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators and original owners of assets (e.g., code, data, models), used in the paper, are properly credited and the license and terms of use explicitly are mentioned and properly respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
22

---

Answer:
[NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non stanndaddrrd components.
Please refer to our LLM policy (https://neut) /cC//aareeecccc/2222//////////// for what should or should not be described.
23

---

A Related works
In-context learning via transformers. The powerful performance of transformers is generally believed to come from its in-context learning ability [14
18
49
44
77
A line of recent works
study the phenomenon of in-context learning from both theoretical 9
64 41 and 39 perspectives diverse settings. 14 first showed that
16 27
empirical [30 4 40,
61 56 55 12
29
GPT-3 can perform in co-totettt learning. [6 studied the role of different heads within transformers
in performing in-context learning focusing the sparse linear regression setting. [27] studied the ability of one-layer linear transformers to perform in-context learning for linear classification tasks. 3 demonstrated linear transformers perform less well on in-context tasks, but the linear transformer here refers to a model that uses a kernel function to approximate the standard attention computation for computational efficiency, albeit the expense of some in-context learning performance. In contrast, the linear transformer in our paper is a model that removes the activation function from the standard transformer. As discussed in the main paper [.2] this is a very common setting in transformer theory research. In addition, while not considering the ICL setting, [83 studied the
training of one-layer transformers on the group-sparse linear classification setting.
Mechanism interpretability of transformers. Among the various theoretical interpretations of 24
transformers [28
17 [3 one of the most widely studied theories is the ability of transformers to implement optimization algorithms such as gradient descent [69 theoretically and empirically proved that transformers
82
43
54
58
33
73 75
87
2. 185 76
learn in-context by implementing a single step of gradient descent per layer. 2 theoretically analyzed that transformers can learn to implement preconditioned gradient descent for in-context learning. [85] considered ICL in the setting of linear regression with a non-zero mean Gaussian prior, more general and common scenario where different tasks share signal, which is highly relevant
19
5 23
84
69
work.
Robustness of transformers. The security issues of large language models have always attracted a
great deal of attention 779 46 57
89 g However. most of the research focuses on jail-breaking
black-box models [21], such as context- context-based adversarial attacks [8 74 78
7
88
20 70). There
is very little white-box interpretation work of attacks the transformer, the foundation model of
32 7 37]. (59] first considered attacking large language models during in-context learning, but they did not study the role of transformers in robustness. 37 proposed the phenomenon of context hijacking, which became the key motivation of our work. They analyzed this problem from the perspective of associative memory models instead of the in-context learning ability of transformers.
LLMS [59 10
B Notations
Given two sequences { } and }, we denote vn = (yy) if there exist some absolute constant
Gi > 0 and N >> 0 such that [xr, S C1 [ynl for all n > N. Similarly, we denote Vy, (y,.) ) if there exist C2 > 0 and N > 0 such that [æ,|| 2 C2ly%1 for all n > N. We say and Tn (2))) both holds. Additionally, we denote x,, o(ya) if, for any > O. there exists some (ee) > 0 such that |z,.| S elyn| for all n 2 (e),, and we denote In = w(yn) if yn o(x,). We use Ó(-), $2(.), and 6(.) to hide logarithmic factors in these notations respectively. Finally, for any n e N., we use [n] to denote the set ,,,
+(yn) if I,, O
C Proof of Proposition 3.1
In this section we provide a proof for Proposition 3.1
Proof of Proposition 3.7 Our proof is inspired by Lemma in 22, while we consider a non-zero initialization. We first provide the parameters We, Pt, Qe E (d++))((++1) of L-layers transformer.
In (0)
Oaaxd 0 Q
1
-
0 where Ie E Rddd
W
P
1
For the linear classification problem, the input sample Zo E R(d+1)x(x+)) consists of {(x,, v:)} 1n and (Xquery, Vquery) in (2.1) which will first be embedded by W g. Let ( [dx(n+1) denote
24

---

the first d rows of Wa W (Zo) and let Yo) E R Ix(++)) denote the (d + 1) 1)-th row of W (Z(Zo). In subsequent iterative updates in (2.3). the values at the same position will be denoted as x and Y for ( = I,..., L. Similarly, define X
(1) E IRdx" and Y E I as matrices that exclude the last query sample (Xquery, Viuuer). That is, they only contain the first columns of the output of the 1-th layer. Let x and y. be the i-th pair of samples output by the I-th layer. Define a function
(1) (2)
g(x,y. v. : Rd x R X Z I: let Xquery x and Viueery y (wed)) x), then g(x,, y, Vquery-
(0)
Next, based on the update formula (2.3) and the parameters constructed above, we have:
x(+11= x
X (0)
y(++)= y - m(( TT)
Then for all i E (I,........
n
y
1+1)
yi
(1)
Σ
T x;
So vi (1+1) does not depend on Yquery. For query position,
((+1)
Σ
x querry 12x,1} (1)
Equery yqueer
Then we obtain g(x, y. and 0,
g(x,v. 1) = squerry
x query Fi الزز 11-1)
j=1
yquezy
query r 11-2813] 2)
query r 1 11
n
n query T (11)
"qquery
x query To Tox;,0
(0)
(1-1)
j=1
/=1 (0)
= y (w (0)
gel Xqueryy
Σ
x query Tox
query T 11-1X,11 (-1)
3=1
g(x, 0, k)- = Viquery
X queer T1-18,y1 11-1)
= (
query T1-2X
(1-2)
query T T 1x,1, (1-1)
12
Vquuery
X Xuuery Tox
(0)
x query 1 (1-1)
1=1
72
= =-(w() (0) gd Xquery)
x query X Fox,u,
(0)
X Ii-IX ;-
j=1
So have 9(x, y, = g(x, 0, 1) + y. Observing g(x, 0,1), can find that it is linear in x for the reason that every term of xx 0, 7) is linear in Xquery, which means we can rewrite it. We verify that there exists 02 E IRd for each / E [4]. such that for all x,
g(x, y. () g(x, 0, k) + y = (0i, x) +
Let / = O, we have (0,, x) = g(x y, 0) y = Vquery - = - (( (0)
Xquery), ) so 0, = -w) (0)
Next,
gd
we will show that for all (x,, vi) E (x1188)) ... (Xn, vn), (Xquery, Vquery) }
g(x,, v1, yi (0), Vb.
25

---

(++1) Observing the update formulas for Di ((+1) and Jquery: if we let Xquery := x, for some i, we can get that V/ (I+1) ( because yi Yquery by definition. This indicates that
Yo = + X.
Finally, we rewrite the update formula for Vi
(n+1)
squery =
X query I !x,,y (1)
Vquery
(TX(((()) (T) Xquery)
Equery
(01+1, Xquery) = (0), Xquery) - XX O + (Y(()) Xquery
Since Xquery is an arbitrary variable, we get the general update formula for O::
6 = 0 (TiX X 0, + (Y(0))
Notice that we use the mean squared error, we have
1
I(w)= =
(w, x1) y12
2
1=1
1
X
(88(0)) 112
2
Then we get its gradient vi(w) X (X (Y(0))T). Let w!ll = 0i, we have
0141 ((+) gd
0: (T X + +()))
x X Tww gd - (P(()))))
w
gd
10) gd
gd
And the output of the 1-th layer Jquery is (querr: Vquery: I) = Vquery + (0), Xquery) =
(1)
Vquery
(w gdl Xquery).
In our settings, we have y/ n+1 - (www (!) (1)
gd Xquery) because the input query label
D Gradient descent updates of parameters
In this section, we provide further details regarding the updating of parameters w ggl which will be utilized in subsequent proof. Besides, it can directly imply Proposition 3.2 Before demonstrating the mathematical, we first introduce several utility notations, which will be used in subsequent technical derivations and proofs. We denote SJ,, as the set of all k-dimensional tuples whose entries are drawn from (0, 1, 1,...,7 I 1} without replacement, i.e.
-jx.- Then given the set of all historical learning rates before or at (-th iteration, i.c. {ao, a...... a4-1) and Si,e defined above, we define A as
Sr,= {(r, f2 ... jk) /1 12, . (0.1,...,1-1
A
Σ
II
aj.
(E k=1
Then we can observe that the permutation of elements of (on, O1,..., 01-1) would not change the value of A Then based on these notations, we present mathematical derivation in the following.
By some basic gradient calculations, we can re-write the iterative rule of gradient descent 3.1 as
(l+1) ( Rd aq ((W (U)
gd
26

---

wwll
((w (1)
gd) x.) - siin((w" xi)))
gd
00
1=1
Ia 00 Σ xxx:
w(( + oi Σ sign((w" x;)) x;
(D.1)
gd
Based on this detailed iterative formula, and the definition of St, and A above, we present and prove the following lemma, which characterizes the closed-form expression for w(!)
(1) Rd
Lemma D.I. For the iterates of gradient descent, i.,
with / 10,,,...,,L 1}. , it holds that
la
Ar. (- Σ xxxx.
w (0)
Σ k=1
AI,
(-E - Exxx] R-1
gd
k=1
i=1
i=1
sign( (w",,;)) xxi)) xi
(D.2)
Proof of Lemma Before we demonstrate our proof, we first present some conclusions regarding St,k and A1,. By directly applying the Binomial theorem and the definition of Ar,s, we can obtain that
II A=
L-
(I
La
+ onk
(- --Exx
(-Exx7)
Ia
ALA - ExxxT)
(D.3)
1=1
k=1
1=1
Additionally, by utilizing the definition of St,k, we can easily derive that
S =((,, ... fi)] ................................ ja E {0, I,..,,............. + (0;
jk
-(Uin.2..... jo) ... ia E 10. 1 f.. * jz} U J. b E 10,
jk-1
13is1 # j2 .... jx-1}
= U {(1.12, jk (1 je. E S1
holds when & S l. This result can further imply that
& II
A141, =
Σ
aj
Σ
ll ajn
Σ
II ag.
(a.Ja..dh)ES 1
Ui 1) kml
Ai, + on A1
(D.4)
holds when & S l. Additionally, it is straightforward that
A++11411 Ar,
A A +
(D.5)
With these conclusions in hands, we will begin proving this lemma by induction. When = 1, by the iterative rule D.II) we can obtain that
I1 on Σ xxxx
+ om Σ sign((w", x xx
W gd
w gd
which follows the conclusion of (D.2) due to (D.3) and the definition of A1,1- By induction, we assume that (D2) holds at 1-th iteration. Then at (1 1)-th iteration, we can obtain that,
(1+1) gd
Ia an Σ (xxx))
w!! (1) gd
+ a
sign( (w", xx))
=0:
sign ((w,,,)))
xx)) x3/
I aj
Σ xxxx
n
Σ A (- n
Σ xxxx R-1
I a,
Σ
xxxp
w
(0)
gd
k=1
27

---

(
sign (w", xx1) xi)
=aq
sign (w" .x,)) xci
+
II
Ia - akk
Σ xxxx ))
w (0)
k=0
i=1
AI,k (-
(- - xix7 (-Exxx) ))
xiX
+
An
1=1
i=1
1=1
i=1
n
sign( (een((w'> (w",, ,x,)) .Xi
Xi
n Ia ar (
xxxx]
(0) + (Ar,1 + az)
Σ sign (w" (x1)) X,
k=0
1=1
Σ
1
(A1,e + on 1)
(-E
Σ xxx7
k-1
Σ sign( (w" ;,w",x;]] xxi) xs)
A=2
i=1
Σ xix]
Σ sign((w" xxi) xi
la
akk
x (xxxx²)
w (0)
gd
+ A141111
Σ
sign ( (w", xii) xi
k=0
i=1
i=1
A141.b
-
>-xxa xxx,
Σ
sign
Xi
2
1=1
A141/41
- Σ xxx]] k
Σ sign( (w" xxi) Xs
i=1
La +
[+1
AL+1,8 (-Exxa xix
w w(0) +
Σ
1+1
ALti,
(
(-Exx²
x
K-1
k=1
i=1
k=1
1=1
Σ sign( (w",, xxx)) Xi
(1) The second equality holds by substituting w Wgd with its expansion from (D. assuming it is valid at the [-th iteration by induction. The third and fourth equalities are established by rearranging the terms. The penultimate equality is derived by applying the conclusions regarding Ar,x from (D.4)
and (D.5). The final equality is obtained by applying (D.3). This demonstrates that (D.2) still holds at [ 1-th iteration given it holds at (-th iteration, which finishes the proof of induction.
(L) Lemma D.I demonstrate that the learning rates ay's will only influence w gd by determining the value of Ar.k s. While as we have discussed above, the values of AL, 's depend solely on the elements in the {ao. Consequently, the permutation of (oo...., 01.-1} will also not affect the value of w gd thereby confirming that Proposition 3.2 holds.
.......... faa..... )1 }, and remain unchanged when the order of these learning rates is rearranged.
(L)
E Proof of Theorem 3.3
In this section, we provide a detailed proof for Theorem 3.3 We begin by introducing and proving a lemma that demonstrates how w gd must align with the direction of B This alignment constrains (0) the choice of w gd to scalar multiple of f specifically in the form of co B². Additionally, in the
3.3
(0)
subsequent sections, we will use the notation Σ Σn Xxx_T
28

---

Lemma E.I. Under the same conditions with Theorem 3.3 to minimize the loss R(wla)) ) w (0) always in the form of co 3
gd is
Proof of Lemma Utilizing the independence among the examples in T {xx,,1)..... :......... Xx, Un), (Squery (quory) J. and w", we can expand Rw by law of total expectation as
R R(w ) =BT [(( (L)
gd Xquerr) ) sign( (oguu(ww. Xquery) ))
gd
=Br,w w [(w (7)
xquarry) sign(( (w", Xquuery )) (w (T)
gd Xquery?] +
=By -1((x119)7
Neuarr) 2 sign({ (w", Xqueery )) )
Th [x
[( (L)
gd Xqueryy? )²
(wi()) gd
Xqquery [(x,, yo) Mi w'] +
Iw!a gd
=B (Ix, w
V π
w
+
2
π
where the last equality holds since Exqqurry ((w, Xquery? )
w Ex Xquory x
query. w
||w l2 when w is independent with Xquery, and Exquery [(W1, Xquery) sign( (W2, Xquery)) )
(1) implied by Lemma G.1 Therefore in the next we attempt to optimize the first 2 (w", gd
term E [(xx,00)
w (L)
gd
w
By applying the closed form of w (L)
gd
Lemma D.1
with all ar a, we have
L-1 [L 05)'
gd
[Ia
2
gd
+
( sign((w> (w", ((", sig
1=0
Based on this, we can further derive that
(( 12
(L) gd
2
w
= (w:")
(0) gdl
T 1 ((x...2))) [Il az)) 1 (0)
π
2
4-1
20 (w
gd
(0))
T "((X..))
Σ (Ia 0151 1+
1=0
Σ sign( (w",,:))) xx))xx +
=1
w
gd
(0)
2
2
202
(w (0)
gd
,8)) +
w (0)
2 3r +
2
(E.I)
gd
where C1, C22 C are some scalar independent of w w0th The second inequality holds since
((xx,.))) w [LL a5)²² 2 Is for some scalar er guaranteed by Lemma G.2 and
((xx,)))) Σl-1 (L. 2 1+ (E xign((w*,,,))x,)
cgB² for some scalar O2.
guaranteed by Lemma G.3 As the result of (D33) is a quartic function of (0) conclude that it achieves the minimum value when w (0) the proof.
gd we can easily grl cof** for some scalar co, which completes
Based on Lemma E. in the following proof, we will directly replace w gd with on and attempt to find the optimal co. Now we are ready to prove the following theorem, a representation of Theorem 3.3
Theorem E.2 (Restate of Theorem 3.3 For training distribution Drr in Definition 2.1 suppose that the training context length is sufficiently large such that 2 [Mmmxx(d dd)). Additionally,
29

---

suppose that the perturbation of w" around its expectation 3 is smaller than 2. i.c. (w",8") > 0.
then for any learning rate a and initialization ( gd it holds that
R( ( gd
S @(1 cm)) gd ci w [B) + ((adL) + +C,
where both c1. C are absolute constants. Additionally, by taking w c1 and ((t)) the upper bound above achieve its optimal rates as
R( R(w ( Rd
θ - +
Proof of Theorem E.2 Utilizing the fact that Ls (La -a)) a It-r (Ia aE)'s and (0)
co.3r we can we-write the close form of w (L)
gd
w (L) gd
la (1a-
V + -aa-a))) c0.3"
1-1
(V2Sw*
"1
Σ
a
Σ
sign( (wnn(ww> xi) xi))x;
1=0
Then by the similar calculation to E.1
Lemma
we have
R( (R(w!l (4) w (L)
gd
=E
gd
V
2
w
+
L-1
[L a5²²² :.
(co3*
1
2
w'
Σ
[La
π
1=0
(V 2 Ew " - sign( (ww', xx) xx)))x])
27
+
[Ia - + 4. (0083 - 2 ---w" w² 1
-v=
<2
+ 2
a²
Σ (Le (ZD'
2 (
Ew-
Σ sign((w (w" xx))xxx
127
+C
π
where the last inequality hols by (a + b)2 < 2a + 2b2. and C is an absolute constant. Therefore, in the following, we discuss the upper-bounds for and II respectively. For I, we have
I | -12 a
22 11
| (<<1)--2~>
E co.3"
2
0(11 ((1naa))2))
{((csP (co0*
2 π
where the first inequality is derived by the independence among x; and w and the submultiplicatiiity of la norm, and the second inequality holds by the concentration results regarding ||El|| provided in Lemma G.4 For IT, we derive that
L-1
sign( (w" x1) )x:) 1
11 US 2
Σ
La
Ew
41.9=0
11.1
11.2
where the inequality is guaranteed by the submultiplicatiiity of l2 norm. Then we discuss I1.1 and TT. 2.2 respectively. For IT.1, we have
aL
50 o aL
aoL log L
11.1 <
Σ 1,
Σ
30

---

The first inequality holds by the fact that 2)
S
k+1 for z E [0, 1]. The second inequality holds
by replace T++ with its upper bound -
The third inequality holds by E4-1 I, S log L
and E|| = (m) demonstrated in Lemma G.4 For II.2, we have
2 Σ V
2 ww) n
I1.2 IE
nIa)w*
sign( (w" xi) x; 12
π
π
Σ sign( (w" (x))xx; BV w | Ö(dd).
E
nIall + 2E
The first equality adds and minuses the same term. The first inequality holds by the submultiplicctiiity of f2 norm, and the fact (a+ 2 2222. The second inequality holds || -I-nlall2 O(v
proved in Lemma G.4 and Σi= sign( (w xx) )xx -n 2 S (V(nd), proved Lemma G.5
Combining all the preceding results, we can obtain that
V w >130 + (odL)) + C.
R(w (L)) S ((1 -am)ak) E
straightforward that when taking co
x² the expectation term will achieve its minimum,
which is the variance of w multiplying by a factor
V This finishes the proof that the optimal
initialization takes the value as wed = / We re-plug this result into the upper-bound above and utilize the fact that the variance is at the constant order. Then to find the optimal learning rate is actually to optimize the summation of (1 an)24 and adL. We can note that the first term will increases. Therefore, minimizing the summation of these two terms is essentially equivalent to finding an optimal a such that both
decrease a increases, while the second term will increase
terms of the same order. Then be bounded as
notice that when consider (ng(n/d) the first term can 2nL
(1 - am)22 = (1 log(nnd)) 2L
d
2L
n
Additionally, is straightforward that adL d'login/d)) conclude that these two terms at the same order. Therefore, the optimal choice of learning rate is O + which optimize the excess risk
When omitting the factors of log, we
R(w(L)) gd - -
(L)
This completes the proof.
Here we provide further discussions regarding the upper bound for the population loss achieved when choosing the optimal learning rate and initialization. The constant C' represents an irreducible term arising from the variance of the model. Such irreducible term always exists when considering least-squares loss, similar to the noise variance in classic linear regression problems. Therefore, when considering the problems with least-square loss function, it is common to define R( ((((()) C as the excess risk and attempt to minimize this term. Consequently, Theorem E.2 reveals that when using the optimal parameters, the excess risk R R(w! (w!))) C will converge the context length goes gd
infinity.
F Proof of Lemma 3.4 and Theorem 3.5
In this section, we provide the proof for both Lemma 3.4 and Theorem 3.5
W.L.O.G, we assume
that > in the subsequent proof. This implies that Vhe -1, Mqueery = 1 and E(w) P((w, Xquery) < 0) for any w. Then we first introduce a lemma providing a closed form for w gd which the parameter vector of the linear model trained by gradient descent with the optimal parameters derived in Theorem 3.3 and data (Xhe> Phe).
31

---

Lemma F.1. For the gradient descent iterates w w()) it holds that
(1)
cf + all)
(F.1)
gd
for all I E (0,1....,,T T.). e is the coefficient of 3 of initialization w (0) Wed and a(I) follows that
aN x IB3') 1 c(x 3*) all) -1- (1-
X
2
Proof of Lemmal F. We prove this lemma by induction. It straightforward that w cB and
W cf aNx1, complying with the formula (F.I) By induction, we assume for 1-th iteration. Then at the + 1-th iteration, we have
F.. still holds
w (l+1)
(La aNxxx]] ) w
gd aNx1
= -Ia--and - -la-aNxip xI] (e3 + +a()xx) -oNx:
= c3 + (a(1)( - x aNN(1 + ,(x1, = c,3* + a(l + 1)x1-
Additionally, by the fact all + +1 = a()) -oN |x_ - -aN(1 c(x1, we can derive that
1 + c(x_, 3*)
aN x 13) a(1) + 1 c(x1, 3*)
(all 1)
x
x 1
11+ c(x Br)
= (1- -11-aNl1x a
| 12
This implies that
a(l) - aaNaaNlx [ 1 1(x1,,
||x 112
which completes the proof.
Based on the closed form of w gd obtained by Lemma F.1 we are ready to prove Lemma 3.4 and
Theorem 3.5
Proof of Lemma 3.4 By Lemma F.I the output of the linear model trained via gradient descent
Xquery can be expanded as
(L) (w
gd Xquery) (cB" + a a(L)x1, tow")
=(," x1) [x]] + ca (ww,, ,
-cBB'xx x1) (1 a1- 123) (1 c(xL, + co(w", (w".,"))
coo("" ,8) 11 (1 a.N x (12) - c(x 37)).
(F2)
By utilizing the independence among w' a. and x and law of total expectation, we can derive that
E(w w (L)
P
co - (1 aN x (12) (1 cxxi,, 1*)) s $0)
gd
E
P C0 w. ,3* - 1 + (1 - aNN |x (12) (1 c(x1,3")) s w", xx xx]]
E Fg -1 (1-1 x ( (1 c(xL, 3)1)
(F33
c(w*,0*
where F-(-) is the cumulative distribution function of a. Similarly, we also have
xL.8")
E(w (0) gd
=
E
Fo
w B-)
32

---

Therefore, by Taylor's first order expansion, we have
E(W(d) - E(w (0)) = E F2(E) (1 (1 aNNll 13) ()) o(x1 171)
gd
1t 1+c(w,
Nd nL
Nd8/2
1
Ō Ö(Vd) S Ō
S Ø(1,
n
where the first inequality utilizing the concentration results that ||x] ([) and (x1, = Ö(vd).
The second inequality holds by the fact (1 1 (M)) = 1 - 0 - by our condition n S
o(d3/2 /n), which also implies the last inequality holds. Therefore, we finish the proof.
L
In the next, we prove Theorem 3.5
Proof of Theorem 3.5 Similar to the proof of Lemma 3.4
we have that
(1 aN || 12: (1 + c(x 3x))
E(w (I)
gd = E
Fo
c/w,,
(1 - NN ||x (12)
(1 - a.V | [x1) " (xx,, / (x11, 3²)))
E
Fo
(5)
c(w", B*)
(w*, B²)
1
(1- 12)
(1 (1- aN [x]
)}
(1
E Fo 1 (1 (1-(1-6(1)))
6 ( c(w",83}
E Fo
(*)
where the third inequality holds as (ign((X1, 3*) is independent with [X ) and || 11 and (E) is a constant. Additionally, let a follows the uniform distribution from a to b, then we can expand the expectation above
El E(w (L) Rd
1
θ (d))
(1 (1-6( (2)
=E
Fo
S a
<(w*,B"
c(ww,Br
(1- θ
Nd
1 1 (1 0 Nd a
nT.
+
Fo
o/w,, ,
c(w', ,3²)
(1 θ Nd
(1 5 o( N)))
F
b
, 3²))
c/w,,B' ,
Nd
1 (1 5 Nd))
S b }
1
1
EE
1
<
}
nL
c(ww Ba)
c/w,,
(1 - O ))
E
1
b
c(w",gy)
1 1- ē
Nd 4
<a
nL
where 61, C2 are two positive scalars solely depending on a, b and the distribution of w". This completes the proof.
G Technical lemmas
In this section, we introduce and prove some technical lemmas utilized in the previous proof. Lemma G.1. Let x ~ (O,,,, La), and W1, w 1d be two vectors independent of x, with ||will2- 1, then it holds that
2
Ex (w, x²) (W1,)) = (w. w1).
33

---

Proof of Lemma G.I Since ||w1l2 1, let r = [w1, W2, wal E R be the orthogonal matrix with Wi being its first column. Then we have
Exl(w, x) bagn(ww)) x))] Ex /w rr (sgg(((1,x)))
2 V (w, w1).
w, wA Ex] x)
where the last equality holds since (WE, x) (O for all & E [d], (Wk, x) and (Wky, x) are independent when k1 +f k2, and B3((,, sign( (w)) x))] = [B| (w1,
This completes
the proof.
For the next lemmas, we follow the notation we used in previous section that Σ Ei= X,X:
Lemma G.2. For any k N, it holds that E. E cla. where is scalar.
Proof of Lemma G.2 Let T be any orthogonal matrix, then we have rx; ~ (O,, Ia). This implies that D (Tx;) Tx;) has the same distribution with Σ. Therefore, we can derive that
TE E {E Txx, (Tx)) E|S"
holds for any orthogonal matrix r, which implies that E must be at the form cla. This completes the proof.
cl, for some scalar by binomial formula it can be expanded as a summation of polynomials of Σ, which all have the expectations with the form clą.
Lemma G.2 implies that [[( Σソ*
Lemma G.3. For any i N, it holds that
E S
Σ x(y:) ) cB*
where c is some scalar.
Proof of Lemma G.3
By binomial theorem, we have
k-
x:yli
xxxx7
E[(x:x7 ym x,yi].
i=1
1=1
=0
p
= cla for some scalar c. In the next, it suffices to show that [[(xxx]) )a1 x:Vi] cB for some scalar Since ||w*||2 1, let T [w", W wa] E Rd be the orthogonal matrix with w being its first column, and let x = r Tx, x, (o, This implies that yi = digu(w",,r coordinate of x/. Based on this, for any fixed w*, we can further derive that
By Lemma G.2 we already obtain that B (Erz xxxx7
= sign(x(,1) which is the first
x th x w*] [X
sign(x(,)) = [ 2* x sign(x((,)) cw". =w
The last equality holds as ||x||| 22r is even function for each coordinate of x/, which implies that finally obtain that
E | 2k, x(,, sxgn
for any j [d] and j + 1. Therefore,
[[xxxx ) f1 x
[xxxx y x
cE [w"] =08,
which completes the proof.
Lemma G.4 (Theorem 9 in
For any 8 > 0, with probability at least 8, it holds that,
d
log(1/5) (1)81
E Ia
max
n
n n2
34

---

Lemma G.5. For any > 0, with probability at least 1 - 8. it holds that,
sign((w", xi))x; - "V w s <0 < (ndl///)))
Proof of Lemma G.3] Similar to the previous proof technique, let r = [w", W2..... Wal E RR be the orthogonal matrix with w being its first column, and let x r X; NN Then we can derive that
Σ sign((w",x,)) VV w
Σ
[x4,1
sign x (x11) wy-
j=2
Since x is a subgaussian random variable with expectation
x by Hoeffding's 's inequality we
can derive that with probability least 8/d,
Σ
[x1,1
2 V < < log(d/o))
Additionally, when 1, sign(x', )x4,, still follows a standard normal distribution (A standard normal random variable times an independent Rademacher random variable is still standard normal random). Therefore, we can also derive that
Σ sign(x(,)x1, S O(V llog(d/0)
holds with probability at least
Then by taking union bound, can finally obtain that
2 w
2
n
sign((w xx))x; ny
[x
+
sign(xX,)
S (nd log(//))
The first equality holds by the orthogonality among w W.. w This completes the proof.
Lemma G.6. For any 5 > 0, with probability at least 1 8, it holds that,
x 11 (d - 1) < 0 dlog((//))
|*, < (Vdllg((1/8)
Proof of Lemma G.6 By the fact that ||x|||3 X²-1 we have E Ei[[x] d 1. Then by the
Bernstein's 's inequality, we can obtain that
xx 112 (d- O(V
holds with probability least 8/2. Besides, since (x1,8%, ~ ~ (*)),,)), by applying the tail bounds Gaussian distribution, we can obtain that
(x1,3g)) S OVV dlog((/5))
holds with probability at least - 6/2. By applying a union bound, we obtain the final result.
H Experimental setup
H.1 Context hijacking in GPT-2
This section will describe our experimental setup for context hijacking on LLMs of dif- ferent depths. We first construct four datasets for different tasks, including language, country, sports, and city. The samples in each dataset consist of four parts: prepend,
35

---

error result, query, and correct result. Each task has a fixed template for the sam- ple. For the language, the template "[People) did not speak ferror result). The native language of (People) is {correct result}". For the country, the tem- plate "(People) does not live ferror result). (People) has a citizenship from fcorrect result}". For the sports, the template is "[People) is not good at playing ferror result}. {People}'s best sport is {correct result}". For the city, the template is "(Landmarks) is not in {error result). {Landmarks} is in the city of {correct result}". We allow samples to have certain deviations from the templates, but they must generally conform to the semantics of the templates. Instance always match the reality, and the main source of instances is the CounterFact dataset [48]. In our dataset, each task contains three hundred to seven hundred specific instances. We conduct experiments of different sizes. Specifically, we consider GPT2, GPT2-MEDIUM, GPT2-LARGE, and GPT2-XI.. They have 12 layers, 24 layers, 36 layers, and 48 layers, respectively. We construct pipeline that test each model on each task, recording the number of prepends for which the context just succeeded in perturbing the output. For those samples that fail to perturb within certain number of prepends (which is determined by the maximum length of the pre-trained model), we exclude them from the statistics. Finally, we verify the relationship between model depth and robustness by averaging the number of prepends required to successfully perturb the output.
GPT2 [60]
H.2 Numerical experiments
We use extensive numerical experiments to verify our theoretical results, including gradient descent and linear transformers.
We use a single-layer neural network as the gradient descent model, which contains only one linear hidden layer. Its input dimension is the dimension d of feature x, and we mainly experiment on d = { (15, 20, 25}. Its output dimension is 1, because we only need to judge the classification result by its sign. We use the mean square error as the loss function and SGD as the optimizer. All data comes from the defined training distribution Drr- The hyperparameters we set include training context length N 50, mean of the Gaussian distribution 3 1, variance of the Gaussian distribution Σ = 0.1 (then normalized). We initialize the neural network to cB, and then perform gradient descents with steps Steps grid search to search for the optimal e and Ir for the loss function. This is equivalent to the trained transformers of layers 1 to 8 learning obtain the shared signal cB3 and the optimal learning rate Ir for the corresponding number of layers. Then they can use in-context data to fine-tune cf to a specific w*.
Gradient descent:
21.2...., 8} and learning rate Ir. We use
After obtaining the optimal initialization and learning rate, we test on the dataset from Din. Again, we set exactly the same hyperparameters as above. In addition, we set the in the test distribution to 0.1.
Linear Transformer: We train on multi-layer single-head linear transformers and use Adam as the optimizer. The training settings for models with different numbers of layers are exactly the same. We use the initial learning rate Ir E (0.0001. 0.0002), and the training steps are 600,000. We use a learning rate decay mechanism, where the learning rate is decayed by half every 50,000 training steps. For training and testing data, we set the data dimension d = 20 and the training context length N = (10, 30, 40). We use batchsize of 5,000 and apply gradient clipping with a threshold 1. Each experiment takes about four hours on average on a single NVIDIA GeForce RTX 4090 GPU.
H.3 Context hijacking in real-world LLMs
we construct four different topics of datasets, including city, country, sport, and language. However, this dataset is more diverse, which helps simulate more realistic situations. Below is a detailed description of the dataset construction process, using an example from the dataset.
Similar to Appendix H.I
First, we will design a fact retrieval problem. It is direct question, such as "Of all the sports, Maria Sharapova is most professional in which one? The answer is". We want the model to predict the next token "tennis".
Next, we will choose topic that is factually correct. For the example above, we can choose the topic that "Maria Sharapova is not professional in rugby".
36

---

0.8
Test Accuracy
0.6
0.4
Layer 7
Layer
0.2
Layer 5
Layer 4
10
15
20
25
30
35
40
Context Length
Figure 6: Standard transformers experiments with different depths. Testing the trained standard transformers (GPT-2 architecture [60]) on the test set, as the number of interference samples increases, the model classification accuracy decreases and gradually converges. The results also show that deeper models are more robust.
Finally, will add factually correct context prefixes of varying lengths before the question. Each sentence of this context prefix will describe the topic that has been determined from a different perspective and with different words. That is, paraphrase the hijacking context instead of repeating them. In our example, these sentences could be "Rugby is not sport that Maria Sharapova is adept playing", "Maria Sharapova's tennis skills do not translate well to rugby", "The physical demands of rugby are not ones with which Maria Sharapova is familiar", etc. The model is then asked the same question. If the model predicts "tennis". then correct. the model predicts "rugby", we call this "label flipping".
The number of samples in each dataset ranges from hundreds thousands. We filter questions that are too difficult based on the model's 's own capabilities and the difficulty of the questions, which means that the model could always correctly answer direct questions without hijacking context. We conduct experiments on Qwen2.5 base models of different sizes (depths) and corresponding instruction fine-tuned versions.
I Additional experiments
L.1 Robustness of standard transformers with different number of layers
To generalize the results to more realistic settings, we transfer the experiments from linear transform- ers to larger and standard transformers, such as GPT-2 (60). We train and test GPT-2 with different numbers of layers based on exactly the same settings as the linear transformers experiments. The results once again verify theory (Figure 6). As the context length increases, the model's 's accuracy decreases, but increasing the number of layers of the model significantly improves the robustness, indicating that our theory has more practical significance. Then we describe the setup of standard transformers experiments briefly.
Setup: We use the standard transformers of the GPT-2 architecture for the experiments, and the main settings are similar 330. We set the embedding size to 256, the number heads to and the batch size to 64. We use a learning rate decay mechanism similar linear transformers experiments, with an initial learning rate of 0.0002, and then reduced by half every 200,000 steps, for total of 600,000 steps. We use Adam the optimizer.
37

---

a=0.8
d=0.4
Layer
Layer
Test Accuracy
Layer
Test Accuracy
Layer:
Layer
Layer
Layer
Layer
0.5
Context Langth
Context Longth
0=0
c=0.1
1.0
Layer
Layer
0.9
Layer
Layer
0.9
Layer
Layer
70t Accuracy
0.8
Layer
Test Accuracy
0.8
Layer
0.7
0.7
0.6
0.5
25
50
75
100
125
150
0
25
50
100 Context Length
75
125
150
Context Length
Figure 7: Linear transformers experiments with different depths and different In real-world semantics, smaller a means stronger interference. Comparing the test performance of the model under different can find that decreases, the robustness of the model decreases significantly, which verifies the rationality our modeling.
1.2 Linear transformers facing different interference intensity
In this section, we mainly discuss how the robustness of the model changes with the interference intensity. In our modeling, the interference intensity is determined only by the distance between the query sample and the similar interference samples defined in the test set, that is, by variable a in Dic. In real-world observations, according to the idea of the induction head [5T], the more similar the context prepend used for interference is to the query, the more likely the model is to use in-context learning output incorrect results. Therefore, examine different model conforms to the actual real-world interference situation, that is, to verify the rationality of our modeling.
determine whether the
Observing the experiment results in Figure when a gradually decreases from 0.8 to 0.1, that is. the interference intensity of the data gradually increases, the classification accuracy of the model decreases significantly. When a is larger and interference context is less, the model can always classify accurately, indicating that weak interference does not affect the performance of the model, which is consistent with real observations. Various experimental phenomena show that our modeling of the context hijacking task by the distance between the interference sample and the query sample is consistent with the real semantics.
1.3 Robustness of nonlinear transformers for nonlinear classification
To extend our theoretical results to more complex situations, we consider nonlinear models and nonlinear tasks. We conduct preliminary experiments on nonlinear classification (Figure B). More specifically, we change (w, x) to (w, x)² - C, where C' is constant. We conduct the experiments on the multi-layer ReLU attention transformers. The results show that even in the nonlinear case, the model still tends to be more robust as it gets deeper, which is consistent with our theoretical results. This strengthens our conclusions and shows that our theory is not limited to the linear case, but is also valid on more complex and practical nonlinear tasks.
38

---

Layer 6
0.8
Layer 5
Layer 4
Test Accuracy
0.6
Layer 3
0.4
0.2
0.0
25
50
100 Context Length
75
125
150
Figure 8: We introduce nonlinear problem by changing (w,x) (w, x)² C, where C is constant. We conduct the experiments on the multi-layer ReLU attention transformers. The results show that the model still tends to be more robust it gets deeper, which is consistent with our thcoretical results.
39


Reviews:
[{"review_id": "2siniYblXJ", "reviewer": "Reviewer_APEh", "summary": "This paper provides a theoretical framework and empirical support for analyzing the robustness of linear transformers against “context hijacking,” a phenomenon where factually correct yet misleading context examples can disrupt model predictions. The authors model in-context learning in linear transformers as multi-step gradient descent and derive how robustness scales with model depth, context length, and hijack token count. Their theory demonstrates that deeper models exhibit stronger robustness, which aligns with experimental results.", "questions": "N/A", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strength\n\n1) The work identifies and rigorously formalizes the problem of context hijacking, a subtle yet impactful threat to LLM robustness.\n\n2) The paper offers a clear and elegant connection between transformer depth and robustness via optimal multi-step gradient descent modeling.\n\n3) Numerical simulations back the theoretical predictions and explain why deeper transformers resist hijacking more effectively.\n\n### Weakness\n\n1) Limited generalization to nonlinear settings: Though nonlinear experiments are mentioned (Appendix I.3), their depth is limited (only showing a depth analysis). It remains unclear how well the conclusions carry over to architectures used in practice.\n\n2) Realism of assumptions: The study is limited to linear transformers and synthetic data. While this allows for tractable theory, it diverges from practical LLM architectures and real-world data distributions, potentially reducing the direct applicability of findings. More empirical experiments are encouraged to address this. Even using synthetic data, try to paraphrase the hijacking contexts instead of repeating them.", "was_revised": true, "final_justification": "I increased the score because the authors connect the paper better to real-world situations", "rebuttal": "Thank you for your informative feedback! We address your comments as follows:\n\n>**Q1:** It remains unclear how well the conclusions carry over to architectures used in practice. More empirical experiments are encouraged to address this.\n\n**A1:**\n\n**Common settings for theoretical work.**\n\nFrom a theoretical perspective, it is a common method to abstract real-world problems into linear problems for analysis, because linear problems have sufficient representation power, supported by many previous works [1-9].\n\n**Experiments on practical LLM architectures and real-world data distributions.**\n\nThen, based on your suggestions, we realize the importance of verifying our conclusions in a more realistic architecture. We conduct extensive supplementary experiments on LLMs of varying depths across diverse topic tasks to demonstrate the validity of our conclusions in real-world contexts. Our dataset is constructed as follows.\n\n**Dataset construction and settings.**\n\n* First, we will design a fact retrieval problem. It is a direct question, such as \"Of all the sports, Maria Sharapova is most professional in which one? The answer is\". We want the model to predict the next token is \"tennis\". \n\n* Next, we will choose a topic that is factually correct. For the example above, we can choose the topic that \"Maria Sharapova is not a professional in rugby\".\n\n* Finally, we will add factually correct context prefixes of varying lengths before the question. **Each sentence of this context prefix will describe the topic that has been determined from a different perspective and with different words. That is, paraphrase the hijacking context instead of repeating them.** In our example, these sentences could be \"Rugby is not a sport that Maria Sharapova is adept at playing\", \"Maria Sharapova's tennis skills do not translate well to rugby\", \"The physical demands of rugby are not ones with which Maria Sharapova is familiar\", etc. The model is then asked the same question. If the model predicts \"tennis\", then it is correct. If the model predicts \"rugby\",  we call this \"label flipping\".\n\nWe design four datasets with different topics, including city, country, sports and language. And the number of samples in each dataset ranges from hundreds to thousands. We divide the context hijacking into eight different levels according to the length of the context prefix, from level 1 to level 8, which means the context has 10 to 80 sentences. We filter out questions that are too difficult based on the model's own capabilities and the difficulty of the questions, which means that the model could always correctly answer direct questions without hijacking context. We conduct experiments on Qwen2.5 base models of different sizes (depths) and corresponding instruction fine-tuned versions. The tables below show the label flipping rates of different models for different levels of context hijacking.\n\n**Experiment results**\n\n1. City.\n\nQwen2.5-Base:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.1320 | 0.2098 | 0.2598 | 0.3096 | 0.3487 | 0.4020 | 0.4337 | 0.4906 |\n| Qwen2.5-1.5B (28 Layers) | 0.0287 | 0.0589 | 0.1005 | 0.1471 | 0.1795 | 0.1950 | 0.2161 | 0.2411 |\n| Qwen2.5-3B (36 Layers) | 0.0230 | 0.0437 | 0.0587 | 0.0780 | 0.0922 | 0.1006 | 0.1067 | 0.1164 |\n\nQwen2.5-Instruct:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B-Instruct (24 Layers) | 0.1181 | 0.1693 | 0.2171 | 0.2684 | 0.2982 | 0.3473 | 0.3836 | 0.4068 |\n| Qwen2.5-1.5B-Instruct (28 Layers) | 0.0221 | 0.0319 | 0.0543 | 0.0823 | 0.1025 | 0.0959 | 0.0983 | 0.0958 |\n| Qwen2.5-3B-Instruct (36 Layers) | 0.0180 | 0.0337 | 0.0453 | 0.0526 | 0.0570 | 0.0690 | 0.0699 | 0.0765 |\n\n2. Country.\n\nQwen2.5-Base:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.4094 | 0.5769 | 0.5936 | 0.5977 | 0.6173 | 0.6603 | 0.6500 | 0.6692 |\n| Qwen2.5-1.5B (28 Layers) | 0.3125 | 0.3906 | 0.5000 | 0.5167 | 0.5547 | 0.5469 | 0.5781 | 0.5781 |\n| Qwen2.5-3B (36 Layers) | 0.1708 | 0.1868 | 0.1893 | 0.2198 | 0.2253 | 0.2527 | 0.2500 | 0.2555 |\n\nQwen2.5-Instruct:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B-Instruct (24 Layers) | 0.6602 | 0.7411 | 0.7624 | 0.7716 | 0.7884 | 0.7884 | 0.8062 | 0.7827 |\n| Qwen2.5-1.5B-Instruct (28 Layers) | 0.4805 | 0.5132 | 0.5171 | 0.5188 | 0.5200 | 0.5286 | 0.5232 | 0.5329 |\n| Qwen2.5-3B-Instruct (36 Layers) | 0.1214 | 0.1947 | 0.2627 | 0.2681 | 0.2745 | 0.2807 | 0.2809 | 0.2872 |\n\n3. Sports:\n\nQwen2.5-Base:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.7489 | 0.7583 | 0.7621 | 0.7738 | 0.7708 | 0.7788 | 0.7914 | 0.8006 |\n| Qwen2.5-1.5B (28 Layers) | 0.5255 | 0.5842 | 0.5856 | 0.5891 | 0.5987 | 0.5910 | 0.6020 | 0.6136 |\n| Qwen2.5-3B (36 Layers) | 0.1103 | 0.1177 | 0.1302 | 0.1336 | 0.1381 | 0.1403 | 0.1398 | 0.1484 |\n\nQwen2.5-Instruct:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B-Instruct (24 Layers) | 0.5419 | 0.5577 | 0.5968 | 0.6122 | 0.6554 | 0.7035 | 0.7131 | 0.7115 |\n| Qwen2.5-1.5B-Instruct (28 Layers) | 0.2500 | 0.2941 | 0.3627 | 0.4265 | 0.4412 | 0.4796 | 0.4853 | 0.4900 |\n| Qwen2.5-3B-Instruct (36 Layers) | 0.1493 | 0.1806 | 0.1967 | 0.2000 | 0.2103 | 0.2161 | 0.2220 | 0.2418 |\n\n4. Language:\n\nQwen2.5-Base:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.3509 | 0.5077 | 0.5666 | 0.6107 | 0.6377 | 0.6441 | 0.6383 | 0.6399 |\n| Qwen2.5-1.5B (28 Layers) | 0.0732 | 0.1279 | 0.1737 | 0.2284 | 0.2485 | 0.2922 | 0.2853 | 0.3013 |\n| Qwen2.5-3B (36 Layers) | 0.0435 | 0.0722 | 0.0740 | 0.0922 | 0.1043 | 0.1024 | 0.1116 | 0.1090 |\n\nQwen2.5-Instruct:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B-Instruct (24 Layers) | 0.3049 | 0.4510 | 0.5205 | 0.5611 | 0.5712 | 0.5639 | 0.5662 | 0.5733 |\n| Qwen2.5-1.5B-Instruct (28 Layers) | 0.0658 | 0.1142 | 0.1228 | 0.1287 | 0.1385 | 0.1410 | 0.1397 | 0.1452 |\n| Qwen2.5-3B-Instruct (36 Layers) | 0.0168 | 0.0288 | 0.0313 | 0.0353 | 0.0396 | 0.0383 | 0.0409 | 0.0375 |\n\nWe can find that **in practical LLMs, longer hijacking context will significantly increase the label flipping rate (leading to lower accuracy), while increasing the model depth can well alleviate this problem**. The experiment results are consistent with our theoretical conclusions, **indicating that our theoretical results can be generalized to deeper and larger LLMs in practice**. Additionally, we find that instruction fine-tuning can improve the model's robustness to context hijacking in most cases, but the effect is not significant, which provides new insights for future work, such as adversarial optimization. This suggests that our work can provide insights into real-world problems.\n\nWe will provide all the experiment results and detailed experimental settings in our revised paper. We believe that our experiments on real-world tasks and architectures fully validate the applicability of our conclusions and hope that these results address your concerns.\n\n[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.\n\n[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.\n\n[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.\n\n[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[6] Mahankali, et al. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ICLR 2024.\n\n[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.\n\n[8] Ren, et al. Learning and Transferring Sparse Contextual Bigrams with Linear Transformers. NeurIPS 2024.\n\n[9] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025."}, {"review_id": "ANyobIORe2", "reviewer": "Reviewer_AmqY", "summary": "I really enjoyed this paper. The paper discusses context hijacking. The terminology works as follows. If a context is adversarial, i.e., the prompt contains false information and adding the context to a prompt hurts the LM's ability to do the task, then that is, in a sense, expected. We do want LLMs to react to context. On the other hand, sometimes adding a correct statement to the context also derails the model. This is called context hijacking.  \n\nAs I understand it, the idea is to come up with a theoretical model of context-hijkacing. There is a simple model that samples contexts and hijacked contexts. A number of propositions are proposed about building transformers that model his.", "questions": "* Is there any way to make the paper more accessible to readers who don't have a strong background in optimization? The answer to this could be no\n* Can the techniques be used for other, similar problems?", "limitations": "None", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 4, "originality": 4, "strengths_and_weaknesses": "This paper has a great premise and is well polished. I enjoyed reading it. The weaknesses are that it's quite difficult to get through. I am sure this is my fault, as I don't have the right background. The whole idea seems to be to cast the context hijacking problem as one of optimization. I was unable to follow most of the technical details.", "was_revised": true, "final_justification": "I think this is a great paper and I would love to see it in the conference.  I was happy with the authors' response.", "rebuttal": "Thank you for your recognition of our work and your constructive questions!\n\n>**Q1:** More accessible way to understand?\n\n**A1:** Thanks for your suggestion. We will try our best to elaborate on our framework with detailed explanations, critical take-home messages, and without any mathematical formulas. \n\n1. **Background of in-context learning (ICL).** Empirical evidence shows that modern LLMs can adapt to diverse tasks solely based on the provided context data, without requiring any task-specific fine-tuning. For example, the same model can solve mathematical equations when given a math problem, yet seamlessly switch to translating text when presented with a sentence in another language, all without explicit retraining for either task. This is the so-called in-context learning capacity of transformers. \n\n2. **A theoretical perspective to understand ICL: single attention layer conducts one-step gradient descent with context data.** One theoretical explanation regarding this phenomenon is that the forward propagation of each layer in transformers is equivalent to conducting one-step gradient descent on context data ([1, 2, 3, 4] and Proposition 3.1 of our paper). You can imagine that every time a pre-trained transformer performs inference for a given query token, it internally constructs an 'implicit linear model' to generate the output. The initial parameters of this model are fixed and determined solely by the pre-trained transformer, independent of the context or query. However, while the initialization remains the same for all 'implicit linear models' corresponding to different tasks/context data/input tokens, when each layer of transformers forward propagates the input query, the parameters of the 'implicit linear model' will be updated one step based on context data. **Intuitively speaking, every forward pass through one attention-layer is effectively a learning step, where the model \"adapts\" to the context data.**\n\n3. **Depth-dependent optimal learning rates of transformer models.** Since the model’s depth determines the number of 'implicit gradient descent' steps available, the optimal learning rate achieving the training loss minimization would vary according to different depths of transformers. Specifically, the optimal learning rate of 'implicit gradient descent' on context data scales inversely with the depth of the transformer model (Theorem 3.3). This is intuitive: A shallow transformer must use a larger learning rate to ensure the 'implicit linear model' learns sufficiently from context data within limited steps. In contrast, a deep transformer, with more iterations available, opts for a smaller learning rate to enable fine-grained adjustments when near the optimum. **In short, when learning from context data, shallow models are more 'aggressive', while deep models are more 'conservative'.** \n\n4. **Context-hijacking and robustness of transformers in terms of depth.** For now, we can explain the phenomenon of 'context hijacking' and our findings of robustness against hijacking. Take Figure 1 as an example. We can conjecture that the context example *'Rafael Nadal is not good at playing basketball'*, and the query *'Rafael Nadal's best sport is'* share similar embeddings in the representation space. When the model learns from contexts through 'implicit gradient descent', the embedding similarities cause the model to conflate their semantic meanings, leading to the incorrect answer 'basketball'. In some sense, context-hijacking occurs when the 'context data' and 'query' belong to different distributions: while their embeddings appear similar, their actual meanings are completely opposite. (In classification terms, this would be cases where inputs ($x$) are similar but labels ($y$) are contradictory.) Because the context data and query follow different underlying patterns, learning from such contexts becomes ineffective or even detrimental to model performance. This explains why deep transformers show greater robustness against hijacking: **shallow models' aggressive learning approach makes them highly susceptible to interference from hijacking examples, whereas deep models employ a more cautious, incremental learning approach that allows for progressive refinement and better resistance to corrupting influences.**\n\nWe hope that such an outline can help you better understand the logic of our theoretical framework, and we will incorporate these discussions into our revisions.  \n\n>**Q2:** Can the techniques be used for other, similar problems?\n\n**A2:** Yes, our framework can be extended to several similar problems.\n\n1. Since we employ the $\\ell_2$ loss for classification tasks, a natural extension would be to adapt it to regression tasks.\n\n2. A key technical advantage of our framework is its inherent capacity to handle distributional shifts between: (1). context data and query data; (2). training data and test data. This makes our framework particularly suitable for investigating:\n\n- Adversarial attacks: A prevalent scenario where context data and query data exhibit distributional discrepancies. Specifically, the context data may contain factually incorrect information.\n\n- Out-of-distribution (O.O.D.) cases: A common situation where training context data and test context data follow different distributions. This typically occurs when applying a pretrained model to specific downstream tasks.\n\n[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.\n\n[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.\n\n[4] Bai, et al. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. NeurIPS 2023."}, {"review_id": "nBx7S4C02i", "reviewer": "Reviewer_kpWy", "summary": "This paper investigates the phenomenon of context hijacking in Transformer-based large language models, where the presence of factually correct context can mislead a model into incorrect predictions. The authors formalize context hijacking as a theoretical linear classification problem and analyze it through a multi-layer linear transformer framework, drawing a connection between transformer layers and multi-step gradient descent optimization. A key theoretical result is that deeper transformers inherently offer greater robustness to context hijacking because they allow for finer-grained optimization, reducing the negative impact of misleading contexts. The paper rigorously derives optimal conditions for transformer initialization and learning rates as functions of model depth and context length, confirming these insights empirically through synthetic experiments. The primary contributions include a comprehensive theoretical framework for analyzing context hijacking, mathematical proofs of the relationship between transformer depth and robustness, and empirical validations that reinforce the theoretical findings.", "questions": "It is a well-established fact that linear transformers perform less well on in-context tasks [1]. So, why was this architecture chosen over the more traditional quadratic transformer?\n\n[1] Aksenov et al, Linear Transformers with Learnable Kernel Functions are Better In-Context Models", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strenghts\n\n* **Novel Theoretical Framework**: The paper introduces the first comprehensive theoretical analysis of context hijacking in transformers using a linear classification setup, which is both novel and practically relevant.\n\n* **Solid Mathematical Grounding**: The authors rigorously connect transformer layers with multi-step gradient descent. This explicit connection provides deep insights into the internal mechanics of transformers.\n\n### Weaknesses\n\n* **Linear Simplification**: While the linear classification setup enables clear theoretical insights, it limits the generalization of the results to complex, real-world nonlinear scenarios and tasks.\n\n* **Assumption on Data Distribution**: The analysis heavily relies on simplified assumptions such as isotropic Gaussian distributions for inputs and uniform distributions for certain parameters. Real-world contexts typically deviate from these idealized conditions.\n\n* **Limited Empirical Scope**: The numerical experiments primarily use synthetic data and simplified transformer structures (linear transformers). It remains unclear whether these findings generalize directly to practical, large-scale models.", "was_revised": true, "final_justification": "Thank you to the authors for the clarifications and additional experiments. Some of my initial misunderstandings stemmed from not being deeply familiar with the area of mechanistic interpretability, and the authors' response has helped resolve these. The new experiments also shed light on the practical relevance of the problem studied in the paper. That said, I believe the paper could benefit from an even broader set of experiments, particularly those exploring more diverse architectures and practical scenarios. Therefore, I will raise my score to 4.", "rebuttal": "We thank the reviewer for detailed comments and suggestions.\n\n>**Q1:** Simplified assumptions on linear tasks and Gaussian data distributions.  \n\n**A1:**\n\nWe would like to first point out that studying linear tasks with Gaussian inputs is a standard setting and has been widely considered in many existing theoretical works regarding transformers, particularly in the in-context learning literature [1-8]. \n\nIn addition, we emphasize that our work is the first to rigorously propose a theoretical framework - through formulating it into linear classification tasks - to investigate 'context hijacking', a theoretically unexplored phenomenon even under standard settings. With this simple yet intuitive mathematical formulation, we establish clear and precise quantitative characterizations in terms of the ICL capacities of transformers with respect to their depths. Specifically, our clear mathematical characterization effectively reveals that: **when learning from context, shallow transformers are more 'aggressive', whereas deep transformers are more 'conservative'**. While the exact mathematical form may not directly transfer to more complicated situations, our findings regarding depth-dependent learning strategies offer valuable theoretical insights that can guide future investigations into non-linear tasks, more complex data structures, or other aspects of multi-layer transformers' ICL. These theoretical conclusions are further supported by our experimental results on more complex scenarios (see A2).\n\nAdditionally, we emphasize that a key technical strength of our framework lies in its inherent ability to accommodate distributional shifts between: (1) context data and query data; (2) training data and test data. This characteristic makes our framework readily adaptable for studying both adversarial attacks and out-of-distribution performance. In summary, we believe this represents a good starting point.\n\n>**Q2:** Limited Empirical Scope: It is unclear whether these findings generalize directly to practical, large-scale models.\n\n**A2:**\n**Experiments on practical LLM architectures and real-world data distributions.**\n\nAlthough we experiment with nonlinear settings and the GPT-2 architecture (Appendix I.1 and I.3), encouraged by your suggestions, we realize the importance of generalizing our results to more realistic architectures. We conduct extensive supplementary experiments on LLMs of varying depths across diverse topic tasks to demonstrate the validity of our conclusions in real-world contexts. Our dataset is constructed as follows.\n\n**Dataset construction and settings.**\n\n* First, we will design a fact retrieval problem. It is a direct question, such as \"Of all the sports, Maria Sharapova is most professional in which one? The answer is\". We want the model to predict the next token is \"tennis\". \n\n* Next, we will choose a topic that is factually correct. For the example above, we can choose the topic that \"Maria Sharapova is not a professional in rugby\".\n\n* Finally, we will add factually correct context prefixes of varying lengths before the question. **Each sentence of this context prefix will describe the topic that has been determined from a different perspective and with different words. That is, paraphrase the hijacking context instead of repeating them.** In our example, these sentences could be \"Maria Sharapova's tennis skills do not translate well to rugby\", \"The physical demands of rugby are not ones with which Maria Sharapova is familiar\", etc. The model is then asked the same question. If the model predicts \"tennis\", then it is correct. If the model predicts \"rugby\",  we call this \"label flipping\".\n\nWe design four datasets with different topics, including city, country, sports and language. And the number of samples in each dataset ranges from hundreds to thousands. We divide the context hijacking into eight different levels according to the length of the context prefix, from level 1 to level 8, which means the context has 10 to 80 sentences. We filter out questions that are too difficult based on the model's own capabilities and the difficulty of the questions, which means that the model could always correctly answer direct questions without hijacking context. We conduct experiments on Qwen2.5 base models of different sizes (depths) and corresponding instruction fine-tuned versions. The tables below show the label flipping rates of different models for different levels of context hijacking.\n\n**Experiment results**\n\n1. City.\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) |0.1320|0.2098|0.2598|0.3096|0.3487|0.4020|0.4337|0.4906|\n| Qwen2.5-1.5B (28 Layers) | 0.0287 | 0.0589 | 0.1005 | 0.1471 | 0.1795 | 0.1950 | 0.2161 | 0.2411 |\n| Qwen2.5-3B (36 Layers) | 0.0230 | 0.0437 | 0.0587 | 0.0780 | 0.0922 | 0.1006 | 0.1067 | 0.1164 |\n\n2. Country.\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.4094 | 0.5769 | 0.5936 | 0.5977 | 0.6173 | 0.6603 | 0.6500 | 0.6692 |\n| Qwen2.5-1.5B (28 Layers) | 0.3125 | 0.3906 | 0.5000 | 0.5167 | 0.5547 | 0.5469 | 0.5781 | 0.5781 |\n| Qwen2.5-3B (36 Layers) | 0.1708 | 0.1868 | 0.1893 | 0.2198 | 0.2253 | 0.2527 | 0.2500 | 0.2555 |\n\n3. Sports:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.7489 | 0.7583 | 0.7621 | 0.7738 | 0.7708 | 0.7788 | 0.7914 | 0.8006 |\n| Qwen2.5-1.5B (28 Layers) | 0.5255 | 0.5842 | 0.5856 | 0.5891 | 0.5987 | 0.5910 | 0.6020 | 0.6136 |\n| Qwen2.5-3B (36 Layers) | 0.1103 | 0.1177 | 0.1302 | 0.1336 | 0.1381 | 0.1403 | 0.1398 | 0.1484 |\n\n4. Language:\n\n| Model | Level 1 | Level 2 | Level 3 | Level 4 | Level 5 | Level 6 | Level 7 | Level 8 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Qwen2.5-0.5B (24 Layers) | 0.3509 | 0.5077 | 0.5666 | 0.6107 | 0.6377 | 0.6441 | 0.6383 | 0.6399 |\n| Qwen2.5-1.5B (28 Layers) | 0.0732 | 0.1279 | 0.1737 | 0.2284 | 0.2485 | 0.2922 | 0.2853 | 0.3013 |\n| Qwen2.5-3B (36 Layers) | 0.0435 | 0.0722 | 0.0740 | 0.0922 | 0.1043 | 0.1024 | 0.1116 | 0.1090 |\n\nWe can find that **in practical LLMs, longer hijacking context will significantly increase the label flipping rate (leading to lower accuracy), while increasing the model depth can well alleviate this problem**. The experiment results are consistent with our theoretical conclusions, **indicating that our theoretical results can be generalized to deeper and larger LLMs in practice**. Additionally, we find that instruction fine-tuning (due to character limitations, please refer to the rebuttal to reviewer APEh) can improve the model's robustness to context hijacking in most cases, but the effect is not significant, which provides new insights for future work, such as adversarial optimization. This suggests that our work can provide insights into real-world problems.\n\nWe will provide all the experiment results and detailed experimental settings in our revised paper. We believe that our experiments on real-world tasks and architectures fully validate the applicability of our conclusions and hope that these results address your concerns.\n\n>**Q3:** It is a well-established fact that linear transformers perform less well on in-context tasks. So, why was this architecture chosen over the more traditional quadratic transformer?\n\n**A3:**\n**Conceptual inconsistency**\n\nWe want to clarify that the linear transformer [9] mentioned in the question is conceptually inconsistent with the linear transformer in our paper, or that the problems they aim to solve have different focuses.\n\nAs defined in the reference paper (Section 3.1), a linear transformer is a model that uses a kernel function to approximate the standard attention computation. The \"linear\" here means that the computational complexity of the approximation is linear in the sequence length. Its purpose is to improve computational efficiency, although this comes at the expense of some in-context learning performance.\n\nIn our definition (Section 2.2), a linear transformer is a model that removes the activation function from the standard transformer. The \"linear\" here means that the forward propagation of the model does not involve nonlinear calculations. This is a very common setup in transformer theory research [1-8, 10], aiming to create a mathematically tractable model that can be used to explain real-world problems from a theoretical perspective. We will discuss in detail the differences and connections between our work and the reference paper in the revision.\n\n[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.\n\n[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.\n\n[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.\n\n[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[6] Mahankali, et al. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ICLR 2024.\n\n[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.\n\n[8] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.\n\n[9] Aksenov, et al, Linear Transformers with Learnable Kernel Functions are Better In-Context Models. ACL 2024.\n\n[10] Ren, et al. Learning and Transferring Sparse Contextual Bigrams with Linear Transformers. NeurIPS 2024."}, {"review_id": "X2FXNA9JVb", "reviewer": "Reviewer_czCk", "summary": "The paper theoretically explores context hijacking within an in-context linear classification problem, utilizing linear transformers. It designs context tokens as factually correct query-answer pairs where queries are similar to the final query but have opposite labels. The authors develop a theoretical analysis on the robustness of linear transformers as a function of model depth, training context lengths, and the number of hijacking context tokens. The key result is a formal equivalence between L-layer transformers and L-step gradient descent from general initialization. The authors derive optimal learning rates and show that deeper models perform finer-grained updates, yielding exponentially stronger robustness against hijacking attacks. A derived error bound explains why models like GPT-4 are more robust than GPT-2. Empirical results confirm the theory. This work offers a principled foundation linking model depth to robustness.", "questions": "1 The analysis assumes repeated hijacking examples that are linearly projected to lie on the decision boundary. Can the framework extend to more realistic hijacking attacks where context examples are diverse, semantically richer, or adversarially optimized?\n2 The current analysis excludes MLPs, normalization, and nonlinearities. Do you expect your robustness characterization to hold for standard Transformers (e.g., GPT-style) if they approximately perform gradient-like steps?\n3 Given the reliance on Gaussian-distributed features, how well do you expect your theoretical findings to transfer to real-world NLP tasks with structured and non-isotropic embeddings?", "limitations": "The theoretical framework is restricted to linear classification under isotropic Gaussian assumptions, which may not generalize well to real-world embeddings or tasks involving long-range semantics.\n\nThe analysis is limited to linear attention-only transformers without MLP, layer norm, or softmax, which reduces relevance to practical architectures.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. This is the first rigorous framework analyzing context hijacking in transformers, reframing the problem via a clean equivalence to multi-step gradient descent.\n\n2. Demonstrates that deeper models achieve exponential robustness gains, offering concrete design guidance beyond vague depth heuristics.\n\n3. Theoretical predictions align precisely with experiments, explaining real-world trends like GPT-4’s superior robustness over GPT-2.\n\nWeaknesses\n1. Analysis is limited to linear classification under isotropic Gaussian assumptions, which limits applicability to realistic NLP tasks.\n\n2. Focuses solely on linear attention-only transformers, omitting nonlinearities, layer norm, and MLPs essential in practice.\n\n3. The analysis relies on the assumption that transformers implement a gradient descent optimizer, and the authors note that more complicated meta-optimization algorithms could lead to different theoretical results. While this is a common approach in theoretical transformer works, it might simplify the true underlying mechanisms.", "was_revised": true, "final_justification": "the responses addressed my questions well. I prefer to keep my current rating.", "rebuttal": "We appreciate your constructive questions and suggestions! We address them as follows:\n\n>**Q1:** Simplified linear tasks, Gaussian data, and transformers' architecture. \n\n**A1:**\n\nWe acknowledge that there exist some gaps between the linear classification tasks, Gaussian data, and linear attention-only transformers adopted in our theoretical framework, and real-world complex models and complex tasks. However, we also point out that due to technical challenges, studying linear tasks with linear transformers on Gaussian data is a standard theoretical setting, and has been widely considered in many existing theoretical works regarding transformers, particularly in the in-context learning literature [1-8]. \n\nIn addition, we would like to emphasize that the purpose of this work is to provide insights towards 'context-hijacking', a theoretically unexplored phenomenon. Proper mathematical simplification allows us to derive clear and precise quantitative characterizations of transformers' in-context learning capacities relative to their depths, rather than being biased by technical challenges. Specifically, our clear mathematical characterization effectively reveals that: **when learning from context, shallow transformers are more 'aggressive', whereas deep transformers are more 'conservative'**. While the exact mathematical form may not directly transfer to more complicated situations, our findings regarding depth-dependent learning strategies offer valuable theoretical insights that can guide future investigations into non-linear tasks, more complex data structures, or other aspects of multi-layer transformers' ICL. In fact, these theoretical conclusions are further empirically supported by our experimental results on more complex scenarios (see A2), guaranteeing the generalization capacity of our theoretical framework. \n\n>**Q2:** Limits applicability to realistic NLP tasks: can the framework extend to more realistic hijacking attacks where context examples are diverse, semantically richer or adversarially optimized?\n\n**A2:**\n**Theoretical extension.**\n\nOur theoretical framework can be extended to more general cases. For example, we can extend it to an out-of-distribution case, where the hijacking context examples and query follow different distributions (so it would be possible that context and query have opposite labels with similar embeddings). Notice that our conclusion that models' depth determines the optimal learning rates still holds (Theorem 3.3). It is natural when there exist significant distinctions between two distributions, deep transformers' 'conservative' learning strategies make them robust to hijacking.  \n\n**Experiments on practical LLM architectures and real-world data distributions.**\n\nWe strongly agree with your opinion and we realize the importance of generalizing our results to more realistic architectures. We conduct extensive supplementary experiments on LLMs of varying depths across diverse topic tasks to demonstrate the validity of our conclusions in real-world contexts. Our dataset is constructed as follows.\n\n**Dataset construction and settings.**\n\n* First, we will design a fact retrieval problem. It is a direct question, such as \"Of all the sports, Maria Sharapova is most professional in which one? The answer is\". We want the model to predict the next token is \"tennis\". \n\n* Next, we will choose a topic that is factually correct. For the example above, we can choose the topic that \"Maria Sharapova is not a professional in rugby\".\n\n* Finally, we will add factually correct context prefixes of varying lengths before the question. **Each sentence of this context prefix will describe the topic that has been determined from a different perspective and with different words. That is, paraphrase the hijacking context instead of repeating them, which makes context examples more diverse and semantically richer.** In our example, these sentences could be \"Maria Sharapova's tennis skills do not translate well to rugby\", \"The physical demands of rugby are not ones with which Maria Sharapova is familiar\", etc. The model is then asked the same question. If the model predicts \"tennis\", then it is correct. If the model predicts \"rugby\",  we call this \"label flipping\".\n\nWe design four datasets with different topics, including city, country, sports and language. And the number of samples in each dataset ranges from hundreds to thousands. We divide the context hijacking into eight different levels according to the length of the context prefix, from level 1 to level 8, which means the context has 10 to 80 sentences. We filter out questions that are too difficult based on the model's own capabilities and the difficulty of the questions, which means that the model could always correctly answer direct questions without hijacking context. We conduct experiments on Qwen2.5 base models of different sizes (depths) and corresponding instruction fine-tuned versions. The tables below show the label flipping rates of different models for different levels of context hijacking.\n\n**Experiment results**\n\n1. City.\n\n|Model|Level 1|Level 2|Level 3|Level 4|Level 5|Level 6|Level 7|Level 8|\n|:---|:---|:---|:---|:---|:---|:---|:---|:---|\n|Qwen2.5-0.5B (24 Layers)|0.1320|0.2098|0.2598|0.3096|0.3487|0.4020|0.4337|0.4906|\n|Qwen2.5-1.5B (28 Layers)|0.0287|0.0589|0.1005|0.1471|0.1795|0.1950|0.2161|0.2411|\n|Qwen2.5-3B (36 Layers)|0.0230|0.0437|0.0587|0.0780|0.0922|0.1006|0.1067|0.1164|\n\n2. Country.\n\n|Model|Level 1|Level 2|Level 3|Level 4|Level 5|Level 6|Level 7|Level 8|\n|:---|:---|:---|:---|:---|:---|:---|:---|:---|\n|Qwen2.5-0.5B (24 Layers)|0.4094|0.5769|0.5936|0.5977|0.6173|0.6603|0.6500|0.6692|\n|Qwen2.5-1.5B (28 Layers)|0.3125|0.3906|0.5000|0.5167|0.5547|0.5469|0.5781|0.5781|\n|Qwen2.5-3B (36 Layers)|0.1708|0.1868|0.1893|0.2198|0.2253|0.2527|0.2500|0.2555|\n\n3. Sports:\n\n|Model|Level 1|Level 2|Level 3|Level 4|Level 5|Level 6|Level 7|Level 8|\n|:---|:---|:---|:---|:---|:---|:---|:---|:---|\n|Qwen2.5-0.5B (24 Layers)|0.7489|0.7583|0.7621|0.7738|0.7708|0.7788|0.7914|0.8006|\n|Qwen2.5-1.5B (28 Layers)|0.5255|0.5842|0.5856|0.5891|0.5987|0.5910|0.6020|0.6136|\n|Qwen2.5-3B (36 Layers)|0.1103|0.1177|0.1302|0.1336|0.1381|0.1403|0.1398|0.1484|\n\n4. Language:\n\n|Model|Level 1|Level 2|Level 3|Level 4|Level 5|Level 6|Level 7|Level 8|\n|:---|:---|:---|:---|:---|:---|:---|:---|:---|\n|Qwen2.5-0.5B (24 Layers)|0.3509|0.5077|0.5666|0.6107|0.6377|0.6441|0.6383|0.6399|\n|Qwen2.5-1.5B (28 Layers)|0.0732|0.1279|0.1737|0.2284|0.2485|0.2922|0.2853|0.3013|\n|Qwen2.5-3B (36 Layers)|0.0435|0.0722|0.0740|0.0922|0.1043|0.1024|0.1116|0.1090|\n\nWe can find that **in practical LLMs, longer hijacking context will significantly increase the label flipping rate (leading to lower accuracy), while increasing the model depth can well alleviate this problem**. The experiment results are consistent with our theoretical conclusions, **indicating that our theoretical results can be generalized to deeper and larger LLMs in practice**. And we are able to transfer our theoretical findings to practical NLP tasks with structured and non-isotropic embeddings. Additionally, we find that instruction fine-tuning (due to character limitations, please refer to the rebuttal to reviewer APEh) can improve the model's robustness to context hijacking in most cases, but the effect is not significant, which provides new insights for future work, such as adversarial optimization. This suggests that our work can provide insights into real-world problems.\n\nWe will provide all the experiment results and detailed experimental settings in our revised paper. We believe that our experiments on real-world tasks and architectures fully validate the applicability of our conclusions and hope that these results address your concerns.\n\n>**Q3:** Does the current conclusion still hold for more complicated architectures if they approximately perform gradient-like steps? How about more complicated meta-\noptimization algorithms?\n\n**A3:**\n\nOur robustness characterization remains valid for any scenario where transformers exhibit gradient-like learning behavior. The foundation of our framework lies in establishing depth-dependent optimal learning rates (Theorem 3.3). Crucially, this core conclusion persists under your assumptions that they approximately perform gradient-like steps. Therefore, our theoretical predictions about depth-mediated robustness continue to hold effectively.\n\nRecent studies have proposed new perspectives on in-context learning, such as transformers' ability to approximate second-order optimization methods like Newton's method [9]. While the exact formulation may differ, we argue that our core theoretical insight in Theorem 3.3, the depth-dependent learning behavior, remains valid across different optimization paradigms. It is natural as long as the connection between the depth and the steps of optimizations still holds. \n\n[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.\n\n[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.\n\n[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.\n\n[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[6] Mahankali, et al. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ICLR 2024.\n\n[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.\n\n[8] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.\n\n[9] Fu, et al. Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression. NeurIPS 2024."}]

Author rebuttal:

Meta-review: This paper theoretically investigates the behavior of Transformers under the so-called context hijacking for linear classification. The authors first establish the equivalence between the L-layer transformers and L-step gradient descent, then derive the optimal learning rates and initialization of the equivalent L-step gradient descent problem, which are shown to be a function of the number of layers and the length of training context. Based this result, the authors analyze the robustness of Transformers under the impact of the number of layers, training context length and test context length. This theorical framework explains a common observation of Transformers -- a deep transformer model is more robust to context hijacking than a shallow transformer model since a deep transformer model employs a more cautious, fine-grained learning approach that allows for progressive refinement and better resistance to corrupting influences. Overall, this is an interesting work. Some major concerns raised by the reviewers include 
1) lack of experiments on practical LLM architectures and real-world data 
2) simplified linear tasks, data distribution and linear transformer architecture.

These concerns have been successfully cleared by the authors in their rebuttal with added new experimental results on Qwen2.5 model and elaborated explanation on the necessity and usefulness of using a simplified mathematical treatment. The authors need to improvement the submission with new results and promised modifications in the revised version.
Decision: Accept (poster)