Основным затруднением при оптимизации весов и активаций в низкой точности является высокая степень шума. Напомню, что операция квантизации недифференцируема, и дабы все можно было оптимизировать градиентными методами применяют трюк имени Бенджио под названием STE (небось, тоже спиздил у Шмидхубера) , где градиент просто пробрасывается через недифференцируемую операцию (как будто вместо нее стоит функция y=x). Но при низких битностнях, такая оценка сильно расходится от истинного градиента и не сходится нормально.
Авторы формулируют задачу оптимизации STE, как минимизацию между оцененным псевдоградиентом и истинным. Предполагая гладкость функции потерь по весам, можно оценить ошибку градиента, как константа на ошибку квантизации весов.
Веса модели разделяют на 2️⃣ группы - с ошибкой квантизации ниже и выше некоторого порога. На шаге оптимизации учитывают только градиенты от весов с ошибкой ниже заданного порога, ибо вторые как раз и вносят шум и нестабильность в обучение.
Далее, дабы работать с более регулярным распределением весов/активаций, которое проще квантизовать, применяют пару трюков:
1️⃣ Чтобы привести распределение в более удобо квантизуемый вид применяют известный старым читателям канала трюк - вращения Адамаровыми матрицами как весов и активаций. В результате получают что-то близкое к гауссиане.
2️⃣ Дабы привести все приблизительно к
N(0, 1) применяют RMS нормализацию к результату шага 1. А для N(0, 1) можно уже численно найти оптимальный скейлинг фактор для решетки квантизации и пользоваться им.То есть в итоге алгоритм выглядит следующим образом:
🎯 На прямом проходе вращаем и нормализуем веса, сохраняя Адамаровы матрицы
🎯 На обратном проходе применяем обратное Адамарово преобразование и маскируем градиент
В Ablation показывают, что trust estimator (отбрасывание градиентов по шумным весам) в связке с Адамаровыми вращениями дает хорошую близость с истинным градиентом, в то время как vanilla STE и без Адамара корреляция низкая.