Entraînement de grands modèles de langage pour raisonner dans un espace latent continu

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

Résumé

Les grands modèles de langage (LLMs) sont limités à raisonner dans l'espace de "langage", où ils expriment généralement le processus de raisonnement avec une chaîne de pensée (CoT) pour résoudre un problème de raisonnement complexe. Cependant, nous soutenons que l'espace de langage n'est pas toujours optimal pour le raisonnement. Par exemple, la plupart des jetons de mots sont principalement destinés à la cohérence textuelle et ne sont pas essentiels au raisonnement, tandis que certains jetons critiques nécessitent une planification complexe et posent d'énormes défis aux LLMs. Pour explorer le potentiel du raisonnement des LLMs dans un espace latent non restreint au lieu d'utiliser le langage naturel, nous introduisons un nouveau paradigme, Coconut (Chaîne de Pensée Continue). Nous utilisons le dernier état caché du LLM comme représentation de l'état de raisonnement (appelé "pensée continue"). Au lieu de le décoder en un jeton de mot, nous le renvoyons au LLM en tant qu'incorporation d'entrée ultérieure directement dans l'espace continu. Les expériences montrent que Coconut peut augmenter efficacement le LLM sur plusieurs tâches de raisonnement. Ce nouveau paradigme de raisonnement latent conduit à des schémas de raisonnement avancés émergents : la pensée continue peut encoder plusieurs étapes de raisonnement suivantes alternatives, permettant au modèle d'effectuer une recherche en largeur (BFS) pour résoudre le problème, plutôt que de s'engager prématurément dans un seul chemin déterministe comme CoT. Coconut surpasse CoT dans certaines tâches de raisonnement logique qui nécessitent un retour en arrière substantiel lors de la planification, avec moins de jetons de réflexion pendant l'inférence. Ces résultats démontrent la promesse du raisonnement latent et offrent des perspectives précieuses pour la recherche future.

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.