Research map / Architectures and efficiency / Attention and transformers
On the Robustness of Transformers against Context Hijacking for Linear Classification
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Theoretically analyzes why deeper linear transformers are more robust to context hijacking in in-context linear classification, showing they implement optimized multi-step gradient descent whose robustness improves with depth and training context length.
Problem
Context hijacking: factually correct but misleading in-context examples can flip LLM predictions; the paper seeks a theoretical understanding of how transformer depth, training context length, and number of hijacking examples affect robustness.
Approach
Models multi-layer linear attention transformers as multi-step preconditioned gradient descent on in-context examples, derives the optimal learning rate and initialization (non-zero mean prior beta*), and uses this to bound classification error under hijacked contexts (boundary-projected, opposite-label examples) as a function of depth L, training context n, and hijack count N.
Key results
- Theorem gives optimal learning rate alpha = Theta~(1/(nL)) and initialization w0 = c*beta* for multi-step GD in linear transformers.
- Classification error bound under hijacking scales with (1 - Theta~(Nd/(nL)))^L, so deeper models and longer training contexts are more robust.
- Synthetic experiments show linear transformers with more layers retain higher test accuracy as hijacking examples increase.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase the number of hijacking context examples | linear transformers · synthetic test set | classification accuracy | Hurt |
| increase transformer depth | linear transformers · synthetic test set | classification accuracy | Helped |
| increase training context length | linear transformers · synthetic test set | classification accuracy under hijacking | Helped |
| increase training context length | linear transformers · synthetic training set | classification accuracy | Helped |
| increase gradient descent steps while using the corresponding optimal learning rate | multi-step gradient descent · hijacking test distribution | optimization performance | Helped |
| increase the number of gradient descent steps | optimal multi-step gradient descent · training-context-length sweep | optimal learning rate | Hurt |
| increase training context length | optimal multi-step gradient descent · gradient descent experiment | optimal learning rate | Hurt |
Limitations
- Analysis restricted to linear attention-only single-head transformers and linear classification.
- Hijacking examples are idealized (placed on the decision boundary, eta to 0).
Experimental setup
Datasets: synthetic Gaussian in-context linear classification data
Metrics: classification error test robust accuracy optimal learning rate
Builds on: linear transformers implement gradient descent context hijacking (Zhang et al.) in-context learning with non-zero mean prior
Closest papers by meaning
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.87)
- Unlabeled Data Can Provably Enhance In-Context Learning of Transformers (NeurIPS 2025 · similarity 0.87)
- Context-Aware Hierarchical Learning: A Two-Step Paradigm towards Safer LLMs (NeurIPS 2025 · similarity 0.86)
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.86)
- Core Context Aware Transformers for Long Context Language Modeling (ICML 2025 · similarity 0.86)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.