Treinando Modelos de Linguagem Grandes para Raciocinar em um Espaço Latente Contínuo

Training Large Language Models to Reason in a Continuous Latent Space

December 9, 2024

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

cs.AI

Resumo

Os modelos de linguagem de grande escala (LLMs) estão restritos a raciocinar no "espaço da linguagem", onde normalmente expressam o processo de raciocínio com uma cadeia de pensamento (CoT) para resolver um problema de raciocínio complexo. No entanto, argumentamos que o espaço da linguagem nem sempre pode ser o mais adequado para o raciocínio. Por exemplo, a maioria dos tokens de palavras são principalmente para a coerência textual e não essenciais para o raciocínio, enquanto alguns tokens críticos exigem um planejamento complexo e representam enormes desafios para os LLMs. Para explorar o potencial do raciocínio dos LLMs em um espaço latente irrestrito em vez de usar linguagem natural, introduzimos um novo paradigma chamado Coconut (Cadeia de Pensamento Contínuo). Utilizamos o último estado oculto do LLM como representação do estado de raciocínio (denominado "pensamento contínuo"). Em vez de decodificá-lo em um token de palavra, o alimentamos de volta ao LLM como o embedding de entrada subsequente diretamente no espaço contínuo. Experimentos mostram que o Coconut pode aumentar efetivamente o desempenho do LLM em várias tarefas de raciocínio. Esse novo paradigma de raciocínio latente resulta em padrões de raciocínio avançados emergentes: o pensamento contínuo pode codificar múltiplas etapas de raciocínio seguintes alternativas, permitindo que o modelo execute uma busca em largura (BFS) para resolver o problema, em vez de se comprometer prematuramente com um único caminho determinístico como o CoT. O Coconut supera o CoT em certas tarefas de raciocínio lógico que exigem um retrocesso substancial durante o planejamento, com menos tokens de pensamento durante a inferência. Essas descobertas demonstram a promessa do raciocínio latente e oferecem insights valiosos para pesquisas futuras.

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.