LLM-Vortraining mit kontinuierlichen Konzepten

LLM Pretraining with Continuous Concepts

February 12, 2025

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

cs.AI

Zusammenfassung

Die Vorhersage des nächsten Tokens war das Standard-Trainingsziel, das bei der Vorbereitung großer Sprachmodelle verwendet wurde. Die Repräsentationen werden als Ergebnis der Optimierung für Token-Level-Perplexität gelernt. Wir schlagen Continuous Concept Mixing (CoCoMix) vor, ein neuartiges Vorbereitungsframework, das diskrete Vorhersage des nächsten Tokens mit kontinuierlichen Konzepten kombiniert. Speziell sagt CoCoMix kontinuierliche Konzepte vorher, die von einem vorab trainierten dünnen Autoencoder gelernt wurden, und mischt sie in den versteckten Zustand des Modells, indem sie mit den versteckten Token-Repräsentationen abwechseln. Durch Experimente an mehreren Benchmarks, einschließlich Sprachmodellierung und nachgelagerten Denkaufgaben, zeigen wir, dass CoCoMix effizienter in der Stichprobenverwendung ist und standardmäßige Vorhersage des nächsten Tokens, Wissensverdichtung und das Einfügen von Pause-Tokens konsequent übertrifft. Wir stellen fest, dass die Kombination von Konzeptlernen und Abwechseln in einem End-to-End-Framework entscheidend für Leistungssteigerungen ist. Darüber hinaus verbessert CoCoMix die Interpretierbarkeit und Steuerbarkeit, indem es eine direkte Inspektion und Modifikation des vorhergesagten Konzepts ermöglicht und somit einen transparenten Weg bietet, um den internen Denkprozess des Modells zu lenken.

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.