Token Sortido: Misturando Tokens Latentes e de Texto para Melhorar o Raciocínio do Modelo de Linguagem

Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning

February 5, 2025

Autores: DiJia Su, Hanlin Zhu, Yingchen Xu, Jiantao Jiao, Yuandong Tian, Qinqing Zheng

cs.AI

Resumo

Grandes Modelos de Linguagem (LLMs) destacam-se em raciocínio e planejamento quando treinados em dados de cadeia de pensamento (CoT), nos quais o processo de pensamento passo a passo é explicitamente delineado por tokens de texto. No entanto, isso resulta em entradas extensas em que muitas palavras sustentam a coerência textual em vez das informações de raciocínio central, e processar essas entradas consome recursos computacionais substanciais. Neste trabalho, propomos uma representação híbrida do processo de raciocínio, onde abstraímos parcialmente os passos iniciais de raciocínio usando tokens discretos latentes gerados por VQ-VAE, reduzindo significativamente o comprimento dos rastros de raciocínio. Exploramos o uso de abstrações de rastros latentes em dois cenários: 1) treinando o modelo do zero para o problema do Labirinto de Encontrar Chaves, 2) ajustando finamente LLMs nesses dados híbridos com um vocabulário estendido incluindo tokens latentes não vistos, tanto para problemas de raciocínio lógico quanto matemático. Para facilitar a aprendizagem eficaz, introduzimos um procedimento de treinamento simples que mistura aleatoriamente tokens latentes e de texto, o que permite uma adaptação rápida a novos tokens latentes. Nossa abordagem supera consistentemente os métodos de referência em vários benchmarks.

English

Large Language Models (LLMs) excel at reasoning and planning when trained on chainof-thought (CoT) data, where the step-by-step thought process is explicitly outlined by text tokens. However, this results in lengthy inputs where many words support textual coherence rather than core reasoning information, and processing these inputs consumes substantial computation resources. In this work, we propose a hybrid representation of the reasoning process, where we partially abstract away the initial reasoning steps using latent discrete tokens generated by VQ-VAE, significantly reducing the length of reasoning traces. We explore the use of latent trace abstractions in two scenarios: 1) training the model from scratch for the Keys-Finding Maze problem, 2) fine-tuning LLMs on this hybrid data with an extended vocabulary including unseen latent tokens, for both logical and mathematical reasoning problems. To facilitate effective learning, we introduce a simple training procedure that randomly mixes latent and text tokens, which enables fast adaptation to new latent tokens. Our approach consistently outperforms the baselines methods in various benchmarks.