ModelSoups: варим суп из моделей

Введение

Для начала вспомним слова известного шеф-повара Гордона Рамзи:

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

Обычно мы подготавливаем данные (основной ингредиент), затем выбираем архитектуру и приправляем эти компоненты дополнительными частями (лосс-функция, оптимизатор и др). А перед подачей сервируем итоговый набор нужными значениями гиперпараметров. И правда, это все выглядит как приготовление еды.

Но что если составить композицию из уже приготовленных блюд? Авторы статьи “Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time” описали свой подход и провели эксперименты по усреднению весов моделей. Сделано это было для создания единственной модели с наилучшими характеристиками по качеству и устойчивости.

Зачем усреднять веса моделей?

Ранее в статье про эффективные ансамбли мы выделили задачу объединения нескольких моделей в ансамбль для улучшения следующих характеристик:

1. Качество решения задачи (качество предсказания);
2. Оценка неопределенности предсказания (можем идентифицировать данные не из распределения + избежать overconfidence модели);
3. Калибровка оценок (приближаем корреляцию точности и оцененной вероятности к 1).

Однако применение ансамблей, как правило, ведет к увеличению времени исполнения. А это для многих приложений существенно ограничит их использование.

Идеальная техника ансамблирования с точки зрения затрат на вычислительные ресурсы — техника, которая позволила бы получить идентичные ансамблю свойства при количестве вычислений, равным одной модели. Есть способ решения подобной задачи: усреднение весов множества моделей. В рамках этого направления исследований можно выделить общий подход:

1. Используем много обученных моделей или чекпоинты, полученные в рамках одного обучения на разных эпохах.
2. Подбираем множество комплиментарных моделей (при усреднении их весов получается наилучшая модель).
3. Выбираем и применяем один из алгоритмов усреднения.

После усреднения у полученной модели улучшаются следующие характеристики:

1. Качество работы на валидационных данных.
2. Устойчивость к сдвигам в распределении и out-of-distribution примерам.

В рамках описанных выше шагов каждое исследование предлагает уникальный подход или развивает уже существующий.

Далее в статье мы рассмотрим:

• подход из статьи Model Soups;
• результаты других авторов и ответ на вопрос: “Почему и при каких условиях данный метод работает?”;
• применение подхода в различных практических задачах.

Как варить суп из моделей, или описание подхода

Задача

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

1. Находим интересующую нас архитектуру и чекпоинт, полученный в режиме supervised обучения на одном из больших наборов данных / в self-supervised режиме.
2. Меняем финальный линейный слой модели на слой, подходящий к задаче.
3. Определяем набор гиперпараметров обучения — настройки оптимизатора, планировщика LR , а также подходящий под задачу набор аугментаций.
4. Делаем finetuning под нашу задачу. Можно дообучить только веса линейного слоя, заморозив прочие параметры, или же произвести end-to-end дообучение модели с настройкой всех имеющихся параметров.

Как правило, при обучении происходит конфигурация множества параметров, влияющих на результаты. Если архитектура и датасет зафиксированы — на начальном этапе варьируются следующие параметры:

• выбор LR, weight decay, параметров планировщика LR;
• выбор набора аугментаций.

После подбора этих параметров вручную / через один из фреймворков по их оптимизации мы получаем множество уникальных запусков (см. Рисунок 1) со своим значением качества на валидационной выборке. Затем мы выбираем финальный чекпоинт.

Авторов интересует вопрос: “Можно ли путем усреднения полученных чекпоинтов получить более робастную и качественную модель?”.

Дизайн эксперимента

Для проверки данной гипотезы был составлен следующий дизайн эксперимента:

1. В качестве предтренированной модели авторы берут CLIP ViT-B/32 (OpenAI).
2. В качестве целевого набор данных — ImageNet.
3. В качестве наборов данных для проверки устойчивости — ImageNet-V2, R, Scetch, A и набор данных ObjectNet.
4. Также используют end-to-end как вариант дообучения (настраивая все имеющиеся параметры).
5. Устанавливают случайный seed и порядок данных.
6. Семплируют случайные значения гиперпараметров для:
   • LR;
   • Weight decay;
   • Augmentation strength;
   • Mixup;
   • Label smoothing.

Чуть позже мы поговорим о влиянии значений гиперпараметров на качество финальной модели после усреднения. А сейчас рассмотрим методы усреднения от авторов статьи. Каждый такой метод лаконично назван Recipe (рецепт).

Рецепты, или алгоритмы усреднения весов моделей

Авторы во время основных экспериментов используют следующий набор рецептов:

• Uniform soup \( f(x,\frac{1}{k}\sum^{k}{i=1}\theta{i}) \) — average весов всех полученных моделей;
• Greedy soup \( \textbf{Recipe 1} \) — выборочное усреднение;
• Learned soup \( \textbf{Recipe 2} \) — усреднение по выученным коэффициентам.

Рассмотрим детальнее Greedy soup:

---

\( \textbf{Recipe 1 : GreedySoup} \)

\( \textbf{Input}: \text{Набор чекпоинтов (ингредиентов супа)}:\ \{{\theta_{1}}, …, {\theta_{k}}\} \)

\( \text{где}\ \{{\theta_{1}}, …, {\theta_{k}}\} — \text{отсортированы по убыванию точности на валидации} \) \( \text{ValAcc}({\theta_{i}}) \)

\( \text{ingredients} ← \{\} — \text{финальный набор чекпоинтов} \)

\( \textbf{for}\ \text{i = 1}\ \textbf{to}\ \text{k}\ \textbf{do} \) :

\( \textbf{if}\ \text{ValAcc}(\text{average}(\text{ingredients} ∪ {\theta_{i}})) ≥ \text{ValAcc}(\text{average}(\text{ingredients})) \)

\( \text{ingredients} ← \text{ingredients} ∪ {\theta_{i}} \)

\( \textbf{return}\ \text{average}(\text{ingredients}) \)

---

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

Преимущества алгоритма:

• он проверяет все полученные чекпоинты;
• чекпоинты, которые ухудшают итоговый результат, не попадают в “суп”: минимальный результат по точности итоговой модели равен точности наилучшей модели после обучения.

Недостатки алгоритма:

• стратегия выбора субоптимальна, ведь мы не исследуем другие последовательности проверки чекпоинтов, кроме перебора по уменьшению точности;
• чекпоинты, при добавлении которых точность финальной модели растет, могут и не улучшить точность на out-of-distribution данных.

В качестве усовершенствованного варианта авторы предлагают дополнительный рецепт:

---

\( \textbf{Recipe 2 : LearnedSoup} \)

\( \textbf{Input}: \text{Набор чекпоинтов (ингредиентов супа)}:\ \{{\theta_{1}}, …, {\theta_{k}}\} \)

\( \text{Обучаемый вектор весов}: W^{1\times k}_{soup}, \beta — \text{Параметр температуры}, \\E — \text{количество итераций} \)

\( {(x_{j}, y_{j})}^n_{j=1} — \text{ валидационный датасет} \)

\( \textbf{for}\ \text{e = 1}\ \textbf{to}\ \text{E}\ \textbf{do} \) :

\( \alpha^{e} = softmax(W^{e}_{soup}) \) — весовые коэффициенты линейной комбинации моделей

\( \mathit{scores} = f(x_{j},\sum^{k}{i=1}\alpha^e{i}\theta_{i}) \) — ответы модели, полученной с весами \( \alpha^{e} \)

\( l^e = \sum^n_{j=1}L(β · \mathit{scores},y_{j}) \) — ошибка на всех примерах валидационного датасета

\( W^{e+1}{soup} = {W^{e}{soup}} -\eta\frac{\partial l^e}{\partial W^{e}_{soup}} \) , где \( \eta \) — скорость обучения

---

При подборе параметров \( W^{e}{soup} \) параметры \( \{{\theta{1}}, …, {\theta_{k}}\} \) заморожены. Еще авторы сообщают: в качестве альтернативного подхода можно подбирать веса не для всей модели, а для каждого отдельного слоя.

Стоит отметить: подобный метод очень похож на подбор температуры во время калибровки модели при работе с ансамблями.

Итак, мы рассмотрели предложенные методы усреднения весов. Давайте перейдем к результатам экспериментов.

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

Посмотрим на таблицу ниже — здесь представлены сравнительные результаты различных рецептов усреднения:

На основе этих результатов мы можем отметить следующее:

• Uniform soup гарантирует наилучшую устойчивость среди рецептов, но ухудшает качество на основной задаче;
• Greedy soup дает 2 результат с наилучшей последовательностью выбора чекпоинтов;
• Порядок выбора моделей в Greedy soup влияет на итоговый результат со значением дисперсии в 0.05 и 0.16 соответственно;
• Наилучшие результаты получаются с помощью Learned Soup (послойного), но он требует подбора NxM параметров, где N — количество моделей кандидатов, а M — количество слоев.

Результаты обучения с различными гиперпараметрами:

В каждой ячейке здесь представлено значение формулы \( Acc(\frac{1}{2}\theta_{a} + \frac{1}{2}\theta_{b}) — max_{\theta} \in \{{\theta_{1}}, …, {\theta_{k}}\}\ Acc(\theta) \) — разницы между супом из 2-х моделей и наилучшей индивидуальной моделью. Каждая из моделей определена своим значением гиперпараметра.

На графике видим:

• при выборе параметров аугментации наилучшие результаты получают комбинации из модели с простыми аугментациями (strenght = 1) и параметром (strenght ≥ 2);
• при выборе скорости обучения — комбинации, где разница в скорости между моделями порядка 10 раз, но оба значения находятся в диапазоне от 3e-5 до 1e -7.

Результаты усреднения моделей, дообученных на различных датасетах:

1. В рамках одного из экспериментов авторы выбрали набор целевых датасетов — ImageNet, CIFAR-10, Food101, SUN397, DTD, Cars;
2. На каждом из них дообучили модель;
3. Сделали усреднение одним из рецептов между начальной инициализацией (CLIP) и обученными моделями;
4. Измерили, как добавление каждой модели в “суп” повлияло на качество работы в наборе данных CIFAR-100 (в режиме zero-shot классификации с помощью матчинга эмбеддингов текстовых описаний классов и изображений).

Этот эксперимент интересен, поскольку показывает, как улучшить CLIP на любом другом наборе данных при помощи только доступных (например, публичных).

Результаты обучения при меньшем начальном наборе данных:

Разберемся с результатами подобного подхода при использовании меньшего начального датасета — ImageNet22k (см. Рисунок 5). Здесь авторы в качестве модели используют ViT-B/32, предтренированый на наборе данных ImageNet-22k, а в качестве целевого датасета — ImageNet.

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

• при Greede soup получается наилучшая по качеству модель на ImageNet, но ее устойчивость существенно ниже результата индивидуального дообучения;
• при Uniform Soup — наилучшая модель по соотношению качества / устойчивости, но максимальные точность и устойчивость характерны для других моделей;
• при ImageNet22k в качестве начального датасета + supervised предобучения результаты у подхода в целом ниже, чем при CLIP.

Почему метод работает?

Итак, в предыдущем разделе мы описали алгоритм улучшения характеристик нейронной сети за счет усреднения весов. И тут возникают вопросы:

1. Почему после усреднения весов нескольких нейронных сетей качество улучшается, а не равняется нулю?
2. Как проверить комплиментарность полученных весов?
3. Какие настройки обучения и гиперпараметры позволяют применять данный метод?

Что ж, обо всем по порядку 🙂

Свойства ландшафта функции потерь

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

Обучение нейронной сети — это не что иное, как подбор параметров на ландшафте функции потерь. Мы ищем такую комбинацию значений весов, где достигается наименьшее (по возможности) значение функции потерь.

Здесь мы имеем дело с невыпуклой оптимизацией. Каждый слой нейронной сети сам по себе — выпуклая функция. Набор операций можно представить так:

Но комбинация этих операций невыпуклая. Поэтому мы сталкиваемся с таким видом ландшафта функции потерь:

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

В дополнение приведем ряд заключений из теоретического анализа в статье “The Loss Surfaces of Multilayer Networks”:

• для нейронных сетей с большим количеством параметров многие локальные минимумы являются эквивалентными и обеспечивают схожую точность на тестовом множестве;
• вероятность попадания в плохой локальный минимум снижается с ростом размера модели;
• попытка найти “глобальный минимум” на тренировочных данных не стоит свеч: это может привести к переобучению (другими словами, нам не столь уж интересен loss = 0, полученный на тренировочных данных).

Что такое плохой локальный минимум?

Во время обучения модели оптимизатор находит минимумы, которые можно разделить на гладкие (flat) и острые (sharp):

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

Исследованию свойств ландшафта функции потерь посвящено много работ [1,2,3,4,5,6,7]. Нас же сейчас интересует возможность связать, то есть усреднить полученные модели после обучения.

В статье “Taxonomizing local versus global structure in neural network loss landscapes” авторы показывают различные типы локальных минимумов. Дизайн эксперимента (архитектура, количество данных, гиперпараметры и др.) здесь обуславливает попадание в них.

Подобные локальные минимумы в статье Model Soups авторы называют Basin (впадина).

Есть вариации локальных минимумов:

1. Globally poorly-connected
   1. Phase I — высокие значения ошибки, в окрестности текущей точки значение ошибки сильно колеблется, различные базины плохо связаны;
   2. Phase III — низкие значения ошибки, колебания в окрестности текущей точки несущественны, различные базины все еще плохо связаны.
2. Globally well-connected
   1. Phase II — высокие значения ошибки, колебания в окрестности текущей точки присутствуют, различные базины связаны, но на пути встречается флуктуация значения ошибки;
   2. Phase IV-A,B — низкие значения ошибки, колебания в окрестности текущей точки несущественны, различные базины хорошо связаны между собой.

Авторы Model Soups отмечают: их подход работает, если все модели лежат в рамках одной базины (графически это похоже на варианты Phase IV-A и Phase IV-B). Для лучшего визуального представления на рисунке ниже показаны ландшафты функции потерь и бирюзовые точки — полученные комбинации весов.

На обеих визуализациях мы видим: между областями с низким значением ошибки есть область в виде барьера. Для нее характерен кратно больший уровень ошибки. Также можно заметить, что этот барьер отличается на 1 и 2 графике, поскольку визуализации относятся к разным настройкам MLP. А еще на левом графике ряд запусков попадает в одну и ту же базину, что относит его к кейсу Phase IV-B.

Чтобы разобраться, как в пространстве функции потерь располагается итоговая точка после обучения (определенная комбинация значений весов нейронной сети), обратимся к результатам двух работ: “Deep Ensembles: A Loss Landscape Perspective” и “Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel”.

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

Авторы запускают обучение модели с разной начальной инициализацией (см. Рисунок 10).

При запуске с различной инициализацией обе модели сходятся к идентичным(практически) значениям ошибки, но с точки зрения функционального пространства полученные функции отличаются. Случайная инициализация приводит к расположению модели в разных базинах. К таким же результатам пришли и в статье “Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs”, где провели 3 независимых запуска обучения модели со случайной начальной инициализацией. Визуализировав полученное пространство, авторы получили следующий результат:

Результат обучения каждой индивидуальной модели приводит к схожим значениям ошибки, но с точки зрения расположения мы находимся в различных базинах. Исходя из этого можно сделать следующий вывод: при работе с методом Model Soups необходимо использовать идентичную начальную инициализацию, иначе мы не получим нужного результата при усреднении весов моделей.

Теперь рассмотрим результаты работы “Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel”. В ней авторы исследуют эволюцию нейронной сети с помощью Neural Tangent Kernel. Помимо экспериментов со случайной инициализацией параметров авторы также рассматривают режим, при котором у нескольких нейронных сетей идентична начальная инициализация. Они представляют результаты экспериментов следующим изображением:

Мы видим два симплекса с расположенными на них точками. Границы симплекса (черная линия) — регионы с низким значением ошибки, белым цветом указан регион с высоким значением ошибки.

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

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

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

Авторы запускают обучение в режиме, похожем на Model Soups. Они стартуют от синей точки и семплируют на различных итерациях обучения по две модели в зеленых точках. Семплированные модели доучиваются до сходимости, их обозначают красными точками на симплексе. Как показано на изображении выше, для ранней итерации создания моделей (первая зеленая точка) они, как и в эксперименте А, находятся на разных границах симплекса. Но если семплирование произошло позже — полученные точки будут находиться на одной границе симплекса. Тогда между ними можно сделать линейную интерполяцию, при которой значение ошибки во всех точках на пути будет низким.

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

На верхней строке рисунка показан результат расположения моделей в функциональном пространстве при условии раннего семплирования. Авторы описывают это как “early chaotic”. На нижней строке представлен результат в фазе “late stable”. При “early chaotic” модели сходятся к разным базинам, а при “late stable” — к одной.

Изучив теоретические исследования, мы приходим к следующим выводам:

• нужна идентичная начальная инициализация (в целом следуем подходу из Model Soups);
• нужна достаточно обученная модель, с которой мы инициализируемся (некоторые эмпирические результаты по этому вопросу рассмотрим далее);
• для получения успешных результатов при усреднении весов моделей также нужна настройка параметров обучения, ведь они влияют на те минимумы, в которые попадет модель.

Таким образом, в этом разделе мы поговорили о связанности весов и возможности их связанности линейным и нелинейным путями. В статье Model Soups авторы опираются на подход Linear mode connectivity, чтобы анализировать связанность результатов различных запусков обучения модели.

Поэтому в следующей части мы рассмотрим такие понятия, как Mode connectivity и Linear mode connectivity, а также эмпирические результаты использования начальной инициализации, полученной на разных эпохах.

Mode connectivity и Linear mode connectivity: анализируем связанность весов

Mode connectivity

В рамках работы с глубокими нейронными сетями анализ свойств связанности весов различных моделей был представлен в статье «Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs» и «Essentially No Barriers in Neural Network Energy Landscape». Здесь авторы предложили методы для поиска пути между независимо обученными нейронными сетями, находящимися в разных базинах. Изначальная постановка задачи требовала реализации такого алгоритма поиска пути, при котором на всем его протяжении значение мат ожидания ошибки будет минимальным. Используя Bezier и Polychain кривую, авторы за счет процесса смещения весов вдоль этих кривых приходят к следующим результатам (подробнее об этом — в видео от автора статьи):

Результаты двух работ отлично применимы для быстрого построения ансамблей. Подробнее об использовании таких путей между весами и для построения ансамблей можно прочитать в статье “Learning Neural Network Subspaces”.

При всем успехе данных методов для работы Model Soups нужна линейная связанность весов. Поэтому давайте обратимся к Linear mode connectivity.

Linear mode connectivity

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

А если мы инициализировали все модели идентично (как, например, в Model Soups) и захотели понять, сходятся ли они в одну базину или принадлежат к разным? Авторы статьи «Linear Mode Connectivity and the Lottery Ticket Hypothesis» формулируют задачу анализа подобных свойств следующим образом:

Пусть нам дано:

• \( \mathit{N} \) — модель
• \( \mathit{W} \) — инициализированные веса модели
• \( \mathit{SGD} \) — оптимизатор
• \( \mathit{U} \) — распределение шума (сюда входит набор агументации, последовательность батчей и другие параметры, которые мы можем задать случайно)

Задаем вопрос:

Какова устойчивость \( \mathit{SGD} \) к случайному шуму, семплированному из \( \mathit{U} \) ? Здесь под устойчивостью подразумевается отсутствие возрастания значений функции потерь при интерполяции между весами \( \mathit{W}_1 \) и \( \mathit{W}_2 \), полученными при идентичной инициализации, но с разными семплами шума из \( \mathit{U} \).

Для измерения данного свойства авторы вычисляют ошибку в точках на промежутке [0,1].

Рассмотрим эксперименты авторов, которые проливают свет на границы работоспособности Model Soups.

1. Результаты обучения моделей с 0: у них идентичная случайная инциализация.

На рисунке выше приведен график интерполяции для двух моделей одной архитектуры, обученных на разных датасетах. X — ошибка при выбранном значении \( a \), Y — значение коэффициента интерполяции.

Согласно результатам исследования получаем следующее:

• для сети LeNet и датасета MNIST оптимизация с SGD является устойчивой к различным шумам и приводит к идентичной базине;
• для всех остальных наборов данных и нейронных сетей обучение с общей случайной инициализацией сводится к различным базинам.

Затем авторы исследовали другой подход: за начальные веса \( \mathit{W} \) берется чекпоинт после k итераций обучения, и каждая из сетей доучивается независимо. Результаты следующие:

На рисунке выше показан график интерполяции для двух моделей одной архитектуры, обученных с чекпоинта на k-той итерации. X — Instability (величина мат ожидания ошибки на пути между моделями), Y — значение k (итерация создания начального чекпоинта).

Согласно результатам исследования получаем следующее:

• для сети LeNet и датасета MNIST значение k не влияет на устойчивость к шуму;
• для набора данных CIFAR-10 и сетей ResNet-20 и VGG-16 устойчивость возникает, если модели инциализируются с весов при k ≥ 2000 для ResNet и k≥1000 для VGG (что эквивалентно 3 и 1.5 процентов от всего обучения при batch size = 128);
• для набора данных ImageNet для сети ResNet-50 устойчивость возникает с 18 эпохи (20% от всего обучения (90 эпох) при batch size = 1024), для Inception-v3 — с 28 эпохи (16%, количество эпох равно 171).

Результаты экспериментов показывают, на какой итерации модель переходит из “early chaotic” к “late stable”.

Исходя из полученных эмпирических результатов, можно сделать следующие выводы, которые позволяют определить требования к инициализации в Model Soups:

1. Для крошечных наборов данных (MNIST) даже случайная инициализация позволит получить линейно связанные веса.
2. С ростом количества примеров в наборе данных растет и требуемое количество эпох обучения (начиная от 3% для CIFAR-10 и до 20% на ImageNet).

Теперь, обладая знаниями о свойствах работы Model Soups, взглянем на известный метод SWA и поймем, в чем связь и отличие данных методов (ранее у нас выходила статья Weight Averaging про базовое описание SWA и имплементацию в коде).

Model Soups VS Stohastic Weight Averaging. А в чем разница?

Кратко суммируем SWA. Перед стартом нам понадобится:

• \( \mathit{N} \) — модель
• \( \mathit{W} \) — проинициализированные веса модели
• \( \mathit{SGD} \) — оптимизатор
• \( \mathit{Cycle\ lenght} \) — длина цикла при циклическом расписании LR (при константном = 1)
• \( \mathit{E_n} \) — количество эпох дообучения

Именно с ранее предобученной модели (75% от времени тренировки для CIFAR-10 и СIFAR-100 , для Imagenet с предтренированного чекпоинта из torchvision) авторы начинают процесс SWA. Проиллюстрировать два разных варианта его настройки можно следующей схемой:

При каждом окончании цикла (на изображении он отмечен оранжевыми точками) авторы применяют следующую формулу для усреднения весов:

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

Давайте кратко рассмотрим свойства циклической скорости обучения — наиболее часто применяемой скорости при запуске SWA.

Cyclical learning rate

Влияние циклического расписания на результаты подробно описаны в статьях «Snapshot Ensembles: Train 1, get M for free” и “Exploring loss function topology with cyclical learning rates”.

Циклическое расписание скорости обучения позволяет эффективно исследовать ландшафт функции потерь как локально (в рамках одной базины), так и более глобально (при поиске новых базин).

Рассмотрим глобальный кейс:

Как видно из рисунка выше (визуализация правой части), изменение скорости обучения в начале каждого нового цикла приводит к переходу в другое пространство ландшафта функции потерь (в другую базину) и до конца цикла, пока скорость обучения снижается, модель постепенно сходится к минимуму этой базины (изменение ошибки показано направлением стрелки). С помощью такого приема можно получать множество независимых чекпоинтов для формирования ансамбля, как это сделали авторы “Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs” и «Snapshot Ensembles: Train 1, get M for free”. После анализа связанности данных весов через Linear mode connectivity авторы статьи “Exploring loss function topology with cyclical learning rates”. получили следующий результат:

На графиках выше показано сравнение пути из W1 и W2, полученных на разных итерациях при стандартном (правая часть) и циклическом (левая часть) LR.

При стандартном обучении на интервале от 0 до 1 наблюдается консистентное значение ошибки. В режиме циклической тренировки мы имеем резкое возрастание ошибки при a = 0.5, что сигнализирует о нахождении весов в разных базинах. Тут может возникнуть вопрос: почему при усреднении по таким чекпоинтам в SWA у нас получается итоговый рост в качестве? Есть ряд ключевых компонентов, которые позволяют перейти из глобального режима в локальный:

• старт с модели, обученной достаточное количество эпох;
• короткий цикл — в SWA используется цикл из 5 эпох, а в SSE — от 20 до 50 эпох (в зависимости от нейронной сети);
• значение LR для цикла — наилучшие результаты SWA достигает при значениях максимального и минимального LR = 5 · 10−2 и 5 · 10−4 соответственно, тогда как в Snapshot Ensembles используется 1 * 10-1 в качестве максимального.

Как видно из рисунка выше, значение LR существенно влияет на результаты SWA.

Итак, остановимся на финальном сравнении SWA и Model Soups:

Model Soups

• стартуем с модели, обученной достаточное количество эпох
• важно, чтобы начальная модель достаточно сошлась
• N независимых дообучений
• исследование базины за счет случайного шума (порядок батчей, аугментации, значение LR)
• есть ограничение на значения гиперпараметров

SWA

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

Таким образом, методы очень схожи, за исключением подхода к исследованию базины и поиску новых решений. Как было сказано ранее, выбор LR существенно влияет на динамику перемещения в рамках базины и получение более разнообразных чекпоинтов.

Финальный вопрос: как выбрать скорость обучения для Model Soups? Для ответа обратимся к результатам статьи “To Stay or Not to Stay in the Pre-train Basin: Insights on Ensembling in Transfer Learning”, авторы которой провели множество реальных экспериментов как по анализу ансамблей в режиме тренировки с предварительным обучением модели, так и по анализу работы Model Soups. Они помогут понять, какой режим тренировки будет оптимальным для получения наилучших результатов при усреднении весов моделей.

Исследуем пространство базины: влияние настроек обучения

Начнем с краткого определения задачи. Авторы формулируют ее в виде следующего вопроса: какой подход к дообучению нейронной сети нужно применить, чтобы из N дообученных сетей получить наилучшее качество ансамбля моделей? При этом предполагается найти такой подход, который мог бы позволить использовать минимальное количество GPU часов, потраченных на обучение подобного ансамбля. Авторы рассматривают два подхода к получению подобных моделей (см. Рисунок 22):

1. SSE — равен варианту из SWA и Snapshot Ensebles (применяем циклическое расписание LR для дообучения модели, чекпоинты сохраняются в конце каждого цикла);
2. StarSSE — для каждой модели используется свой независимый запуск дообучения. Авторы характеризуют его как параллельную версию SSE. Этот вариант максимально схож с Model Soups.

Переходим к практическим результатам метода Model Soups и получаем следующее:

На графиках представлены эксперименты по применению Model Soups для чекпоинтов, полученных с помощью SSE и StarSSE для модели CLIP. При использовании параллельного циклического дообучения (StarSSE) итоговое усреднение дает наилучшую точность, и все полученные модели улучшают итоговую точность.

Теперь разберемся с некоторыми обозначениями на графиках:

• Local DE soup — Model Soups из моделей, обученных при случайной инициализации последнего (линейного) слоя с идентичными гиперпараметрами.
• Soup size — количество индивидуальных моделей, усредненных между собой с помощью Uniform soup.

Далее рассмотрим обозначения more loca, optimal и more semi-local exp. Их можно представить в виде следующей таблицы:

Каждый из режимов представляет собой множители гиперпараметров циклического расписания обучения. Рассмотрим 2D-проекцию получившихся решений и их расположение в пространстве базины:

На графиках мы видим расположения полученных весов и значения функции потерь для них. Эксперименты с большими циклами и высоким LR дают модели, расположенные на границе базины (semi-local). Их точность хуже тех, где использовались маленькие (local) и оптимальные (optimal) значения параметров LR.

Результаты работы добавляют следующие пункты в пайплайн использования Model Soups:

• для получения оптимальных чекпоинтов нужно использовать OneCycle с начальным значением LR от x2 до x4 при снижении количества эпох до 0.25 от их количества;
• Star-SSE можно использовать как drop-in замену SWA в пайплайне дообучения и получать лучшие результаты;
• авторы не варьировали параметры аугментаций, mixup-добавление этих переменных вместе со Star-SSE может дополнительно улучшить результат Model Soups.

Промежуточные выводы

Стоит отметить: приведенные выше результаты направлены на формирование “интуиции” о принципах работы Model Soups и границах применимости метода. Многие результаты получены эмпирическим путем, поэтому относиться к ним нужно с долей скептицизма.

Итак, давайте закрепим финальные выводы по разделу:

1. Случайная инициализация приводит модели в разные участки функционального пространства. Их не получится усреднить.
2. Для получения весов, подходящих Model Soups, нужна общая начальная инциализация.
3. Эта инициализация должна быть обучена достаточное количество эпох (здесь все зависит от размера сети и сложности датасета).
4. Модели, подходящие Model Soups, — это модели в одной базине. При линейной интерполяции между ними не должно возникать существенного увеличения значения ошибки во всех точка интерполяции.
5. SWA — аналог Model Soups, но с другим подходом к исследованию базины.
6. Для получения оптимальных чекпоинтов необходимо использовать OneCycle LR sheduler с начальным значением LR от x2 до x4 от базового, снизив при этом количество эпох обучения (минимально до 0.25 от начального).
7. Применение StarSSE вместе с вариацией пайплайна аугментаций может быть бейзлайном при построении Model Soups.

Использование Model Soups в различных задачах

Улучшение робастности при файнтюнинге foundation models

Для знакомства с этой задачей можно почитать две статьи: “Robust fine-tuning of zero-shot models” и “Patching open-vocabulary models by interpolating weights”. Мы кратко рассмотрим первую из них.

Основная идея методов: получить более робастную модель за счет усреднения весов начальной (zero-shot) и дообученной на конкретной задаче модели. Поскольку у модели уже есть неплохая точность из коробки, а дообучение двигает ее веса только в рамках базины — между такими моделями есть линейная связь. Давайте рассмотрим, как авторы используют подобные свойства в методе Robust fine-tuning of zero-shot models.

Формулируют они задачу так: пусть в качестве in distribution набора данных выступает ImageNet, а в качестве distribution shift данных (со специальными сдвигами в распределении) — наборы данных, идентичные Model Soups.

Вопрос: “Может ли усреднение весов начальной и дообученной модели дать наилучшую устойчивость на distribution shift при хорошем качестве на in distribution?”.

Анализ результатов пайплайна:

1. В качестве начальной модели берется CLIP.
2. Модель файнтюнится на ImageNet (в end-to-end режиме).
3. С помощью выражения \( \mathcal{\theta_{a}} = (1-a) *\theta_{zero-shot} + a * \theta_{fine-tuned} \) и выбранного значения коэффициента \( a \) формируется итоговая модель.
4. Далее измеряется точность полученной модели in distribution и distribution shift данных.

Обозначения на графике выше:

• оранжевый кубик — значение точности дообученной CLIP модели в end-to-end режиме;
• синяя линия — точность всевозможных supervised моделей на ImageNet;
• фиолетовая линия — точность всевозможных CLIP моделей, где дообучались только веса линейного слоя;
• розовая линия — результаты усреднения при различном коэффициенте \( a \).

Мы видим: для всевозможных CLIP моделей при дообучении только линейного слоя точность может быть выше, однако начальная робастность к distribution shift ниже в сравнении с end-to-end дообучением.

Результаты Model Soups при различном коэффициенте \)a \):

• при небольших значениях \( a \) качество растет на всех датасетах;
• при значении \( a \), близком к середине, получается наилучшее distribution shift качество при идентичном in distibution.

Model Soups между такими моделями дает следующие численные результаты:

Возможность гибко настраивать поведение LLM после стадии RLHF

Современный пайплайн обучения LLM предполагает две основных стадии:

1. Начальную предтренировку на огромных наборах данных из интернета.
2. Дообучение модели с помощью RL в задаче выполнения инструкций.

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

Следовательно, возникает следующая проблема: если на стадии RLHF мы не учли что-то в рамках функции вознаграждения — для части пользователей поведение модели и ее ответы не будут оптимальными. Для возможности более гибкой настройки модели после обучения авторы “Personalized Soups: Personalized Large Language Model Alignment via Post-hoc Parameter Merging” придумали свой подход: за счет усреднения весов после обучения с помощью Model Soups мы решаем задачу более тонкой настройки модели согласно предпочтениям конкретного пользователя. Схематично это выглядит так:

Идея метода:

1. Сначала делаем процедуру RLHF, формулируя Reward так, чтобы итоговое поведение модели удовлетворяло среднему пользователю на планете, получаем General веса модели.
2. Далее под каждый интересующий нас вариант поведения формулируем особый Reward и также получаем финальные (персонализированные) веса.
3. В зависимости от требований пользователя усредняем веса разных персонализированных моделей и General, используя при этом весовые коэффициенты для контроля вклада каждой модели.

Таким образом, мы получаем удобное Post-hoc решение, при котором можем расширять базу персонализированных весов и добавлять больше различных настроек для персонализации модели под пользователя.

Адаптация LLM под новую задачу без обучения

Сегодня для тюнинга LLM и других гигантских моделей под конкретную пользовательскую задачу наиболее эффективно используется LoRA для обучения только небольшой добавки к существующим весам модели.

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

Важно отметить: речь идет именно про zero-shot работу модели, ведь использование множества примеров “запрос-ответ” внутри входного промта — наиболее простое решение подобной задачи. Однако оно требует постоянного процессинга дополнительных токенов перед запросом пользователя.

Подобная парадигма работы со множеством ранее обученных адаптеров представлена в работе “LoraHub: Efficient Cross-Task Generalization via Dynamic LoRA Composition”.

Авторы предлагают новую парадигму адаптации LLM без обучения:

Авторы формулируют начальные условия:

1. Есть заранее выбранная LLM для решения задачи;
2. Есть хаб с множеством уникальных LoRA параметров, полученных при дообучении под различные датасеты;
3. Есть набор пар “запрос-ответ” в качестве примеров запросов на инференсе модели.

Задача: использовать множество уникальных LoRA параметров для решения новой задачи.

Алгоритм решения:

1. Создадим вектор весовых коэффициентов под каждый из адаптеров;
2. Усредним все адаптеры с заданными весовыми коэффициентами (аналогично Model Soups);
3. Добавим полученную добавку к весам модели;
4. Оценим ошибку модели при текущих весах для каждого адаптера.

Затем мы повторяем все шаги и изменяем при этом вектор весов таким образом, чтобы достичь наименьшей ошибки на имеющихся у нас примерах. Процесс похож на Learned Soup (о нем мы рассказывали в первом разделе), за исключением проведения оптимизации с помощью неградиентного метода (https://facebookresearch.github.io/nevergrad/).

Улучшение доменной генерализации при семантической сегментации

Идея получить модель, которая будет хорошо работать сразу с несколькими доменами, не нова: в этом направлении каждый год появляются разные методы. В статье “A Re-Parameterized Vision Transformer (ReVT) for Domain-Generalized Semantic Segmentation” представлен подход, где Model Soups используется для получения модели, хорошо работающей сразу на нескольких набора данных.

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

• выбираем сеть для сегментации (в экспериментах авторов это Segformer B2, B3, B5);
• берем в качестве начального набора данных синтетический датасет GTA5 для задачи семантической сегментации в домене автономных машин;
• берем в качестве целевых наборов данных для валидации модели Cityscapes Mapillary Vistas, BDD100k, ACDC, KITTI из того же общего домена.

Обучение модели по следующему сценарию:

1. Инициализируем энкодер весами модели, предтренированной на ImageNet.
2. Инициализируем декодер случайным образом.
3. Используем специфический набор аугментаций под каждую отдельную модель.

Пайплайн аугментаций выглядит так:

Фактически модели получают разный по силе набор аугментаций:

• базовые — Resize, Random Crop, Flip;
• усиленные — PhotoAug, Bilateral Filter;
• наиболее сильная аугментация — PixMix.

Подробнее о настройках обучения и аугментаций можно посмотреть в Supplementary Material, расположенном после источников в статье на arxiv.

Пайплайн для усреднения следующий:

1. Выбираем три модели комбинации, которые дают наилучшее качество;
2. Усредняем их;
3. В качестве декодера выбираем один из имеющихся согласно метрикам.

В сравнении с другими моделями мы получаем следующие результаты (здесь они показаны только для Segformer-B2):

Опираясь на данные таблицы, мы можем отметить:

• этот вариант обучения гораздо лучше бейзлайна на всех целевых наборах данных (в ряде случаев прирост около 5% mIOU);
• лучшие метрики показывает комбинация из модели с базовыми аугментациями, одной сильной аугментацией и моделью со средними и сильными аугментациями (что повторяет эмпирические результаты Model Soups);
• решение помогает обходить сети с более сложными пайплайнами обучения (например, в SAN+SAW используют различные преобразования над пространством фичей, в WildNet — контрастивное обучение вместе со стилизацией).

Применение в неструктурированном прунинге

В рамках этой задачи Model Soups используют для улучшения точности спарсифицированной модели. В статье “Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging” авторы предложили применять усреднение весов моделей для улучшения результатов процесса неструктурированного прунинга. Кстати, о прунинге мы писали в одной из наших статей 😊

Для прунинга авторы применили алгоритм Iterative Magnitude Pruning (IMP). Он состоит из следующих шагов:

1. Начинаем с ранее обученной модели;
2. Для всех весов в каждом из сверточных / линейных слоев вычисляем L1 норму;
3. Зануляем заранее заданный % наименьших по L1 норме весов для каждого из слоев;
4. Переобуваем сеть.

Весь процесс визуально выглядит так:

Авторы предложили для алгоритма дополнение в виде Model Soups — оно позволило улучшить итоговые результаты прунинга. Идея следующая: производить не одно дообучение модели, а несколько параллельных, стартующих от единой инициализации, но с разными настройками обучения. Это повторяет процесс Model Soups. По завершению обучения веса моделей усредняются. Схематически это выглядит так:

Кроме того, авторы написали псевдокод алгоритма:

---

\( \textbf{Algorithm 1} \)

\( \textbf{Input}: \text{Pretrained model}\ \theta \)

\( \textbf{Output}: \text{Sparse model soup} \)

\( \textbf{for}\ \text{each prune-retrain cycle}\ \textbf{do} \) :

\( \textbf{Prune}\ \theta \)

\( \textbf{for}\ i ← 1\ \textbf{to}\ m\ \textbf{do} \)

\( ▷\ \text{Fully parallelizable}\ \theta_{i} ← \theta​ \)

\( \textbf{Retrain}\ \theta_{i}\ \text{with specific hyperparameters} \)

\( \textbf{end for​} \)

\( \textbf{return​}\ \theta​ \)​

Авторы рассматривали качество модели на разных стадиях, каждая из которых закреплена за определенным уровнем спарсификации (количество зануленных весов). Еще они варьировали параметр mm — количество параллельных обучений, которое будет запускаться на каждой итерации тренировки после прунинга. Авторы проверили два рецепта усреднения: Greedy и Uniform. Получили следующие результаты:

Опираясь на данные таблицы, мы можем отметить:

По сравнению с метриками IMP, авторам удалось улучшить результаты приблизительно на 1%. Для задачи прунинга — это существенное улучшение.

Применение 10 параллельных обучений обеспечивает наилучшие результаты на всех уровнях спарсификации.

Рецепт Uniform предпочтителен для 1 и 3 фазы, Greedy — для 2 фазы.

Выводы по разделу

Здесь мы рассмотрели пять примеров использования Model Soups в различных задачах. Сразу перейдем к выводам:

Подход Model Soups популярен, он применяется в различных областях: CV и NLP. Все статьи, к которым мы обращались в этом разделе, датированы 2021-2023 годами, что подтверждает его актуальность.

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

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

Заключение

Итак, в статье мы подробно рассмотрели подход Model Soups. А точнее — познакомились со следующими аспектами:

В первой части разобрали содержание статьи о Model Soups и предложенные авторами алгоритмы усреднения весов, выделили основные эмпирические результаты.

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

В третьей части рассмотрели примеры практического использования Model Soups в различных задачах. Выделили направления применения (CV и NLP) и, наконец, разобрали основные идеи, стоящие за представленными алгоритмами.