On the Robustness of Transformers against Context Hijacking for Linear Classification

We appreciate your constructive questions and suggestions! We address them as follows:

Q1: Simplified linear tasks, Gaussian data, and transformers' architecture.

A1:

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

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

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?

A2: Theoretical extension.

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

Experiments on practical LLM architectures and real-world data distributions.

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

Dataset construction and settings.

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

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

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

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

Experiment results

1. City.

2. Country.

3. Sports:

4. Language:

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

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

Q3: Does the current conclusion still hold for more complicated architectures if they approximately perform gradient-like steps? How about more complicated meta- optimization algorithms?

A3:

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

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

[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.

[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.

[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.

[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.

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

[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.

[8] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.

[9] Fu, et al. Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression. NeurIPS 2024.

We thank the reviewer for detailed comments and suggestions.

Q1: Simplified assumptions on linear tasks and Gaussian data distributions.

A1:

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

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

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

Q2: Limited Empirical Scope: It is unclear whether these findings generalize directly to practical, large-scale models.

A2: Experiments on practical LLM architectures and real-world data distributions.

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

Dataset construction and settings.

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

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

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

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

Experiment results

1. City.

2. Country.

3. Sports:

4. Language:

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

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

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?

A3: Conceptual inconsistency

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

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

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

[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.

[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.

[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.

[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.

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

[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.

[8] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.

[9] Aksenov, et al, Linear Transformers with Learnable Kernel Functions are Better In-Context Models. ACL 2024.

[10] Ren, et al. Learning and Transferring Sparse Contextual Bigrams with Linear Transformers. NeurIPS 2024.

Thank you for your recognition of our work and your constructive questions!

Q1: More accessible way to understand?

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.

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

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

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

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

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

Q2: Can the techniques be used for other, similar problems?

A2: Yes, our framework can be extended to several similar problems.

1. Since we employ the $\ell_2$ loss for classification tasks, a natural extension would be to adapt it to regression tasks.

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

• Adversarial attacks: A prevalent scenario where context data and query data exhibit distributional discrepancies. Specifically, the context data may contain factually incorrect information.

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

[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.

[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.

[4] Bai, et al. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. NeurIPS 2023.

We are truly grateful for your recognition of our work! As the discussion phase draws to a close, we wish to ensure that our responses have addressed your concerns. It would be our pleasure to provide any further clarification you may desire. We greatly appreciate your valuable time and feedback.

Thank you for your informative feedback! We address your comments as follows:

Q1: It remains unclear how well the conclusions carry over to architectures used in practice. More empirical experiments are encouraged to address this.

A1:

Common settings for theoretical work.

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

Experiments on practical LLM architectures and real-world data distributions.

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

Dataset construction and settings.

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

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

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

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

Experiment results

1. City.

Qwen2.5-Base:

Qwen2.5-Instruct:

2. Country.

Qwen2.5-Base:

Qwen2.5-Instruct:

3. Sports:

Qwen2.5-Base:

Qwen2.5-Instruct:

4. Language:

Qwen2.5-Base:

Qwen2.5-Instruct:

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

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

[1] Von Oswald, et al. Transformers learn in-context by gradient descent. ICML 2023.

[2] Ahn, et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[3] Zhang, et al. In-context learning of a linear transformer block: benefits of the mlp component and one-step gd initialization. NeurIPS 2024.

[4] Frei, et al. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.

[5] Zhang, et al. Trained transformers learn linear models in-context. JMLR 2024.

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

[7] Chen, et al. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. NeurIPS 2024.

[8] Ren, et al. Learning and Transferring Sparse Contextual Bigrams with Linear Transformers. NeurIPS 2024.

[9] Huang, et al. Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.