Addestramento preliminare di LLM con concetti continui

LLM Pretraining with Continuous Concepts

February 12, 2025

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

cs.AI

Abstract

La previsione del token successivo è stata l'obiettivo di addestramento standard utilizzato nella preformazione di grandi modelli linguistici. Le rappresentazioni sono apprese come risultato dell'ottimizzazione della perplessità a livello di token. Proponiamo Continuous Concept Mixing (CoCoMix), un nuovo framework di preformazione che combina la previsione discreta del token successivo con concetti continui. In particolare, CoCoMix prevede concetti continui appresi da un autoencoder sparso preaddestrato e li mescola nello stato nascosto del modello alternandoli alle rappresentazioni nascoste dei token. Attraverso esperimenti su diversi benchmark, inclusi modelli linguistici e compiti di ragionamento successivi, dimostriamo che CoCoMix è più efficiente in termini di campioni e supera costantemente la previsione standard del token successivo, la distillazione della conoscenza e l'inserimento di token di pausa. Abbiamo scoperto che combinare sia l'apprendimento dei concetti che l'alternanza in un framework end-to-end è fondamentale per ottenere miglioramenti delle prestazioni. Inoltre, CoCoMix migliora l'interpretabilità e la guidabilità consentendo l'ispezione diretta e la modifica del concetto previsto, offrendo un modo trasparente per guidare il processo di ragionamento interno del modello.

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.