System 2 Attention (is something you might need too)

Abstract: Soft attention in Transformer-based LLMs is susceptible to incorporating irrelevant information from the context into its latent representations, which adversely affects next token generations. To help rectify these issues, we introduce System 2 Attention (S2A), which leverages the ability of LLMs to reason in natural language and follow instructions in order to decide what to attend to. S2A regenerates the input context to only include the relevant portions, before attending to the regenerated context to elicit the final response. In experiments, S2A outperforms standard attention-based LLMs on three tasks containing opinion or irrelevant information, QA, math word problems and longform generation, where S2A increases factuality and objectivity, and decreases sycophancy.

• The paper proposes a novel method that improves LLM performance by regenerating context to filter out irrelevant information.
• Experiments demonstrate that S2A nearly restores factual accuracy in TriviaQA tasks, achieving results close to oracle performance.
• The approach enhances objectivity in longform generation and math problem solving, offering a scalable solution for reducing distractor impact.

System 2 Attention: Informed Context Regeneration for Enhanced LLM Performance

The paper "System 2 Attention (S2A)" by Jason Weston and Sainbayar Sukhbaatar presents a novel approach to improve the performance of Transformer-based LLMs by addressing the susceptibility of these models to spurious correlations and irrelevant context information. The core proposition of S2A is to leverage the inherent reasoning capabilities of LLMs to regenerate a context that emphasizes relevant information, thereby mitigating common issues such as decreased factuality and increased sycophancy in generated responses.

Introduction

LLMs, while highly capable, exhibit weaknesses in handling irrelevant context which can adversely affect their performance, particularly in tasks requiring significant reasoning. The authors introduce S2A to address this by making the LLMs act as natural language reasoners. By prompting the model to regenerate a context with only pertinent information, S2A aims to align the model's attention more closely with the task requirements, thereby improving response quality.

System 2 Attention Mechanism

Motivation

Soft attention mechanisms in LLMs assign probability weights to portions of the input context, including irrelevant segments, which can lead to erroneous outputs. Examples include instances of sycophancy where the model agrees with opinions embedded in prompts. The motivation for S2A stems from these observations, emphasizing the need for a deliberate mechanism akin to system 2 reasoning in humans, which focuses effortfully on relevant aspects to avoid errors seen in automatic (system 1) operations.

Implementation

S2A operates through a two-step process:

1. Context Regeneration: The LLM is prompted to produce a refined context $x'$ from the original context $x$, excluding irrelevant parts. This is achieved using an instruction-tuned LLM with prompts designed to differentiate between useful context and the primary query.
2. Response Generation: The model then generates the final response based on the regenerated context $x'$.

The implementation leverages zero-shot prompting techniques to instruct the LLM to extract only relevant context segments, structuring prompts to enhance clarity and reduce potential bias.

Experiments and Results

The efficacy of S2A was tested across three diverse tasks: factual question answering (TriviaQA with SycophancyEval), longform generation evaluation, and math word problems (GSM-IC).

Factual Question Answering

In the modified TriviaQA task, where questions were augmented with opinionated prompts, S2A improved the factual accuracy of responses. Standard LLMs experienced a drop in accuracy from 82.0% (oracle) to 62.8%. S2A significantly mitigated this, achieving 80.3%, close to oracle performance. This demonstrates S2A's ability to reduce the impact of spurious correlations, ensuring the model's responses remain factual.

Longform Generation

For longform generation tasks, S2A improved objectivity without compromising overall response quality. While the baseline and oracle prompts produced similar quality scores (4.6-4.7), S2A outperformed in objectivity, scoring 3.82 versus 3.0 for oracle prompts and 2.23 for baseline.

Math Word Problems

In math problem tasks with embedded irrelevant sentences, S2A showed a substantial performance gain. Baseline accuracy significantly dropped due to distractors, but S2A effectively improved performance, especially with in-topic distractors, providing accuracy gains of up to 10-12%.

Discussion

The S2A approach represents a significant enhancement for LLMs by tackling fundamental issues of context irrelevance and spurious correlations. By regenerating context with a deliberate attention mechanism, S2A aligns model outputs more closely with task objectives, reducing errors and improving overall performance.

Implications and Future Directions

Practically, S2A offers a scalable approach applicable across various domains requiring precise and reliable language understanding. Theoretically, it presents a blueprint for integrating high-level reasoning within LLMs, potentially informing the development of more robust attention mechanisms. Future research may explore optimizing S2A further through fine-tuning methods, reinforcement learning, or more sophisticated prompting approaches.

Conclusion

System 2 Attention (S2A) provides a mechanism to leverage LLMs' reasoning capabilities to refine context and improve output quality. By focusing on relevant information and mitigating spurious correlations, S2A enhances the factuality and objectivity of model responses, representing an important step forward in the development of more intelligent and reliable AI systems.