Token Assortiert: Mischen von Latenten und Text-Token zur Verbesserung der Sprachmodell-Begründung

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

February 5, 2025

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) zeichnen sich durch ihre Fähigkeit zur Schlussfolgerung und Planung aus, wenn sie mit Chain-of-Thought (CoT) Daten trainiert werden, bei denen der schrittweise Gedankengang explizit durch Text-Token dargelegt wird. Dies führt jedoch zu langen Eingaben, in denen viele Wörter die textuelle Kohärenz unterstützen, anstatt Kern-Schlussfolgerungsinformationen zu liefern, und die Verarbeitung dieser Eingaben erfordert erhebliche Rechenressourcen. In dieser Arbeit schlagen wir eine hybride Darstellung des Schlussfolgerungsprozesses vor, bei der wir die anfänglichen Schlussfolgerungsschritte teilweise abstrahieren, indem wir latente diskrete Token verwenden, die von VQ-VAE generiert werden, was die Länge der Schlussfolgerungsspuren erheblich reduziert. Wir untersuchen den Einsatz von latenten Spurenabstraktionen in zwei Szenarien: 1) das Training des Modells von Grund auf für das Keys-Finding Maze Problem, 2) das Feintuning von LLMs auf diesen hybriden Daten mit einem erweiterten Vokabular, das bisher ungesehene latente Token für logische und mathematische Schlussfolgerungsprobleme einschließt. Um ein effektives Lernen zu erleichtern, führen wir ein einfaches Trainingsverfahren ein, das latenten und Text-Token zufällig mischt, was eine schnelle Anpassung an neue latente Token ermöglicht. Unser Ansatz übertrifft konsistent die Basismethoden in verschiedenen 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.