Inside System 2 Attention: Meta AI New Method to Improve Reasoning in LLMs
The method is inspired in behavioral psychology and improves reasoning across different LLM problems.
I recently started an AI-focused educational newsletter, that already has over 160,000 subscribers. TheSequence is a no-BS (meaning no hype, no news, etc) ML-oriented newsletter that takes 5 minutes to read. The goal is to keep you up to date with machine learning projects, research papers, and concepts. Please give it a try by subscribing below:
Reasoning is shaping up to be the next major area of focus for Large Language Models (LLMs). Despite their advanced capabilities, most LLMs often stumble over simple errors, showing limitations in their reasoning. These models can be misled by irrelevant details in their context or influenced by the biases in the input prompts. This latter tendency, known as sycophancy, results in the model agreeing with the input, regardless of accuracy. Various efforts have been made to address these shortcomings, including increasing supervised training data or applying reinforcement learning methods. In a recent research, Meta AI suggests that the root of the problem lies in the fundamental design of the transformer architecture used in these models, particularly the attention mechanism. The research drives inspiration from Daniel Kahneman and Amos Tversky research on behavioral psychology that was brilliantly captured in the book Thinking Fast and Slow. Not surprisingly, they coined the term System 2 Attention.
The Problem
LLMs excel in reasoning and knowledge accumulation through their extensive pre-training. They are designed to focus intensely on the current context for predicting the next word. For instance, if a particular entity appears in a text, the model anticipates its recurrence. Transformer-based LLMs, with their soft-attention mechanism, are adept at identifying similar words and concepts within their context. While this enhances their prediction accuracy, it also leaves them vulnerable to misleading correlations in the context they analyze.
S2A
Meta AI proposes using LLMs as natural language reasoners to focus attention. This method, named System 2 Attention (S2A), involves prompting LLMs to create a context stripped of irrelevant information that could distort reasoning. The concept draws inspiration from human cognitive processes, where ‘System 2’ represents conscious, effortful mental activity, particularly employed when error-prone ‘System 1’ reasoning isn’t sufficient. S2A aims to replicate this by directing the LLM’s reasoning capabilities to overcome the flaws inherent in the transformer’s soft attention mechanism.
In a typical scenario, an LLM is presented with a context (x) and tasked with generating a high-quality output (y). S2A modifies this process through a two-step method. First, S2A reformulates the given context (x) into a refined version (x′) by removing elements that might negatively impact the output. This is represented as x′ ∼ S2A(x). Then, the LLM generates the final response (y) using this revised context (x′), instead of the original, symbolized as y ∼ LLM(x′).
S2A encompasses a range of techniques for implementing this first step. Meta AI’s specific approach takes advantage of instruction-tuned LLMs already skilled in reasoning and generation tasks. They use these capabilities to instruct the LLM through prompting to perform the S2A task. In practice, this involves creating a zero-shot prompt that guides the LLM to apply System 2 Attention to the given context, denoted as S2A(x) = LLM(PS2A(x)), with PS2A being the function generating the prompt.
Meta AI has chosen LLaMA-2–70B-chat as their primary model for evaluation. They assess its performance in two distinct scenarios:
- Baseline Setting: In this approach, the input prompt from the dataset is directly supplied to the model, which then generates a response in a zero-shot manner. This method is straightforward but has a drawback. The model’s outputs are susceptible to being influenced by any biases, opinions, or irrelevant details present in the input prompt.
- Oracle Prompt Setting: Here, the model is given a cleaned-up version of the prompt, stripped of any superfluous opinions or irrelevant sentences. The model then responds in a zero-shot fashion to this refined prompt. This setting is used to gauge the ideal performance level of the model if it were able to perfectly filter out irrelevant information from the input. This scenario acts as a hypothetical best-case benchmark for the model’s capabilities.
The Results
Meta AI conducts research to assess the effectiveness of System 2 Attention (S2A) in three distinct scenarios that hghlights LLM reasoning abilities.
- Factual Question Answering: The team utilizes a version of TriviaQA from SycophancyEval1, designed for straightforward fact-based queries, but with additional opinions included in the prompts. They experiment with prompts that either suggest a correct or incorrect answer, or refute the correct answer. These additions are known to influence standard models’ responses. The findings show that S2A significantly enhances accuracy in such opinion-influenced scenarios, closely matching the performance of unopinionated, or ‘oracle’, prompts.
2. Longform Generation of Arguments: For this, the team uses argument prompts from SycophancyEval, which include opinions in the context of provided arguments. These prompts come with comments that express liking, disliking, authorship, or non-authorship of the argument. Standard models tend to skew their responses based on these sentiments. However, S2A shows a marked improvement in maintaining objectivity in the generated arguments, even surpassing the objectivity level of the oracle prompts.
3. Math Word Problem Solving: Here, Meta AI tests S2A on the GSM-IC task, which involves math word problems from GSM8K, but with added irrelevant sentences. These distracting sentences can significantly decrease the accuracy of LLMs. Meta AI experiments with two types of distractors: random and topic-related. This setup aims to evaluate how effectively S2A can filter out irrelevant information and focus on the essential aspects of the problem to provide accurate solutions.
S2A represents an important milestone in the evolution of reasoning methods in LLMs. The method closely ressembles human reasoning and avoid distractions. We should expect S2A to be an important baseline in reasoning research in recent months.
