Token Assorti : Mélange de Jetons Latents et de Jetons de Texte pour une Amélioration de la Raisonnement des Modèles de Langue

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

February 5, 2025

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

cs.AI

Résumé

Les grands modèles de langage (LLMs) excellent dans le raisonnement et la planification lorsqu'ils sont entraînés sur des données de chaîne de pensée (CoT), où le processus de pensée étape par étape est explicitement décrit par des jetons de texte. Cependant, cela entraîne des entrées longues où de nombreux mots soutiennent la cohérence textuelle plutôt que les informations de raisonnement de base, et le traitement de ces entrées consomme des ressources de calcul substantielles. Dans ce travail, nous proposons une représentation hybride du processus de raisonnement, où nous abstrayons partiellement les étapes de raisonnement initiales en utilisant des jetons discrets latents générés par VQ-VAE, réduisant significativement la longueur des traces de raisonnement. Nous explorons l'utilisation d'abstractions de traces latentes dans deux scénarios : 1) entraîner le modèle à partir de zéro pour le problème du labyrinthe de recherche de clés, 2) affiner les LLMs sur ces données hybrides avec un vocabulaire étendu comprenant des jetons latents invisibles, pour des problèmes de raisonnement logique et mathématique. Pour faciliter l'apprentissage efficace, nous introduisons une procédure d'entraînement simple qui mélange de manière aléatoire des jetons latents et textuels, ce qui permet une adaptation rapide à de nouveaux jetons latents. Notre approche surpasse de manière constante les méthodes de base dans divers 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.