Don't Take Things Out of Context: Attention Intervention for Enhancing Chain-of-Thought Reasoning in Large Language Models

TL;DR

This work identifies a fragility in few-shot Chain-of-Thought prompting: isolated tokens in demonstrations can distract LLMs by attracting attention before sufficient information aggregation. It introduces Few-shot Attention Intervention (FAI), a lightweight method that analyzes per-layer attention patterns to locate tokens with limited aggregation but strong influence on predictions, and then suppresses their flow by zeroing their attention to the output at the relevant layer. Through extensive experiments across GSM8K, AQuA, CSQA, Date Understanding, Sport Understanding, and Last Letter tasks, FAI yields consistent improvements (notably +5.91% on AQuA) while preserving the beneficial CoT effects and incurring minimal overhead. The findings advance robust CoT reasoning by making it less sensitive to distractor tokens, with practical implications for in-context learning pipelines and future scaling to larger models.

Abstract

Few-shot Chain-of-Thought (CoT) significantly enhances the reasoning capabilities of large language models (LLMs), functioning as a whole to guide these models in generating reasoning steps toward final answers. However, we observe that isolated segments, words, or tokens within CoT demonstrations can unexpectedly disrupt the generation process of LLMs. The model may overly concentrate on certain local information present in the demonstration, introducing irrelevant noise into the reasoning process and potentially leading to incorrect answers. In this paper, we investigate the underlying mechanism of CoT through dynamically tracing and manipulating the inner workings of LLMs at each output step, which demonstrates that tokens exhibiting specific attention characteristics are more likely to induce the model to take things out of context; these tokens directly attend to the hidden states tied with prediction, without substantial integration of non-local information. Building upon these insights, we propose a Few-shot Attention Intervention method (FAI) that dynamically analyzes the attention patterns of demonstrations to accurately identify these tokens and subsequently make targeted adjustments to the attention weights to effectively suppress their distracting effect on LLMs. Comprehensive experiments across multiple benchmarks demonstrate consistent improvements over baseline methods, with a remarkable 5.91% improvement on the AQuA dataset, further highlighting the effectiveness of FAI.

Figures (10)

• Figure 1: An example of a Few-shot Chain-of-Thought demonstration distracting the reasoning of LLMs. Questions are collected from the GSM8K cobbe2021training dataset, and the responses are generated by Llama-3-8B-Instruct llama3modelcard.
• Figure 2: Attention salience analysis example. For each sample, the upper part of the figure shows the salience scores of the demonstration tokens, the tokens in the question, and the generated tokens for the current prediction position; darker colors indicate stronger salience. Subsequently, we select one token that has a significant impact on the output, and the lower part of the figure displays the salience of the preceding tokens for that selected token. In each case, we choose a layer with pronounced phenomena to demonstrate the characteristics of attention salience more intuitively. (a)(c): Cases correctly answered. (b)(d): Cases with wrong responses.
• Figure 3: Overview of the proposed FAI method. In each layer, FAI identifies the positions that require intervention by analyzing the attention matrix from the demonstration. It then applies these interventions to the attention matrix of the subsequent layer, relative to the positions of the output tokens.
• Figure 4: (a) Comparison of the Overall Accuracy on $GSM_{bad}$ and $GSM_{good}$. (b) Comparison of the Rate of Answer Following Rationale on $GSM_{bad}$ and $GSM_{good}$.
• Figure 5: (a) The overall accuracy of Llama-3-8B-Instruct on the test set of GSM8K with 45 various one-shot CoT demonstration randomly selected from the training set of GSM8K using different seeds. (b) The IoU of correctly answered questions between every groups.