Het trainen van grote taalmodellen om te redeneren in een continue latente ruimte

Training Large Language Models to Reason in a Continuous Latent Space

December 9, 2024

Auteurs: Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian

cs.AI

Samenvatting

Grote taalmodellen (LLM's) zijn beperkt tot redeneren in de "taalruimte", waar ze doorgaans het redeneerproces uitdrukken met een keten van gedachten (CoT) om een complex redeneerprobleem op te lossen. We betogen echter dat de taalruimte niet altijd optimaal is voor redeneren. Zo zijn de meeste woordtokens voornamelijk voor tekstuele samenhang en niet essentieel voor redeneren, terwijl sommige kritieke tokens complexe planning vereisen en enorme uitdagingen vormen voor LLM's. Om het potentieel van LLM-redeneren te verkennen in een onbeperkte latente ruimte in plaats van natuurlijke taal te gebruiken, introduceren we een nieuw paradigma Coconut (Keten van Continue Gedachte). We gebruiken de laatste verborgen toestand van de LLM als een representatie van de redeneertoestand (genaamd "continue gedachte"). In plaats van dit te decoderen tot een woordtoken, voeren we het direct terug naar de LLM als de daaropvolgende invoer-embedding in de continue ruimte. Experimenten tonen aan dat Coconut effectief de LLM kan versterken bij verschillende redeneertaken. Dit nieuwe latente redeneerparadigma leidt tot opkomende geavanceerde redeneerpatronen: de continue gedachte kan meerdere alternatieve volgende redeneerstappen coderen, waardoor het model een breedte-eerst zoeken (BFS) kan uitvoeren om het probleem op te lossen, in plaats van zich vroegtijdig te committeren aan een enkel deterministisch pad zoals CoT. Coconut presteert beter dan CoT bij bepaalde logische redeneertaken die aanzienlijke terugtracking vereisen tijdens de planning, met minder denktokens tijdens inferentie. Deze bevindingen tonen de belofte van latente redenering en bieden waardevolle inzichten voor toekomstig onderzoek.

English

Large language models (LLMs) are restricted to reason in the "language space", where they typically express the reasoning process with a chain-of-thought (CoT) to solve a complex reasoning problem. However, we argue that language space may not always be optimal for reasoning. For example, most word tokens are primarily for textual coherence and not essential for reasoning, while some critical tokens require complex planning and pose huge challenges to LLMs. To explore the potential of LLM reasoning in an unrestricted latent space instead of using natural language, we introduce a new paradigm Coconut (Chain of Continuous Thought). We utilize the last hidden state of the LLM as a representation of the reasoning state (termed "continuous thought"). Rather than decoding this into a word token, we feed it back to the LLM as the subsequent input embedding directly in the continuous space. Experiments show that Coconut can effectively augment the LLM on several reasoning tasks. This novel latent reasoning paradigm leads to emergent advanced reasoning patterns: the continuous thought can encode multiple alternative next reasoning steps, allowing the model to perform a breadth-first search (BFS) to solve the problem, rather than prematurely committing to a single deterministic path like CoT. Coconut outperforms CoT in certain logical reasoning tasks that require substantial backtracking during planning, with fewer thinking tokens during inference. These findings demonstrate the promise of latent reasoning and offer valuable insights for future research.