Token Assortiti: Mescolare Token Latenti e Testuali per un Miglioramento del Ragionamento del Modello Linguistico

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

February 5, 2025

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

cs.AI

Abstract

I Large Language Models (LLM) eccellono nel ragionamento e nella pianificazione quando vengono addestrati su dati di catena di pensiero (CoT), dove il processo di pensiero passo dopo passo è esplicitamente delineato da token di testo. Tuttavia, ciò porta a input lunghi in cui molte parole supportano la coerenza testuale piuttosto che le informazioni di ragionamento principale, e il processamento di questi input richiede notevoli risorse computazionali. In questo lavoro, proponiamo una rappresentazione ibrida del processo di ragionamento, in cui astraiamo parzialmente i passaggi di ragionamento iniziali utilizzando token discreti latenti generati da VQ-VAE, riducendo significativamente la lunghezza delle tracce di ragionamento. Esploriamo l'uso di astrazioni di tracce latenti in due scenari: 1) addestrare il modello da zero per il problema del Labirinto della Ricerca delle Chiavi, 2) raffinare i LLM su questi dati ibridi con un vocabolario esteso che include token latenti non visti, per problemi di ragionamento logico e matematico. Per facilitare l'apprendimento efficace, introduciamo una procedura di addestramento semplice che mescola casualmente token latenti e di testo, consentendo un'adattamento rapido ai nuovi token latenti. Il nostro approccio supera costantemente i metodi di riferimento in vari benchmark.

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.