Метод ускорения претрейна LLM через усреднение токенов

⚙️ Метод 💡 Идея довольно проста: некоторую долю 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, так как использует огрублённые токены, тем самым теряя часть информации. 💡 Выводы Дешёвый и простой способ удешевить претрейн. Вопрос — в масштабируемости на большие сетапы. Кажется, что сильное огрубление токенов работает, пока сама модель ещё очень плоха и ничего не умеет. А когда она начинает выдавать адекватные метрики, слишком много ценной информации может теряться, ограничивая достижимое качество.