noise_step: Training in 1.58b With No Gradient Memory
[Манускрипт] [Репозиторий]

Введение

Первый пост данного года будет несколько комедийного содержания, как раз в самый раз для прочтения после нескольких бокалов шампанского 🥂 (или чего покрепче 🥃).

Некто Уилл Брикнер выложил на гитхаб презанятнейший опус про обучение тернарной сети в 1.58 бит без необходимости выделения памяти 😱 на градиенты и состояния оптимизатора.

Метод

Товарищи из Мелкософта в серии работ про BitNet показали, что обучая сеть с тернарными весами (принимающими значения только -1, 0, 1 и умноженными на некий скаляр), и низкобитными активациями (4/8 бит) можно выжать качество, сравнимое с fp обучением при тех же бюджетах обучения. Однако, во время само обучения приходится хранить floating-point веса, и состояния оптимизатора, как для fp модели. То есть обучение все равно требует значительных затрат памяти.

Автор данного опуса, вспоминая статью Gradients without Backpropagation, замечают, что операция умножения якобиана по выходу модели на фиксированный вектор не требует backpropagation.

Потому предлагается делать случайные возмущения, причем для случая тернарных весов возмущения это -1, 0, 1. Для улучшения сходимости предлагается отбрасывать слишком малые возмущения (т.е своего рода прунить обновление).

Так как на практике мы используем псевдослучайные числа, то для параметризации модели достаточно хранить только случайные зерна со всех шагов оптимизации. И для обучения GPT-3, взяв данные из техрепорта (тогда еще ClosedAI еще не совсем Closed), получают ~100к шагов оптимизации, и всего несколько мегабайт на хранение 175B весов 🤪. А как вы будете эти сиды превращать в веса - это ваши проблемы)

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

Предложенный метод валидируют на 4-слойной MLP c hidden_size = 256, и данный метод (о, боже!) даже сходится и выдает космические 🚀 почти 90% качества 😱 на MNIST.

Единственный недостаток всей этой красоты, в том, что авторы не релизнули эффективные кернелы для обучения и инференса. Что ж поделать, не все познали дзен куды и тритона (в том числе и пишущий сии строки).

Вывод

Это, наверное, самый забавный каламбур на моей памяти в данной области) Интересно, автор сам дошел до этого или воспользовался помощью всесильного оракула в виде LLM. Я в полном восхищении 😱, в любом случае.