⚙️ Метод

💡 Идея довольно проста: некоторую долю r от обучения подаём на вход и пытаемся предсказать несколько следующих токенов, а именно:

- 🔤 Токенизируем текст как обычно.
- 📊 После прогона через матрицу эмбеддингов усредняем s подряд идущих токенов. Проходимся по тексту непересекающимися окнами размера s.
- 🎯 На выходе предсказываем s следующих токенов без учёта порядка.

Итоговый лосс — просто усреднение кросс-энтропии для каждого следующего токена.

💸 Экономия достигается за счёт того, что на обучении мы гоняем в s раз меньше токенов через модель. Т. е. за тот же бюджет прогонов токенов через LLM модель видит в r·s + (1 − r) раз больше данных из корпуса.

После окончания этой стадии переходим к стандартному обучению на исходных токенах.

🧪 Эксперименты

Метод валидируют на трёх dense llama-like моделях (270M, 600M, 3B) и (10B-A1B) MoE.

При заданном количестве FLOPs и даже Wall time по итоговому лоссу оно оказывается лучше бейзлайна, который просто учится на токенах. И вроде даже метрики чуть получше, хоть и все едва лучше рандома.

📈 Оптимальнее всего ставить r в диапазоне от 0.2 до 0.4. Если меньше, то слишком мала доля обучения на усреднённых токенах, чтобы дать выигрыш; если больше — модель не успевает перестроиться на стандартное обучение.

Оптимально по качеству ставить s в диапазоне 5–8.

Можно по отдельности усреднять на входе или предсказывать bag-of и даже так получать выигрыш, но и то и другое вместе работает ещё эффективнее.

Output-усреднение сродни MTP, а MTP вроде бы считается полезным для ускорения сходимости.

🤔 А вот почему input-усреднение хорошо работает, не очень очевидно. Авторы говорят про какую-то регуляризацию, smoothing латентного пространства. Но без чёткого понимания.

Из интересного ещё оказывается, что предсказывать мешок следующих токенов работает лучше, чем пытаться задать какой-то порядок. Явная добавка позиционных эмбеддингов на выход только портит качество.

Всякие манипуляции с RoPE, попытки адаптировать их под усреднение токенов ни к чему не привели.

⚠️ Кроме того, авторы отмечают, что в случае ограниченного количества данных их метод менее data-efficient, так как использует огрублённые токены, тем самым теряя часть информации.

💡 Выводы

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