トークンの組み合わせ：改善された言語モデル推論のための潜在トークンとテキストトークンの組み合わせ

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

February 5, 2025

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

cs.AI

要旨

大規模言語モデル（LLMs）は、chain-of-thought（CoT）データで訓練されると、推論や計画立案に優れています。ここでは、ステップバイステップの思考プロセスがテキストトークンによって明示的に示されます。しかしながら、これにより、多くの単語が本質的な推論情報ではなくテキストの一貫性を支えるため、入力が長大になり、これらの入力を処理するのに膨大な計算リソースが必要となります。本研究では、推論プロセスのハイブリッド表現を提案し、VQ-VAEによって生成された潜在的な離散トークンを使用して初期の推論ステップを部分的に抽象化することで、推論トレースの長さを大幅に削減します。私たちは、潜在的なトレースの抽象化の使用を、以下の2つのシナリオで探求します：1）Keys-Finding Maze問題のためにモデルをゼロから訓練すること、2）このハイブリッドデータでLLMsをファインチューニングすること。ここでは、未知の潜在トークンを含む拡張語彙を使用し、論理的および数学的推論問題に対応します。効果的な学習を促進するために、潜在トークンとテキストトークンをランダムに混合する簡単なトレーニング手順を導入し、新しい潜在トークンへの迅速な適応を可能にします。私たちのアプローチは、さまざまなベンチマークでベースライン手法を一貫して上回ることが示されています。

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.