Tokens Variados: Mezclando Tokens Latentes y de Texto para Mejorar el Razonamiento del Modelo de Lenguaje

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

Resumen

Los Modelos de Lenguaje Grandes (LLMs) destacan en el razonamiento y la planificación cuando se entrenan con datos de cadena de pensamiento (CoT), donde el proceso de pensamiento paso a paso está explícitamente delineado por tokens de texto. Sin embargo, esto resulta en entradas extensas donde muchas palabras respaldan la coherencia textual en lugar de la información de razonamiento central, y el procesamiento de estas entradas consume recursos computacionales sustanciales. En este trabajo, proponemos una representación híbrida del proceso de razonamiento, donde abstraemos parcialmente los pasos de razonamiento iniciales utilizando tokens discretos latentes generados por VQ-VAE, reduciendo significativamente la longitud de las trazas de razonamiento. Exploramos el uso de abstracciones de trazas latentes en dos escenarios: 1) entrenando el modelo desde cero para el problema del Laberinto de Búsqueda de Claves, 2) ajustando finamente LLMs en estos datos híbridos con un vocabulario extendido que incluye tokens latentes no vistos, tanto para problemas de razonamiento lógico como matemático. Para facilitar un aprendizaje efectivo, introducimos un procedimiento de entrenamiento simple que mezcla aleatoriamente tokens latentes y de texto, lo que permite una rápida adaptación a nuevos tokens latentes. Nuestro enfoque supera consistentemente a los métodos de referencia en varios 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.