Large language models (LLM) have been making significant strides in various domains, yet their ability to reason effectively remains a subject of ongoing research. Several studies have explored different prompting techniques to enhance the logical problem-solving capabilities of LLMs. The latest technique from researchers at Meta, named System 2 Attention (S2A), borrows concepts from psychological research. S2A meticulously revises the user’s prompt, eliminating any misleading or irrelevant information. By focusing solely on the task-relevant data, S2A allows LLMs to perform more accurately in question-answering and reasoning tasks. Initial experiments indicate a notable improvement in the performance of language models using S2A, which can be useful for applications that require reliable reasoning capabilities.
Understanding the Problem with Transformer-Based Models
The performance of LLMs in reasoning is a mixed bag. While certain prompt engineering techniques can enhance their performance, these models can falter when the prompt includes irrelevant or opinionated information. For instance, if a user’s question contains their personal guess or opinion, the model is prone to merely confirm or echo the user’s input rather than providing the correct answer. This behavior can be traced back to the training and attention mechanisms of transformers, the deep learning architecture employed in language models. Transformers are trained on next-token prediction, making them highly sensitive to contextual information. Consequently, if an entity is mentioned in a context, the model is likely to predict the same entity appearing later in the same context, leading to an overemphasis on repeated tokens in the model’s prompt.
“We posit that the underlying problem is inherent in the way the transformer itself is built, and in particular its attention mechanism,” the researchers write. “Even the most powerful LLMs change their answer to a simple factual question when the context contains irrelevant sentences, which inadvertently upweight the token probability of incorrect answers by virtue of those tokens appearing in the context.”
The System 2 Attention Technique
The researchers investigate a different approach to attention mechanisms that uses the LLM as a natural language reasoner. “We leverage the ability of LLMs to follow instructions, and prompt them to generate the context that they should pay attention to, such that it contains only relevant material that will not skew its reasoning,” they write. By using instruction-tuned LLMs to rewrite their context, the model can effectively eliminate irrelevant text. This enables the model to determine which parts of the input to concentrate on before generating a response. This process is referred to as System 2 Attention (S2A), a nod to the System 1 and System 2 thinking concept introduced by psychologist Daniel Kahneman in his book Thinking, Fast and Slow.
“System 1 thinking is rapid, intuitive, and automatic, often driven by instinct and prior learning. It’s the system we engage when performing tasks that don’t require much conscious thought, such as navigating familiar surroundings, engaging in casual conversation, or making quick, instinctive decisions. While efficient, System 1 thinking can lead to biases and errors in judgment due to its reliance on mental shortcuts and lack of deep analysis. In contrast, System 2 thinking is slow, deliberate, and analytical. It’s the system we engage when solving complex problems that require careful consideration and planning, or when learning new skills. System 2 demands more cognitive effort and can process information in a more logical and less biased manner than System 1. However, it’s also more mentally taxing and time-consuming.”
The S2A mechanism in LLMs can help mitigate the issues that arise when the model uses its standard attention mechanism to reason about its task. The researchers note, “S2A can produce more factual and less opinionated or sycophantic generations than standard attention-based LLMs.” The System 2 Attention technique is a straightforward two-step process. Initially, S2A modifies the original context, removing irrelevant parts that could negatively influence the output. Subsequently, the altered context is passed to the main LLM to generate its output. There are several ways to implement the initial step of S2A. Using instruction-tuned LLMs allows for control over the model’s attention focus based on the task or the model’s fine-tuning process.
The researchers implemented a function that sends a zero-shot prompt to the LLM, instructing it to perform the desired S2A task over the original prompt. For instance, they generate a prompt that instructs the LLM to regenerate the context, extracting the part that provides relevant context for a given query.
“In this implementation it specifically asks to generate an x′ [the modified prompt] that separates useful context from the query itself in order to clarify these reasoning steps for the model,” the researchers note.
In their paper, the researchers introduce several S2A variants. For instance, they find that for short contexts or strong LLMs, partitioning the context and question isn’t necessary. An S2A prompt that simply asks for a non-partitioned rewrite of the query should suffice. Another variant keeps the original prompt and adds the S2A-generated query to it, so both the original context and its reinterpretation are available for the model to access.
The researchers tested S2A on a variety of problems including question answering, long-form reasoning, and math word problems that either contain irrelevant information, misleading facts, or opinionated sentences. The S2A system must answer the question objectively and remove irrelevant information to guide the model toward using the data points that will provide the most accurate answer. Their experiments demonstrate that S2A is “not easily swayed by opinion,” and it enables LLMs to perform nearly on par with situations where it is provided with a clean prompt free from distracting or opinionated information. The experiments also show that LLMs equipped with S2A are better at maintaining their objectivity in long-form generation tasks. Despite the impressive results of S2A, the researchers acknowledge that it doesn’t always succeed and models “will still sometimes be affected by spurious correlations.” S2A also increases the costs of LLM generation because it adds steps to the generation task and requires the extraction of the contextual information from the original prompt. These are some of the areas where the technique can be improved in the future to become a good addition to the toolbox of reasoning techniques for LLM applications.
