Stack More Layers Differently: High-Rank Training Through Low-Rank Updates
[Статья][Код]

Обучение всех параметров больших языков моделей весьма прожорливо по памяти из-за необходимости хранить кроме самой тяжеловесной модели еще и состояния оптимизатора (8 байт на параметр).

LoRA, один из самых ходовых методов PEFT, заключающийся в обучении низкоранговых добавок к весам позволяет сильно сэкономить по памяти, демонстрируя при этом хорошее качество при обучении предобученной модели на downstream задачах. Но низкоранговые представления имеют место при дообучении, в то время как для эффективного предобучения на разнообразных данных желательно использовать все имеющуюся в распоряжении емкость сети - то есть обучение должно быть высокоранговым.

В данной статье авторы предлагают метод последовательного обучения низкоранговых добавок к весам линейных слоев нейронной сети с последующим их слиянием с основными весами. И как утверждается, подобная процедура для достаточно больших сетей (самая большая обученная сеть имеет 350M параметров - сущий пустяк по современным меркам), работает ненамного хуже стандартной полноранговой процедуры обучения.

Метод

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

Однако, чтобы метод заработал, авторам пришлось учесть ряд нюансов и применить пару трюков.

Во-первых, используемый при обучении трансформеров Adam хранит скользящие статистики градиентов, и при переходе к обучению новой низкоранговой добавки, если не предпринимать никаких действий, оптимизация будет проводиться в том же подпространстве, что и у предыдущей LoRA добавки, нивелируя всякий смысл в итеративной процедуре. Для предотвращения такого сценария, авторы зануляют 99% состояний оптимизатора с меньшей абсолютной величиной (почему не все? почему не любую другую долю?) при инициализации новой добавки.

Кроме того, learning rate в момент начала обучения новой добавки зануляется и потом быстро разогревается до примерно того же значения, с которым закончила обучение прошлая добавка (используется cosine annealing learning rate). Без короткой warmup фазы обучение расходится.

Предложенная cтратегия именуется ReLoRA.