Эффективное предварительное обучение с суперпозицией токенов

Efficient Pre-Training with Token Superposition

May 7, 2026

Авторы: Bowen Peng, Théo Gigant, Jeffrey Quesnelle

cs.AI

Аннотация

Предварительное обучение больших языковых моделей часто оказывается непомерно дорогостоящим и неэффективным при масштабировании, требуя сложных и инвазивных модификаций для достижения высокой пропускной способности данных. В данной работе мы представляем метод обучения с суперпозицией токенов (Token-Superposition Training, TST) — простой метод, не требующий изменений в существующей инфраструктуре, который значительно повышает пропускную способность данных на единицу FLOPs в ходе предварительного обучения без модификации параллелизма, оптимизатора, токенизатора, данных или архитектуры модели. TST выполняется в два этапа: (i) высокоэффективный этап суперпозиции, на котором мы объединяем несколько следующих подряд токенов в один мешок и обучаем модель с помощью целевой функции многогорячей кросс-энтропии (MCE), и (ii) этап восстановления, на котором мы возвращаемся к стандартному обучению. Мы всесторонне оцениваем TST на масштабах 270M и 600M параметров, а также проводим валидацию на моделях с 3B и 10B параметров, использующих смесь экспертов A1B, что демонстрирует высокую устойчивость метода в различных условиях. В итоге TST стабильно превосходит базовые значения функции потерь и оценки на последующих задачах, а при равных значениях функции потерь обеспечивает сокращение общего времени предварительного обучения до 2,5 раз на масштабе 10B A1B.

English

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.