Pré-entraînement LLM avec Concepts Continus

LLM Pretraining with Continuous Concepts

February 12, 2025

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

cs.AI

Résumé

La prédiction du jeton suivant a été l'objectif d'entraînement standard utilisé dans la préformation de grands modèles de langage. Les représentations sont apprises en optimisant la perplexité au niveau du jeton. Nous proposons Continuous Concept Mixing (CoCoMix), un nouveau cadre de préformation qui combine la prédiction discrète du jeton suivant avec des concepts continus. Plus précisément, CoCoMix prédit des concepts continus appris à partir d'un autoencodeur clairsemé préappris et les mélange à l'état caché du modèle en les entrelaçant avec les représentations cachées des jetons. À travers des expériences sur plusieurs référentiels, y compris la modélisation de langage et des tâches de raisonnement en aval, nous montrons que CoCoMix est plus efficace en termes d'échantillonnage et surpasse de manière constante la prédiction standard du jeton suivant, la distillation des connaissances et l'insertion de jetons de pause. Nous constatons que combiner à la fois l'apprentissage de concepts et l'entrelacement dans un cadre de bout en bout est crucial pour les gains de performance. De plus, CoCoMix améliore l'interprétabilité et la dirigibilité en permettant l'inspection directe et la modification du concept prédit, offrant ainsi un moyen transparent de guider le processus de raisonnement interne du modèle.

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.