BitNet a4.8: 4-bit Activations for 1-bit LLMs
[Статья] [Репозиторий (пока не обновили код)]

Продолжение истории с BitNet от Microsoft.
1️⃣ В первой части была предложили бинарную квантизацию весов.
2️⃣ Во второй части вместо бинарной квантизации тернарная, но зато качество заметно возросло.

При этом активации квантовались absmax в INT8.

В данной статье фокусируются на сжатии активаций (в частности, KV-кэшей) в 4 и бита и меньше.

Метод

Авторы строят гистограммы весов активаций после линейных слоев. Выходы q, k, v, gate, и up неплохо описываются нормальным распределением. А вот out, down имеют тяжелые хвосты (потому тяжело квантизуются).

Поэтому предлагается квантовать в 4 бита активации только первой группы слоев, а для выхода out_proj (следуя работе QSparse, которую сам пока не читал 😂) прореживают активации через topk (с 50% sparsity).

Кроме того вместо swish в FFN gate применяют ReLU^2, чтобы форсировать прореживание активаций в down projection. Впрочем, так как паттерн случайный, использовать его для ускорения сходу не получится.

Также рассматривают FP4 квантизацию вместо INT4 с большим динамическим диапазоном.

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

Обучают семейство 🦙-подобных моделей от 700M до 7B параметров на Красной Пижаме. Причем обучают в 2 стадии, сначала с 8-битными активациями большее время обучения (как в BitNet1.58), а затем в перечисленных выше группы слоев квантуют в 4 бита (или прореживают).

⭐️ BitNet a4.8 (INT4/FP4) почти не уступает по качеству fp бейзлайну (не очень сильному, впрочем).
⭐️ ReLU^2 дает значительную (около 80%) степень разреженности в down_proj.
⭐️ Обучение чисто в INT4 не сходится в принципе. С FP4 гораздо хуже, чем предложенный вариант.
⭐️ ReLU^2 дает +/- тоже качество, что более традиционный Swish.
⭐️ Кэши квантизуют в 3 бита кроме <bos>, <eos> токенов, которые более чувствительны и поддерживаются 4 битными, без просадки в качестве по сравнению с 4-битными кэшами.

Определив хорошую конфигурацию, 2B модель обучают на 2T токенах и получают метрики, конкурирующие с хорошими fp16 моделями в той же весовой категории (Llama-3.2, gemma-2, Qwen-2.5). Впрочем, ребята из Мелкософта известны тем, что добавляют тест в train (как было в случае с моделями Phi) 😈.

Вывод

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