Token Assortiment: Het Mengen van Latente en Tekst Tokens voor Verbeterde Taalmodelredenering

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

Samenvatting

Grote Taalmodellen (LLM's) excelleren in redeneren en plannen wanneer ze getraind worden op keten-van-gedachten (CoT) data, waarbij het stapsgewijze denkproces expliciet wordt uiteengezet door teksttokens. Dit resulteert echter in lange invoeren waarbij veel woorden de tekstuele samenhang ondersteunen in plaats van kernredeneerinformatie, en het verwerken van deze invoeren vereist aanzienlijke rekencapaciteit. In dit werk stellen we een hybride representatie van het redeneerproces voor, waarbij we gedeeltelijk de initiële redeneerstappen abstraheren met behulp van latente discrete tokens gegenereerd door VQ-VAE, wat de lengte van redeneersporen aanzienlijk vermindert. We onderzoeken het gebruik van latente spoorabstracties in twee scenario's: 1) het model vanaf nul trainen voor het Sleutelvindingslabyrintprobleem, 2) LLM's finetunen op deze hybride data met een uitgebreide woordenschat inclusief ongeziene latente tokens, voor zowel logische als wiskundige redeneerproblemen. Om effectief leren te vergemakkelijken, introduceren we een eenvoudige trainingsprocedure die latent- en teksttokens willekeurig mixt, wat snelle aanpassing aan nieuwe latente tokens mogelijk maakt. Onze aanpak presteert consequent beter dan de basismethoden in verschillende 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.