В основе Quantization Aware Training (QAT) обыкновенно лежит STE (или его модификация), когда градиент просто пробрасывается через не дифференцируемую операцию квантизации. Трюк рабочий, но не имеет под собой теоретических гарантий.
Авторы предлагают рассматривать QAT как задачу оптимизации с ограничениями, и переходят к задаче безусловной оптимизации с множителем Лагранжа:
min L(w) = min L_orig (w) + λ (Q(x) - x)величина λ определяет баланс между лоссом задачи и ограничением. Такая добавка математически эквивалента добавлению error feedback. И рассматривают два варианта - coupled/decoupled - где добавка подается в градиент или момент, но в итоге выбирают decoupled, как более удобный.
В качестве базового алгоритма оптимизации берут AdamW и CAGE отличается от него только наличием error feedback.
Для лучшей сходимости метода константа регуляризации λ разогревается от 0 до максимального значения.
🧪Эксперименты
Метод валидируют, обучая семейство моделей Llama от 30 до 800 M параметров. CAGE стабильно дает некоторое улучшение по лоссу по сравнению с базовым алгоритмом QuEST для разных битностей.
Так называемая эффективная емкость (просадка по лоссу в scaling laws нормализованная на битность) примерно на 0.5 лучше по сравнению с QuEST.
CAGE успешно работает и с MXFP4.
На модельной квадратичной задаче SGD/Adam с STE не могут попасть в оптимум при 4-битной квантизации, а CAGE могет.
💡 Выводы
Метод выглядит вполне себе рабочим и интересным. Было бы интересно посмотреть его в действии на больших обучениях моделей на Blackwell чипах для MXFP4/NVFP4 форматов.