Метод

Рассматривают обучение 1-, 1.58-, 2-, 3-, 4- битных моделей MobileLLM.

По существу в работе 2️⃣ основные идеи.

⭐️ Можно какую-то часть времени обучать модель в исходной точности, а затем врубить QAT. Оптимальнее всего оказывается 90% учить во float16, а последние 10%. Т.е большую часть обучения динамика совпадает с floating-point моделью, а оставшихся 10% хватает для восстановления качества после сжатия.
⭐️ Под каждую битность подбирают оптимальную сетку. При бинаризации весов оптимально (да больше ничего и не сделаешь) брать знак от весов, при 1.58 и 2-х битах брать определенную в данной работе Stretched Elastic Quant сетку (небольшую модификацию LSQ), и LSQ при 3-х и 4-х битах.

Причем важно скейлы квантизации держать как обучаемые параметры. min-max масштабирование работает совсем плохо на низких битностях, обрезка (непонятно каким образом) лучше, но все же заметно хуже learnable scales.

Для обучения всего хозяйства используют STE без каких-то наворотов.

Кроме того, еще пробуют QAT на MobileLLM не с нуля, а с готовых чекпоинтов, и это оказывается ожидаемо лучше.

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

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

Сначала подход проверяют на семействе моделей MobileLLM (от 125M до 1.5B). Битности 1.58 и 2 оказываются Парето-оптимальными. Качество оценивают как средне по нескольким бенчам из lm-eval-harness.

Затем ParetoQ проверяют на Llama-3-8B, где предложенный подход заметно опережает LLM-QAT, EfficientQAT, и PTQ подходы (даже с AQLM и QuIP сравниваются). Учат модель на 30B токенах, что не так уж мало, но и не так уж много. C PV-Tuning и QTIP тактично не сравниваются, хитрецы. При всей дороговизне PV, он задействует все же меньше токенов).

На тернарной квантизации авторы показывают превосходство над прошлыми QAT подходами - Spectra, BitNet - при фиксированном бюджете обучения.

Кроме того, авторы написали CPU kernel, дающий хорошее ускорение при 2-битной квантизации, благодаря чему 2-битные модельки оказываются оптимальными с точки зрения баланса между скоростью инференса и качеством.

Вывод

Подход выглядит разумным и рабочим, без rocket 🚀 science. Для полного счастья не хватает квантизации активаций и замеров на более сложных задачах.