Итоговая идея -- надо сконцентрироваться на обновлении малых по магнитуде весов. Во-первых, благодаря оверпараметризации этого может быть достаточно. Во-вторых, хоть пересечение и неидеальное, всё равно малые веса дают хороший шанс на пересечение с большими градиентами. Наконец, большие веса, вероятно, содержат в себе что-то важное с предобучения, так что трогать их опасно. Такой подход к обновлению весов авторы называют nano gradient descent. Соответственно основанный на этой идее оптимизатор называется NanoAdam. В двух словах идея в том, что мы выбираем подмножество параметров, которые будем обучать. Они задаются маской, выбирающей из всех параметров отобранные, для них и выполняются стандартные апдейты Adam. В дополнение к маске здесь присутствует так называемый density scheduler, который диманически подстраивает долю параметров, участвующих в обучении (по дефолту линейное затухание). В итоге по некоторому графику обновляется доля параметров (k штук) и генерится маска (bottom k по магнитуде веса), это происходит каждые d и m итераций соответственно. Каждый раз обновлять маску не имеет смысла, параметры с малой магнитудой имеют тенденцию оставаться таковыми на протяжении всей оптимизации. Кроме того благодаря этому можно сохранить и динамику обучения с моментом. Абляции показывают, что предложенная процедура имеет смысл. Если вместо малых весов выбирать большие или рандомные, то результаты стабильно хуже. Особенно плохо всё с большими. Если выбирать веса с большими градиентами, то тоже работает похуже. Провели кучу экспериментов по файнтюну на NLP задачах с BERT-Base/Large и OPT-1.3B. В среднем NanoAdam лучше альтернатив в виде MicroAdam, Galore и AdamW обычный/8bit. Доверительные интервалы, правда, непонятны. Самое главное, что памяти требует меньше. На OPT-1.3B это 11.6G против 13G у MicroAdam/Adam-8b, 14G у GaLore и 18G у AdamW. Нормальная экономия. Отдельно проверили catastrophic forgetting на CV задачах c ViT и ResNet при continual learning, когда сначала зафайнтюнили на CIFAR-10 (Task 1), а затем на Flowers102 (Task 2). NanoAdam -- чемпион. Я только не уловил, почему у резнетов такая большая просадка на Flowers102. NanoAdam также привёл к наименьшему изменению параметров (за исключением классификационных голов) по L2 метрике (и несмотря на больший learning rate). Выглядит интересно. И мне нравится подход from first principles.