Pré-treinamento de LLM com Conceitos Contínuos

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

Resumo

A previsão do próximo token tem sido o objetivo de treinamento padrão utilizado no pré-treinamento de grandes modelos de linguagem. As representações são aprendidas como resultado da otimização da perplexidade ao nível do token. Propomos o Continuous Concept Mixing (CoCoMix), um novo framework de pré-treinamento que combina a previsão discreta do próximo token com conceitos contínuos. Especificamente, o CoCoMix prevê conceitos contínuos aprendidos a partir de um autoencoder esparsamente pré-treinado e os mistura no estado oculto do modelo intercalando com as representações ocultas do token. Através de experimentos em múltiplos benchmarks, incluindo modelagem de linguagem e tarefas de raciocínio subsequentes, mostramos que o CoCoMix é mais eficiente em amostras e consistentemente supera a previsão padrão do próximo token, destilação de conhecimento e inserção de tokens de pausa. Descobrimos que combinar tanto a aprendizagem de conceitos quanto a intercalação em um framework de ponta a ponta é crucial para ganhos de desempenho. Além disso, o CoCoMix aprimora a interpretabilidade e a capacidade de direção, permitindo a inspeção direta e modificação do conceito previsto, oferecendo uma maneira transparente de guiar o processo de raciocínio interno do 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.