Title: On the Robustness of Transformers against Context Hijacking for Linear Classification
Venue: NeurIPS 2025
Accepted: yes
Keywords: in-context learning

transformers

robustness

deep learning theory

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\n1 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\nIt 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\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

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

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

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.

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
* Can the techniques be used for other, similar problems?

LIMITATIONS:
None
TITLE: Official Review
CLARITY: 2
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-07-03

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

STRENGTHS AND WEAKNESSES:
### Strength

1) The work identifies and rigorously formalizes the problem of context hijacking, a subtle yet impactful threat to LLM robustness.

2) The paper offers a clear and elegant connection between transformer depth and robustness via optimal multi-step gradient descent modeling.

3) Numerical simulations back the theoretical predictions and explain why deeper transformers resist hijacking more effectively.

### Weakness

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

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

QUESTIONS:
N/A

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-04