Preentrenamiento de LLM con Conceptos Continuos

LLM Pretraining with Continuous Concepts

February 12, 2025

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

cs.AI

Resumen

La predicción del siguiente token ha sido el objetivo estándar de entrenamiento utilizado en el preentrenamiento de modelos de lenguaje a gran escala. Las representaciones se aprenden como resultado de optimizar la perplejidad a nivel de token. Proponemos Continuous Concept Mixing (CoCoMix), un nuevo marco de preentrenamiento que combina la predicción discreta del siguiente token con conceptos continuos. Específicamente, CoCoMix predice conceptos continuos aprendidos de un autoencoder disperso preentrenado y los mezcla en el estado oculto del modelo intercalándolos con representaciones ocultas de tokens. A través de experimentos en múltiples referencias, incluyendo modelado de lenguaje y tareas de razonamiento posteriores, demostramos que CoCoMix es más eficiente en muestras y supera consistentemente a la predicción estándar del siguiente token, destilación de conocimiento e inserción de tokens de pausa. Encontramos que combinar tanto el aprendizaje de conceptos como la intercalación en un marco de extremo a extremo es fundamental para obtener mejoras de rendimiento. Además, CoCoMix mejora la interpretabilidad y la capacidad de control al permitir la inspección directa y modificación del concepto predicho, ofreciendo una forma transparente de guiar el proceso de razonamiento interno del modelo.

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.