Instruction Following by Principled Boosting Attention of Large Language Models

1. Problem Statement

Large Language Models (LLMs) are increasingly deployed as general-purpose assistants where adherence to system instructions (e.g., safety constraints, persona definitions, refusal boundaries) is critical. However, in practice, models often fail to follow these instructions, particularly when:

• Context Conflict: User-provided context or long prompts conflict with the system instructions.
• Brittleness: Models satisfy instructions in a way that degrades task precision or ignores necessary context.
• Existing Limitations: Current inference-time interventions have trade-offs:
  • Latent Steering: Modifying hidden states often leads to "fluency collapse" (repetitive, nonsensical output) and inconsistent performance across tasks.
  • Prior Attention Steering (PASTA, SpotLight): While effective, these methods either require computationally expensive head selection (PASTA) or enforce rigid, state-dependent attention targets that can suppress benign, task-relevant context, leading to "instruction over-focus" and relevance degradation (SpotLight).

The core challenge is to strengthen instruction adherence without retraining the model, while preserving the model's ability to utilize task-relevant context and maintaining generation quality.

2. Methodology: INSTABOOST

The authors propose INSTABOOST (Instruction Attention Boosting), a simple, efficient inference-time intervention.

Theoretical Foundation

The work is grounded in a unifying theory based on the Logicbreaks abstraction, which models instruction following as a rule-based competition:

• Instruction Rules ($\Gamma$): Derived from system prompts and explicit instructions.
• Competing Rules ($\Delta$): Derived from user queries and additional context.
• Mechanism: Attention mediates this competition. The paper proves that adding a constant additive bias ($B$) to the pre-softmax attention logits for instruction tokens exponentially increases the influence of instruction rules relative to competing rules.
• Trade-off: While boosting helps instruction adherence, excessive boosting can suppress "benign" competing rules (necessary context), causing relevance loss.

The INSTABOOST Algorithm

Unlike previous methods that select specific heads or use dynamic, state-dependent biases, INSTABOOST applies a uniform, constant additive bias across all layers and all attention heads.

• Operation: For an input sequence $x' = p \oplus x$ (where $p$ is the instruction prefix of length $K$ and $x$ is the query), the method modifies the pre-softmax attention logits $S_{ij}$:
  $$T_B(S_{ij}) = \begin{cases} S_{ij} + B & \text{if } 0 \le j < K \quad (\text{Instruction tokens}) \\ S_{ij} & \text{if } K \le j < N \quad (\text{Context tokens}) \end{cases}$$
• Implementation: This is equivalent to multiplying the softmax numerator for instruction keys by a factor $M = e^B$.
• Key Advantage: It requires no offline profiling to select heads (unlike PASTA) and does not enforce a rigid minimum attention mass that might crush context (unlike SpotLight). It acts as a single, interpretable knob ($B$ or $M$) to navigate the robustness-relevance tradeoff.

3. Key Contributions

1. Unifying Theory: Formalized instruction following as a competition between instruction rules and context-derived rules, mediated by attention. The authors proved that additive attention boosting systematically tilts this competition in favor of instructions, making it exponentially harder for context to override them.
2. Theoretical Characterization of Failure Modes: The theory predicts two regimes:
   • Robustness: Boosting increases the "subversion budget" required for competing rules to override instructions.
   • Benign Suppression: Excessive boosting can downweight necessary context, leading to relevance loss.
3. INSTABOOST Proposal: A computationally efficient, head-agnostic intervention that applies a constant bias to instruction keys.
4. Comprehensive Evaluation: Evaluated across 15 diverse tasks (Emotion, AI Persona, Jailbreaking, Toxicity, Truthfulness, QA) on Meta-Llama-3-8B-Instruct and Gemma-7b-it, comparing against prompting, latent steering, and prior attention methods.

4. Experimental Results

The evaluation demonstrates that INSTABOOST achieves a superior steering-quality tradeoff.

• Performance vs. Baselines:
  • INSTABOOST matches or outperforms all baselines (including latent steering and PASTA/SpotLight) across 15 tasks.
  • Jailbreaking: INSTABOOST is the only method that significantly improves performance on jailbreaking tasks (e.g., AdvBench, JailbreakBench) while maintaining high fluency, whereas latent methods often fail or degrade quality.
  • Consistency: Unlike latent steering methods, which show high variance in performance across different task types (e.g., good on Emotion, poor on Jailbreak), INSTABOOST is consistently effective.
• Side Effects (Fluency & Relevance):
  • Fluency: Latent steering methods suffer from severe fluency collapse as steering strength increases. INSTABOOST maintains high fluency (scores $\approx 1.5-2.0$) even at high strengths.
  • Relevance: SpotLight suffers from "instruction over-focus," where the model ignores the user's query to satisfy the instruction (relevance drops to 0). INSTABOOST preserves relevance, successfully answering user queries while adhering to constraints.
• Ablation Studies: Increasing the bias $B$ makes failure modes (fact amnesia, rule suppression) harder to induce in a controlled "learned reasoner" setting, validating the theoretical predictions.

5. Significance and Conclusion

This work provides a principled, mechanistic understanding of how attention manipulation influences instruction following. By framing the problem as a rule-based competition, the authors demonstrate that a simple, uniform additive bias is sufficient to achieve robust instruction adherence without the computational overhead of head selection or the quality degradation of latent steering.

Key Takeaways:

• Simplicity: INSTABOOST is a "drop-in" intervention requiring only a single hyperparameter ($B$ or $M$) applied globally.
• Safety & Reliability: It offers a reliable method to enforce safety constraints and system instructions without compromising the model's ability to answer questions or generate fluent text.
• Future Direction: The paper suggests that guiding attention is a more effective interface for controlling LLM behavior than manipulating hidden states, offering a promising path toward safer and more controllable AI systems.

The code for INSTABOOST is intended to be open-sourced, facilitating immediate adoption in inference pipelines.