連続的概念を用いたLLMの事前学習

LLM Pretraining with Continuous Concepts

February 12, 2025

著者: Jihoon Tack, Jack Lanchantin, Jane Yu, Andrew Cohen, Ilia Kulikov, Janice Lan, Shibo Hao, Yuandong Tian, Jason Weston, Xian Li

cs.AI

要旨

次のトークン予測は、大規模言語モデルの事前学習で使用される標準的なトレーニング目標です。表現は、トークンレベルのパープレキシティを最適化することによって学習されます。私たちは、連続概念混合（CoCoMix）という新しい事前学習フレームワークを提案します。これは、離散的な次のトークン予測と連続的な概念を組み合わせます。具体的には、CoCoMixは、事前学習された疎なオートエンコーダから学習した連続的な概念を予測し、それらをモデルの隠れ状態に混ぜ込みます。これは、トークンの隠れ表現と交互になります。言語モデリングや下流の推論タスクを含む複数のベンチマークでの実験を通じて、CoCoMixがよりサンプル効率的であり、標準的な次のトークン予測、知識蒸留、および一時停止トークンの挿入を一貫して上回ることを示します。概念学習と交互の両方を組み合わせることがパフォーマンス向上には重要であることがわかりました。さらに、CoCoMixは、予測された概念の直接的な検査と修正を可能にすることで、解釈可能性と誘導可能性を向上させ、モデルの内部推論プロセスを誘導する透明な方法を提供します。

English

Next token prediction has been the standard training objective used in large language model pretraining. Representations are learned as a result of optimizing for token-level perplexity. We propose Continuous Concept Mixing (CoCoMix), a novel pretraining framework that combines discrete next token prediction with continuous concepts. Specifically, CoCoMix predicts continuous concepts learned from a pretrained sparse autoencoder and mixes them into the model's hidden state by interleaving with token hidden representations. Through experiments on multiple benchmarks, including language modeling and downstream reasoning tasks, we show that CoCoMix is more sample efficient and consistently outperforms standard next token prediction, knowledge distillation and inserting pause tokens. We find that combining both concept learning and interleaving in an end-to-end framework is critical to performance gains. Furthermore, CoCoMix enhances interpretability and steerability by allowing direct inspection and modification of the predicted concept, offering a transparent way to guide the model's internal reasoning process.