SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining
[Статья] [Код]

Введение

Низкоранговые приближения хорошо себя зарекомендовали при дообучении моделей, но их выразительности не хватает при обучении модели с нуля. Ранее было предложено по ходу обучения вливать LoRA в веса модели (статья ReLoRA) и инициализировать низкоранговую добавку, или накладывать низкоранговость на градиенты и состояния оптимизатора (GaLoRE). Но в первом случае, для получения качества, близкому к бейзлайну был необходим этап обучения всей модели целиком.

Возникает вопрос:

возможно ли добиться хорошего качества, имея всегда малое подмножество обучаемых параметров?

И авторы предлагают обучать Low-Rank + Sparse компоненту.

Метод

Идея, естесна, далеко не нова. По существу, это своего рода RobustPCA. Для дообучения LLM ранее подобное реализовали в RoSA.

Авторы мотивируют использование Sparse компоненты в дополнение к Low-Rank, тем что она обладает несколько иными свойствами (при не слишком малой sparsity матрица почти всегда полноранговая) и вместе они могут более гибко фитировать исходные матрицы весов.

Сэмплируется рандомная sparse маска и фиксируется по ходу обучения. Sparse матрица прибавляется к Low-Rank компоненте, потому sparse matrix операции не требуются при такой постановке, а для получения градиента достаточно проиндексировать градиент по всему весу.

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

Авторы обучают семейство моделей с архитектурой Llama от 60M до 7B параметров на С4 корпусе на бюджете в несколько миллиардов токенов.

В качестве бейзлайнов выступают:
1️⃣ Full-Rank обучение
2️⃣ Low-Rank training
3️⃣ ReLoRA
4️⃣GaLoRE

Качество оценивают по валидационной перплексии.

SLTrain по качеству чуть лучше ReLoRA и GaLoRE, несколько экономичнее по памяти GaLoRE, и примерно такой же по скорости.

Рассматривают разные конфигуации Sparsity/Low-Rank. 1-3% ненулевых весов в Sparse компоненте более менее оптимально.

Анализируя спектр выученных весов, они показывают, что первые собственные значения Full-Rank обученных матриц соотвествуют Low-Rank компоненте, а последующие - Sparse компоненте. То есть, сумма Low-Rank + Sparse более точно описывает спектр при полном обучении.

Вывод

Идея разумная и логичная, но непонятно, насколько идея масштабируется на более практически интересные сценарии. Отсутствуют даже замеры 0-шотов на lm-eval (по всей видимости, при таких бюджетах обучения они мало будут отличаться от random guess).