Instruction Following by Principled Boosting Attention of Large Language Models

Abstract

Large language models' behavior is often shaped by instructions such as system prompts, refusal boundaries, privacy constraints, and tool-use rules that must hold at inference time. Yet in practice these constraints can be violated under long contexts or when user-provided context conflicts with them, creating reliability and safety risks. This motivates inference-time interventions that strengthen instruction influence without retraining. One such intervention is attention steering, which biases attention toward instruction tokens. In this work, we present a unifying theory for attention steering methods by formalizing instruction following as rule-based competition between instruction rules and context-derived rules, with attention mediating which rules dominate. We prove that boosting attention to instruction tokens tilts this competition, making it harder for context to override instruction-following. However, excessive boosting can suppress task-relevant context that should be incorporated alongside the instruction. Guided by this theory, we propose Instruction Attention Boosting (InstABoost), a simple intervention that applies a constant additive bias to instruction-key attention logits across all layers and heads. We evaluate InstABoost against prompting, latent steering, and prior attention steering methods across 15 tasks. InstABoost matches or outperforms all baselines while avoiding the fluency collapse of latent methods and the instruction over-focus of prior attention methods, achieving a stronger steering-quality tradeoff.

Figures (20)

• Figure 1: Illustration of InstABoost. We prepend an instruction prefix to the user query and run the transformer as usual, except that at each attention head we add a constant bias $B$ to the pre-softmax attention logits for keys corresponding to instruction tokens. This increases the attention mass allocated to instruction tokens during generation.
• Figure 2: Attention allocation while processing the last input token, averaged across heads. We show the base model’s attention (left) and the change in attention induced by each method relative to base ($\Delta$; right three panels). All methods increase attention on the instruction (above the dashed line), but differ in how the boost is applied: SpotLight uses a query-dependent bias to enforce a minimum instruction attention mass; PASTA downweights non-instruction keys on a selected subset of heads; InstABoost applies a fixed bias to the instruction span across all heads.
• Figure 3: Increasing instruction-attention bias makes the failure modes harder to induce in a learned reasoner. Using the same learned-reasoner setting and suffix constructions as Logicbreaks, we apply InstABoost and vary only the bias $B$. Larger $B$ requires more suffix repetitions to induce fact amnesia and rule suppression, and slows the variance collapse for state coercion. The clean reasoning accuracy decreases only from $99.6\%$ at $B=0$ to $99.3\%$ at $B=2.5$.
• Figure 4: InstABoost outperforms or is competitive with all evaluated interventions. For each task, we show the steering success of the model without intervention (brown), the best-performing latent steering method on each task (purple), the instruction-only intervention (red), the attention-based methods PASTA (green) and SpotLight (orange), and InstABoost (blue). Error bars show a standard deviation above and below the mean, computed by bootstrapping. Full results are in Appendix \ref{['app:additional-results']}.
• Figure 5: The effect of steering strength on steering success, fluency, and relevance for the Sadness (top) and Power QA (bottom) tasks. Latent methods show a clear trade-off, where increasing strength improves steering success but collapses fluency. While attention-based methods preserve fluency, SpotLight severely harms relevance. InstABoost is unique in its ability to achieve high steering success while maintaining both high fluency and relevance.

Theorems & Definitions (24)

• Proposition 3.1: Update decomposition (informal)
• Theorem 3.2: Subversion-budget inflation (informal)
• Theorem 3.3: Benign correctness (informal)
• Definition A.3: Additive attention boosting
• Proposition A.5: Biased softmax concentration
• proof
• Proposition A.6: Update decomposition
• proof
• Theorem A.7: Monotonicity robustness (fact-amnesia)
• proof