Thinking LLMs: General Instruction Following with Thought Generation (2410.10630v1)

Abstract: LLMs are typically trained to answer user questions or follow instructions similarly to how human experts respond. However, in the standard alignment framework they lack the basic ability of explicit thinking before answering. Thinking is important for complex questions that require reasoning and planning -- but can be applied to any task. We propose a training method for equipping existing LLMs with such thinking abilities for general instruction following without use of additional human data. We achieve this by an iterative search and optimization procedure that explores the space of possible thought generations, allowing the model to learn how to think without direct supervision. For each instruction, the thought candidates are scored using a judge model to evaluate their responses only, and then optimized via preference optimization. We show that this procedure leads to superior performance on AlpacaEval and Arena-Hard, and shows gains from thinking on non-reasoning categories such as marketing, health and general knowledge, in addition to more traditional reasoning & problem-solving tasks.

• The paper introduces Thought Preference Optimization (TPO) to enable LLMs to generate internal thought processes before responding.
• The paper utilizes a reinforcement learning from AI feedback framework to iteratively optimize these thought processes without relying on human-supervised data.
• The paper shows that TPO-trained models achieve superior benchmark performance, expanding LLM applicability beyond conventional reasoning tasks.

Summary of "Thinking LLMs: General Instruction Following with Thought Generation"

This paper addresses the challenge of equipping LLMs with the capability to "think" before generating a response, a feature absent in standard alignment frameworks. The authors propose a method termed Thought Preference Optimization (TPO), which allows LLMs to engage in thought processes in natural language for general instruction-following tasks.

Motivation and Methodology

The premise of the work is that current LLMs, based on the Transformer architecture, have a static compute budget per token, regardless of instruction complexity. This limitation inhibits the generation of responses requiring extensive reasoning or planning. Inspired by human cognitive processes, the paper introduces TPO, which enhances the model's ability to internally process thoughts before responding. This method does not rely on additional human-supervised data, which is often scarce, particularly for internal thought processes.

The TPO process involves the LLM generating a sequence that comprises both thought and response segments. Training begins with a generic or specific prompt instructing the model to articulate its thought process. The model's output is then evaluated using a judge model that only considers the response, optimizing thoughts to improve subsequent responses. This is achieved through iterative training in a Reinforcement Learning from AI Feedback (RLAIF) framework without explicit thought supervision.

Experimental Results

The proposed approach demonstrated superior performance compared to traditional direct-response LLMs across several benchmarks. On datasets like AlpacaEval and Arena-Hard, TPO-trained models showed improved win rates, suggesting that the thought process aids in tackling a broader range of tasks beyond math and reasoning, as traditionally associated with Chain-of-Thought (CoT) prompting.

The experimental setup involved multiple iterations of training with diverse data sources and judge models, such as ArmoRM and Self-Taught Evaluator (STE). Fine-grained evaluations revealed that TPO can improve performance in categories once considered non-reasoning, such as language translation, marketing, and health, broadening the applicability of thinking models to general instruction following.

Implications and Future Directions

The advancement outlined in this paper marks a significant step in harnessing the benefits of thought processes in LLMs, potentially leading to more versatile AI applications. The research suggests that internal thought mechanisms could be optimized to enhance the quality of responses without direct oversight of the thought content itself.

Future research directions could explore diverse thought prompts and their effects on various task categories, aiming to further refine the process. Moreover, applying this approach to larger-scale models could provide insights into scalability and effectiveness. Addressing the challenges of optimizing thoughts for math-intensive tasks might require integrating more specialized training data and evaluative mechanisms.

In conclusion, this paper adds a promising layer to the development of thinking AI systems, laying groundwork for future innovations in enhancing LLM capabilities in complex, multi-domain interactions without the need for explicit human-generated supervisory signals.