EDEN: Communication-Efficient and Robust Distributed Mean Estimation for Federated Learning
[Статья] [Код]

Квантизовать можно веса, можно активации. А что еще можно?

Правильно, градиенты! 👻

И это очень даже нужно и полезно в контексте распределенной оптимизации.
Если обучение на многих хостах, а скорость передачи информации между серверами не очень высока, то на синхронизацию между процессами может уходить основная доля времени, а не на сами вычисления. Дабы уменьшить трафик можно сжимать 🗜 градиенты - прунить, квантовать - что душе угодно. При этом, само собой, хочется их не слишком испортить, иначе все обучение пойдет прахом.

Метод

Идея квантовать градиенты для распределенного не нова - датируется еще 2016 годом (QSGD). Но данная работа предлагает ряд интересных идей (некоторые из которых мы узнаем в литературе вышедшей позднее, а сама статья 2021 года).

1️⃣ Случайные повороты
2️⃣ Оптимальные решетки квантования
3️⃣ Случайный прунинг
4️⃣ Энтропийное кодирование

В тензоре значения могут быть распределены, вообще говоря, самым произвольным и поганым (с точки зрения удобства квантования) образом. Но, известно, что если “повернуть” тензор некоей случайной матрицей, то с большой вероятностью полученные значения будут распределены как N(0, sigma^2). В качестве случайных поворотов используют (кто бы мог подумать!) - случайные Адамаровы матрицы. Каждый процесс хранит seed генерации, а главный сервер, зная этот seed на деквантизации, может применить обратное преобразование.

Решетку квантизации (а тут дело идет о квантовании в 1 и 2 бита) надо выбирать с умом. Потому берут не абы какие, а MSE-оптимальные (для 1 бита находится аналитически, для 2-ух численно).

Дабы увеличить сжатие - можно еще случайным образом запрунить ☠️ часть весов. Маску передавать при этом не надо - она как и вращение задается состоянием случайного генератора.

Решетки квантизации можно еще дотюнить при заданном ограничении на битность посредством Entropy-Constrained Vector Quantization.

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

И вся эта кухня идет с теоретическими гарантиями на ошибку аппроксимации.

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

Метод валидируют на синтетических распределениях и распределенном обучении. В качестве бейзлайнов рассматривают QSGD, Hadamard + SQ (Stohastic Quantization), Kashin + SQ. В качестве бенчей для распределенного обучения рассматривают EMNIST (с маленькой CNNкой) и Shakespeare (с LSTM). Не самые свежие и впечатляющие бенчи, но что поделаешь 😞.

EDEN наиболее точно приближает исходные данные (в сравнении с бейзлайнами), при этом по скорости не такой быстрый как QSGD, но быстрее Kashin + SQ (ближайший с точки зрения качества подход).

В 4 бита разницы вообще нет между EDEN и float обучением, в 2 и 1 бит появляется некий зазор в качестве, но не критичный, и EDEN стабильно лучше конкурентных подходов.

Вывод

Интересная с точки зрения математики и методологии статья. Кажется, что если бы публикация чуть дотерпела до тех времен (и у авторов был на то ресурс), когда квантовать LLMки стало мейнстримом предложенные идеи были бы обязательно применены в контексте сжатия LLM.