Оглавление

Препроцессинг данных и архитектурные ухищрения обычно идут рука об руку, но иногда появляются решения, которые элегантно обходят привычные ограничения. Исследовательская группа Nous Research представила метод Token Superposition Training (TST), который позволяет существенно сократить время предварительного обучения больших языковых моделей (LLM) без изменения их архитектуры или используемых оптимизаторов.

Как пишет MarkTechPost, технология демонстрирует впечатляющие результаты на моделях разного масштаба — от компактных решений на 270 миллионов параметров до систем Mixture-of-Experts (MoE) объемом 10 миллиардов. В частности, для MoE-модели время обучения сократилось с 12 311 до 4 768 GPU-часов на ускорителях NVIDIA B200, что эквивалентно ускорению в 2,5 раза при достижении более низкого уровня потерь.

Механика двух фаз: от суперпозиции к восстановлению

Идея TST заключается в разделении процесса обучения на два этапа, что напоминает своего рода «интенсивный курс» перед основной программой. На первой фазе, занимающей от 20% до 40% общего времени, модель работает не с отдельными токенами, а с их «пакетами» (bags). Эмбеддинги нескольких идущих подряд токенов усредняются в один латентный вектор, который и подается на вход трансформеру.

Чтобы сохранить объем вычислений (FLOPs) на прежнем уровне, разработчики пропорционально увеличивают длину последовательности данных. Таким образом, за одну операцию модель «проглатывает» значительно больше полезной информации. На выходе система предсказывает не следующий токен, а сразу весь следующий пакет, используя модифицированную функцию потерь — multi-hot cross-entropy (MCE), которая распределяет вероятности между всеми целевыми элементами пакета.

Вторая фаза — восстановление — подразумевает возврат к классическому предсказанию следующего токена. В этот момент происходит резкий, но кратковременный скачок функции потерь, который стабилизируется через несколько тысяч шагов. Важно, что на выходе получается стандартная модель, полностью совместимая с существующими методами инференса, поскольку все специфические механизмы TST удаляются после завершения обучения.

Практическая реализация и ограничения

Для тех, кто решит внедрить этот метод в свои пайплайны на PyTorch, процесс выглядит достаточно прозрачным. Основные изменения касаются формирования входных данных и логики вычисления эмбеддингов. Например, при использовании пакета размером s, входной тензор необходимо переформатировать, чтобы объединить соседние токены, а в слое эмбеддингов реализовать суммирование векторов с последующим усреднением.

В условиях избытка обучающих текстов это выглядит как триумф инженерии, однако для узкоспециализированных доменов, где каждый токен на счету, TST может оказаться избыточным. Технология выигрывает за счет «грубого» поглощения информации в начале пути, но вопрос о том, не теряются ли при таком усреднении тонкие семантические связи, остается открытым. Это отличный костыль для тех, кто завален данными, но не готов бесконечно платить за аренду кластеров.

Эксперименты подтвердили эффективность подхода на популярных бенчмарках, таких как HellaSwag и ARC. Модели, обученные с применением TST, не только быстрее достигали целевых показателей, но и зачастую превосходили базовые варианты по точности. Однако стоит учитывать «железный» нюанс: преимущество сохраняется только в сценариях, где обучение ограничено вычислительными ресурсами, а не объемом доступных данных, так как метод потребляет в несколько раз больше токенов за единицу времени.

Исследователи также отметили ряд неудачных попыток, предостерегая коллег от добавления позиционного кодирования RoPE до этапа усреднения или использования нескольких независимых голов предсказания. Как выяснилось, простейшая реализация с усреднением векторов и единой головой работает стабильнее всего, доказывая, что в мире глубокого обучения кратчайший путь не всегда самый очевидный.