Предварительное обучение LLM с непрерывными концепциями

LLM Pretraining with Continuous Concepts

February 12, 2025

Авторы: Jihoon Tack, Jack Lanchantin, Jane Yu, Andrew Cohen, Ilia Kulikov, Janice Lan, Shibo Hao, Yuandong Tian, Jason Weston, Xian Li

cs.AI

Аннотация

Прогнозирование следующего токена было стандартной целью обучения, используемой в предварительном обучении больших языковых моделей. Представления изучаются в результате оптимизации на уровне токенов. Мы предлагаем Continuous Concept Mixing (CoCoMix), новую предварительную обучающую структуру, которая объединяет дискретное прогнозирование следующего токена с непрерывными концепциями. Конкретно, CoCoMix предсказывает непрерывные концепции, изученные из предварительно обученного разреженного автокодировщика, и вкладывает их в скрытое состояние модели, чередуя с скрытыми представлениями токенов. Через эксперименты на нескольких бенчмарках, включая языковое моделирование и задачи логического вывода, мы показываем, что CoCoMix более эффективен по образцам и последовательно превосходит стандартное прогнозирование следующего токена, дистилляцию знаний и вставку токенов паузы. Мы обнаружили, что комбинирование как обучения концепций, так и чередования в единой структуре критично для увеличения производительности. Кроме того, CoCoMix улучшает интерпретируемость и управляемость, позволяя непосредственно проверять и изменять предсказанный концепт, предлагая прозрачный способ направлять внутренний процесс рассуждений модели.

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.