Почти нулевые эмбеддинги токенов обнулили 20% шагов дообучения gated-моделей на телугу

Если вы дообучаете открытую модель с гейтами — Mamba, Gated DeltaNet или xLSTM, то есть рекуррентную архитектуру, где доступ к накопленному состоянию регулируют «вентили», — заметная часть шагов оптимизации может не менять веса вообще. Виновата не ошибка в данных, а токены, чьи входные эмбеддинги имеют почти нулевую норму. В таких gated-моделях один подобный токен усиливает градиенты на десятки порядков, а стандартный клиппинг градиентов превращает этот взрыв в полное обнуление обновления. При continual pretraining на телугу так пропадает примерно 20% шагов, причём ни исключений, ни предупреждений в логах нет.

Откуда берутся «нулевые» строки в чекпойнтах

Авторы проверили 291 открытый репозиторий Hugging Face (280 различных таблиц эмбеддингов). Веса и токенизатор читали напрямую через HTTP byte-range запросы. Чтобы отсеять служебные и недостижимые записи, каждую подозрительную строку прогоняли через проверку: декодировали ID и токенизировали заново — отдельно и в восьми шаблонах фраз. Токен считался достижимым, только если получался тот же ID.

Из 291 репозитория 92 содержат строки с очень малой нормой, а 67 (23,0%) — достижимые околонулевые токены. Среди 101 модели со связанными входной и выходной матрицами таких токенов нет ни одного. Среди 190 моделей с раздельными матрицами они есть в 67 (35,3%). У обычных Transformer доля составляет 30,9%, у gated и гибридных моделей — 52,6% (20 из 38).

Причина кроется в претрейне. Токен, который ни разу не встретился в тексте, не получает градиента во входной строке. Но стандартный AdamW применяет weight decay ко всем двумерным матрицам, и неиспользуемая строка шаг за шагом сжимается к нулю. Выходная строка при этом получает ненулевые градиенты softmax на каждом шаге и остаётся здоровой: у 92% из 979 проверенных токенов (901) выходные нормы не ниже медианы по словарю. При связанных весах выходные градиенты сами удерживают нормальный масштаб. В выборке на 364 млн токенов все 251 околонулевая строка принадлежат токенам с нулём вхождений. Затронуты обычные строки: подслова телугу и баскского, монгольские слова на кириллице, фрагменты кода с CRLF. Дообученные и дистиллированные модели наследуют эти строки от базовых чекпойнтов.

Почему страдают именно gated-модели

В Transformer на непервой позиции проблемы нет. Уже первый слой внимания усредняет значения предыдущих токенов нормального масштаба и возвращает сигнал в рабочий диапазон. То же происходит, если в начале последовательности стоит BOS.

В gated-рекуррентностях чтение исторического состояния умножается на функцию от текущего токена — гейт или линейный запрос. При нулевом входе множитель тоже уходит в ноль, и вклад истории сжимается. Остаточный вектор остаётся почти нулевым через несколько блоков подряд, а Pre-RMSNorm на каждом слое умножает обратные градиенты на большой коэффициент. Эффект копится по глубине: градиенты входных эмбеддингов растут более чем на 30 порядков (до 35). Гибриды упираются в потолок, чистые gated-модели растут дальше.

Причинность авторы проверили вмешательством: открыли гейт первого слоя в прямом проходе, сохранив все обратные производные, — и градиент резко упал. Одно лишь отключение обратной производной гейта проблему не убрало.

Тихое обнуление батча и два способа лечения

Дальше срабатывает клиппинг по глобальной норме. PyTorch, HuggingFace Trainer и FSDP2 суммируют квадраты градиентов в FP32. Если хоть одна запись слишком велика, при возведении в квадрат получается бесконечность. Градиенты делятся на бесконечную норму и превращаются в нули. AdamW делает шаг только по старому моменту и weight decay на остальной сети, а новая околонулевая строка с нулевым моментом остаётся замороженной и ломает следующие батчи. В инференсе и on-policy RL-роллаутах этот сбой не возникает, но полнопараметрическое RL-обновление на промптах с околонулевыми чат-тегами его запускает.

Авторы предлагают два решения без накладных расходов по времени шага и пиковой памяти:

• Раздельный клиппинг входных эмбеддингов. Норма считается отдельно от остальных весов, в 64-битной арифметике или с перемасштабированием, и это три строки кода. Потери обновлений исчезают, затронутая строка «исцеляется» за один шаг. Сканировать словарь и менять архитектуру не нужно.
• Положительное смещение гейта для претрейна с нуля. Гейт остаётся положительным при нулевом входе, и первый слой пропускает ненулевое состояние. Околонулевые градиенты падают более чем на 28 порядков на моделях до 370M параметров, а валидационная perplexity сопоставима с базовой.

Отключение weight decay для эмбеддингов в уже выпущенных чекпойнтах не поможет: взрыв происходит на первом шаге. А при обучении с нуля одного такого отключения тоже мало — блоки всё равно усиливают градиенты при малой инициализации норм.

Практический вывод простой. Если вы дообучаете несвязанный gated-чекпойнт на низкоресурсном языке или со своими чат-тегами, проверьте нормы входных строк эмбеддингов и вынесите их клиппинг отдельно.