Nous Research представила метод Token Superposition Training для ускорения предобучения LLM в 2-3 раза.
Исследовательская группа Nous Research опубликовала описание метода Token Superposition Training, который, по утверждению авторов, ускоряет предобучение больших языковых моделей в 2–3 раза по реальному времени при неизменном объёме вычислений.
Суть метода
На первые 20–40% обучения модель обрабатывает не отдельные токены, а батчи из нескольких идущих подряд токенов. Их векторные представления усредняются на входе, а на выходе применяется модифицированная функция потерь.
Никакого отбора по смыслу, частоте или содержанию нет — деление чисто механическое, по позиции в тексте. Внутри батча порядок токенов отбрасывается. Модель не учится предсказывать, какой токен стоит на какой конкретной позиции в ближайшем будущем, — только то, какие токены вообще там встречаются.
Размер батча — один из двух ключевых параметров метода. Оптимальное значение растёт вместе с размером модели: для 270 млн параметров оно лежало в диапазоне от 3 до 8 токенов, а для 10 млрд равнялось 16.
Оставшуюся часть обучения модель переводят на обычный режим предсказания следующего токена. По словам Nous Research, готовая модель на инференсе ничем не отличается от обученной стандартным способом — архитектура, оптимизатор, токенизатор или набор данных не изменяются.
Метод протестировали на моделях четырёх размеров: 270 млн, 600 млн и 3 млрд параметров, а также на 10-миллиардной MoE. В самом крупном эксперименте модель с TST достигла более низкого значения лосса, чем сопоставимая по вычислениям базовая модель, примерно за 40% времени и показала лучшие результаты на тестах HellaSwag, ARC и MMLU.
Ограничения
TST расходует обучающие данные быстрее обычного, поскольку модель переваривает больше текста. Если обучающих данных мало, метод становится контрпродуктивным.

