SwitchLoRA: Переключаемая адаптация низкого ранга может изучать информацию полного ранга

Основная информация

• ID статьи: 2406.06564v3
• Название: SwitchLoRA: Switched Low-Rank Adaptation Can Learn Full-Rank Information
• Авторы: Kaiye Zhou, Shucheng Wang, Jun Xu (China Mobile (Suzhou) Software Technology Co. Ltd.)
• Классификация: cs.LG, cs.AI, cs.CL
• Дата публикации: 2 января 2025 г. (препринт arXiv)
• Ссылка на статью: https://arxiv.org/abs/2406.06564v3

Аннотация

В данной статье предлагается метод SwitchLoRA для решения проблемы параметрически эффективного обучения на этапе предварительной подготовки больших языковых моделей. Традиционные методы адаптации низкого ранга, такие как LoRA, показывают отличные результаты на этапе тонкой настройки, но их прямое применение к предварительной подготовке приводит к значительному снижению производительности. Существующие методы ReLoRA и GaLore пытаются решить эту проблему путём обновления подпространства низкого ранга, но всё ещё не достигают точности обучения полного ранга. SwitchLoRA решает эту проблему путём частой и плавной замены обучаемых параметров адаптеров LoRA, постепенно обновляя подпространство низкого ранга, каждый раз затрагивая только несколько измерений, чтобы минимизировать влияние на состояние оптимизатора. Экспериментальные результаты показывают, что SwitchLoRA снижает недоумение на модели LLaMA 1.3B с 15.23 до 15.01, превосходя обучение полного ранга, одновременно снижая коммуникационные издержки на 54% и использование памяти на 13%.

Предпосылки и мотивация исследования

Основная проблема

С появлением архитектуры Transformer размер больших языковых моделей резко увеличился, и распределённое обучение триллионных моделей сталкивается с огромными коммуникационными издержками между узлами. Хотя параметрически эффективные методы, такие как LoRA, показывают отличные результаты на этапе тонкой настройки, их прямое применение к этапу предварительной подготовки приводит к значительному снижению производительности.

Ограничения существующих методов

1. ReLoRA: Для сохранения согласованности состояния оптимизатора ограничивает частоту обновлений, не может полностью приблизиться к поведению обучения полного ранга
2. GaLore: Полагается на SVD-аппроксимацию пространства полного ранга, вводя потери точности в процессе аппроксимации

Мотивация исследования

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

Основные вклады

1. Предложение метода SwitchLoRA: Путём частой плавной корректировки обучаемых параметров матриц LoRA снижается использование памяти и коммуникационные издержки при сохранении точности обучения полного ранга
2. Стратегия управления состоянием оптимизатора: Разработаны механизмы сброса и временной заморозки состояния оптимизатора при переключении параметров, снижающие влияние несогласованности состояния
3. Улучшенные правила инициализации: Предложены новые стратегии инициализации для параметров адаптеров LoRA и их кандидатных векторов, повышающие эффективность обучения
4. Комплексная экспериментальная проверка: Метод проверен на моделях LLaMA различных размеров, а способность к выводу проверена с помощью эталонного теста GLUE

Подробное описание метода

Определение задачи

Для матрицы весов $W \in \mathbb{R}^{m \times n}$ предварительно обученной модели традиционный LoRA преобразует её в $W + \frac{\alpha}{r}BA$, где $B \in \mathbb{R}^{m \times r}$, $A \in \mathbb{R}^{r \times n}$, $r \ll \min(m,n)$. SwitchLoRA динамически переключает векторы в B и A для увеличения эффективного ранга.

Архитектура модели

Основной механизм переключения

1. Разложение векторов: Матрица B разлагается на векторы-столбцы $b_k \in \mathbb{R}^{m \times 1}$, матрица A разлагается на векторы-строки $a_k^T \in \mathbb{R}^{1 \times n}$
2. Набор кандидатных векторов: Поддерживаются наборы кандидатных векторов $C(B)$ и $C(A^T)$, содержащие $\min(m,n)$ векторов
3. Динамическое переключение: На этапе обучения $b_k$ и $a_k$ заменяются на кандидатные векторы $b_k' \in C(B)$ и $a_k' \in C(A^T)$

Стратегия корректировки весов

При замене векторов соответствующим образом корректируется матрица весов: $W \leftarrow W + b_k a_k^T - b_k' a_k'^T$

Проектирование частоты переключения

Используется функция экспоненциального затухания: $frequency = Ce^{-\theta \cdot step}$, отражающая естественную эволюцию модели от полного ранга к низкому рангу.

Технические инновации

1. Минимизация влияния на состояние оптимизатора

• При переключении $a_k$ сбрасывается состояние оптимизатора для $b_k$
• При переключении $b_k$ сбрасывается состояние оптимизатора для $a_k$
• После сброса соответствующие параметры временно замораживаются на N шагов (N=5)

2. Улучшенная стратегия инициализации

На основе идей инициализации Xavier и Kaiming разработано новое стандартное отклонение: $std[B] = std[b] = \left(\frac{r}{\sqrt{mn}}\right)^{\frac{1}{4}} gain^{\frac{1}{2}}$$std[A] = std[a] = \left(\frac{\sqrt{mr}}{\sqrt{nn}}\right)^{\frac{1}{4}} gain^{\frac{1}{2}}$

3. Оптимизация памяти

Резервные кандидатные векторы выгружаются на CPU с использованием неблокирующей передачи для параллельной обработки процесса переключения.

Экспериментальная установка

Наборы данных

• Предварительная подготовка: Набор данных C4, использование первых 46M обучающих образцов и полного набора валидации
• Оценка: Оценка потерь валидации на 10M токенов каждые 1000 шагов
• Тонкая настройка: Множественные задачи эталонного теста GLUE

Конфигурация модели

Эксперименты охватывают модели LLaMA различных размеров:

• 130M (768 измерений, 12 голов, 12 слоёв)
• 250M (768 измерений, 16 голов, 24 слоя)
• 350M (1024 измерения, 16 голов, 24 слоя)
• 1.3B (2048 измерений, 32 головы, 24 слоя)

Метрики оценки

• Предварительная подготовка: Недоумение (Perplexity)
• Тонкая настройка: Точность, коэффициент корреляции Пирсона, коэффициент корреляции Мэтьюса

Методы сравнения

• Обучение полного ранга
• Стандартный LoRA
• ReLoRA
• GaLore

Детали реализации

• Оптимизатор: Adam (β₁=0.9, β₂=0.999)
• График обучения: Косинусное отжигание, 100 шагов разминки
• Общее количество шагов обучения: 40,000
• Оборудование: 8×NVIDIA A800 80GB PCIe GPU

Результаты экспериментов

Основные результаты

Сравнение производительности предварительной подготовки

Результаты недоумения на модели 1.3B:

• Full-rank: 15.23
• SwitchLoRA (rank=512): 15.01 (превосходит обучение полного ранга)
• SwitchLoRA (rank=256): 15.89

Сравнение потребления ресурсов

На примере модели 1.3B:

• Использование памяти: Снижение на 13% по сравнению с обучением полного ранга (36.1GB → 31.9GB)
• Коммуникационные издержки: Снижение на 54% (обучаемые параметры с 1339M до 610M)
• Время обучения: Примерно одинаково (21.6s vs 22.5s)

Сравнение с существующими методами

vs ReLoRA

На модели 250M, когда ReLoRA использует 5000 шагов предварительной подготовки полного ранга, а SwitchLoRA использует только 200 шагов, SwitchLoRA всё ещё показывает лучшие результаты. При одинаковых 1000 шагах предварительной подготовки полного ранга SwitchLoRA значительно превосходит ReLoRA.

vs GaLore

На модели 350M:

• GaLore: 20.29 недоумения
• SwitchLoRA: 19.58 недоумения

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

Абляционные исследования

Влияние частоты переключения

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

Влияние количества шагов заморозки

Выбор количества шагов заморозки N влияет на результаты обучения, N=5 является оптимальной установкой.

Проверка стратегии инициализации

Новый метод инициализации значительно улучшает скорость сходимости по сравнению с традиционной инициализацией LoRA.

Проверка способности к выводу

Результаты эталонного теста GLUE

На модели 350M:

• Модель предварительной подготовки SwitchLoRA в среднем на 3.0 балла выше, чем модель предварительной подготовки GaLore
• На 0.3 балла выше, чем модель предварительной подготовки полного ранга

На модели 1.3B:

• Модель предварительной подготовки SwitchLoRA в среднем на ~1.0 балл выше, чем модель предварительной подготовки полного ранга

Экспериментальные находки

1. Анализ распределения ранга: Распределение сингулярных значений SwitchLoRA более близко к обучению полного ранга, тогда как стандартный LoRA показывает патологическое распределение
2. Эффект масштаба: По мере увеличения размера модели преимущество SwitchLoRA над стандартным LoRA становится более выраженным
3. Способность к обобщению: Модели, предварительно подготовленные с помощью SwitchLoRA, демонстрируют более сильные способности к выводу и обобщению на нижестоящих задачах

Связанные работы

Методы разложения низкого ранга

Ранние работы реализовали низкоранговую аппроксимацию матриц весов с помощью методов, таких как SVD, в основном применяемых к CNN и небольшим языковым моделям.

Варианты LoRA

• Класс объединения параметров: Chain of LoRA, ReLoRA увеличивают эффективный ранг путём периодического объединения параметров
• Класс улучшения инициализации: Улучшение стратегий инициализации матриц B и A и настроек скорости обучения
• Класс модификации структуры: Изменение процесса обучения LoRA и механизма обновления параметров

Другие методы сжатия

Включая квантизацию, обрезку, сжатие градиентов и другие методы, где GaLore реализует эффективное по памяти обучение путём проекции градиентов.

Заключение и обсуждение

Основные выводы

1. Прорыв в производительности: SwitchLoRA впервые достигает производительности, превосходящей обучение полного ранга на этапе предварительной подготовки
2. Ресурсная эффективность: Значительно снижает использование памяти и коммуникационные издержки при сохранении сопоставимого времени обучения
3. Улучшенная обобщаемость: Модели предварительной подготовки демонстрируют более сильные способности к выводу на нижестоящих задачах

Ограничения

1. Чувствительность к гиперпараметрам: Гиперпараметры, такие как частота переключения, требуют тщательной настройки
2. Выбор ранга: Для достижения точности обучения полного ранга по-прежнему требуется достаточно большой ранг LoRA
3. Выбор кандидатных векторов: В настоящее время используется случайный или последовательный выбор, что может оставить место для оптимизации

Направления будущих исследований

1. Адаптивная частота: Разработка более интеллектуальных стратегий корректировки частоты переключения
2. Оптимизация на уровне слоёв: Разработка дифференцированных стратегий переключения для различных типов слоёв (матрицы Q, K, V)
3. Оптимизация выбора кандидатных векторов: Исследование более эффективных стратегий выбора и обновления кандидатных векторов

Глубокая оценка

Преимущества

1. Теоретическая инновация: Предлагает новый подход к постепенному обновлению подпространства низкого ранга, эффективно решая проблему низкоранговой подготовки на этапе предварительной подготовки
2. Инженерная реализация: Тщательно рассматривает управление состоянием оптимизатора, оптимизацию памяти и другие практические вопросы, обладая сильной практической ценностью
3. Полнота экспериментов: Проверяет эффективность метода с множества углов, включая производительность предварительной подготовки, потребление ресурсов, способность к выводу и т.д.
4. Теоретический анализ: Предоставляет теоретические объяснения независимости обновления векторов, обоснованности сброса состояния оптимизатора и т.д.

Недостатки

1. Увеличение сложности: Увеличивает сложность реализации по сравнению со стандартным LoRA, требуя дополнительного управления кандидатными векторами
2. Настройка гиперпараметров: Множество гиперпараметров (частота переключения, скорость затухания, количество шагов заморозки) требуют тщательной настройки
3. Проверка масштаба: Хотя протестированы модели различных размеров, максимум составляет 7B, применимость к более крупным моделям требует дальнейшей проверки
4. Полнота теории: Хотя предоставлен некоторый теоретический анализ, отсутствует глубокое теоретическое объяснение того, почему метод может превосходить обучение полного ранга

Влияние

1. Академическая ценность: Предоставляет новый подход к параметрически эффективному обучению на этапе предварительной подготовки, потенциально вдохновляя дальнейшие исследования в этой области
2. Практическая ценность: Значительно снижает потребление ресурсов при сохранении производительности, имеет важное значение для практического обучения больших моделей
3. Воспроизводимость: Статья предоставляет подробные детали реализации и настройки гиперпараметров, облегчая воспроизведение и применение

Применимые сценарии

1. Предварительная подготовка больших моделей: Особенно подходит для сценариев с ограниченными ресурсами, но требующих высококачественной предварительной подготовки
2. Распределённое обучение: Может значительно снизить коммуникационные издержки при обучении на нескольких узлах
3. Дополнительное обучение: Подходит для сценариев, требующих продолжения обучения на основе предварительной подготовки

Библиография

Статья цитирует большое количество связанных работ, включая:

• Hu et al. 2022: Оригинальная статья LoRA
• Lialin et al. 2023: Метод ReLoRA
• Zhao et al. 2024: Метод GaLore
• Vaswani et al. 2017: Архитектура Transformer
• Rajbhandari et al. 2020: Оптимизатор ZeRO

Общая оценка: Это высококачественная исследовательская статья, демонстрирующая отличные результаты в теоретической инновации, экспериментальной проверке и практической ценности. Метод SwitchLoRA умело решает проблему низкоранговой подготовки на этапе предварительной подготовки, не только сохраняя эффективность обучения, но и достигая значительной экономии ресурсов. Хотя существуют некоторые ограничения, её вклад достаточен для продвижения развития этой области.