Instruction Following by Principled Attention Boosting of Large Language Models

Vitoria Guardieiro, Avishree Khare, Adam Stein, Eric Wong

Keywords: instruction following, attention steering, inference-time intervention, attention steering theory, LLM safety

TL;DR: We develop a theory for how attention steering improves instruction following, then use it to propose InstABoost, a simple inference-time method that improves instruction following without hurting generation quality.

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 theoretical formalization of instruction following as rule-based competition between instruction rules and context-derived rules, with attention mediating which rules dominate, unifying existing attention-steering methods. 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 methods across 15 steering 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.

Submission Number: 190