Мой коллега Никита Габдуллин работает в Отделе перспективных исследований ИТ-компании «Криптонит». Он автор библиотеки Loss Landscape Analysis (LLA) и научной статьи о ней, препринт которой доступен на английском языке. Здесь мы подробнее рассказываем о самой библиотеке, в то время как научная публикация в основном посвящена исследованию разных свойств ландшафтов функций потерь. Также в русскоязычной статье мы допускаем некоторые лингвистические вольности, которые не приняты в академической среде, но упрощают восприятие текста.

При работе с нейросетями-классификаторами у всех на слуху какие-то известные архитектуры, которые характеризуются числом параметров, скоростью вычислений (инференса), точностью выполнения той или иной известной задачи. Популярны соревнования, посвящённые тому, насколько точно можно решить задачу классификации на типовых датасетах, и часто борьба уже идёт за доли процента [PWC]. Однако в реальных задачах нейросети часто показывают себя куда хуже, чем в «лабораторных» условиях, что переводит акццнт внимания с тренировочных и тестовых (train-test) задач на проверку обобщающей способности (generalization) нейросетей.

В наших работах мы столкнулись с тем, что нейросети одного типа могут иметь практически идентичные показатели train-test, но демонстрировать кардинально отличающиеся результаты на датасетах, отличных от тренировочного. Без углублённого анализа непонятно, за счёт чего возникают такие эффекты. Поэтому для таких нейросетей очень сложно выполнить оценку их реальной обобщающей способности. Это вдохновило нас на поиски методов, которые позволили бы проанализировать обобщающую способность нейросети с теми или иными весами, среди которых метод построения ландшафта функции потерь (loss landscape) показался интересным кандидатом.

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

Это послужило мотивацией для разработки библиотеки и проведения теоретической работы, результаты которой были недавно опубликованы на ArXiv [lla]. Библиотека, о которой речь пойдёт в данной статье, доступна на GitHub [llagh] и GitFlic [llagc]. Сразу хочу сказать, что изначально не стояло цели делать свою библиотеку.  Нам хотелось найти готовую, чтобы поскорее получить результаты для своих задач. Однако этого достичь не удалось, и в какой-то момент стало понятно, что с ростом числа модификаций существовавших библиотек стоит доработать полученный код до новой библиотеки с улучшенным и обновлённым функционалом.

1. Откуда берутся ландшафты функций потерь?

По традиции введение в данную тему начинается с картинки из оригинальной статьи [LLO], репродуцированной на Рис 1. На ней изображён ландшафт функции потерь ResNet56 на датасете Cifar10. Видно, что Рис 1 (а) и (б) сильно отличаются при отсутствии и наличии транзитных соединений (skip connections) в архитектуре. В чём же отличие между этими двумя картинками и поведением нейросетей, которым они соответствуют? Попробуем разобраться.

Данные ландшафты получены следующим образом: для фиксированных данных определённым образом немного варьируются веса модели с целью определить, какой эффект это будет иметь на выход нейросети — значение целевой функции потерь. Для этого выбираются два случайных ортогональных вектора размерности весов, которые и будут использоваться для изменения весов (уравнения можно найти в разделе 3.2). При этом если малые изменения весов приводят к значительным изменениям значений функции потерь, как на Рис. 1 (а), можно говорить о нестабильности сети или даже о невозможности её обучить. А сеть с ландшафтом Рис. 1 (б) как раз куда лучше обучается и обобщается, что и показали авторы оригинальной статьи.

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

2. Устройство библиотеки Loss Landscape Analysis

Наша библиотека Loss Landscape Analysis (LLA) в значительной мере опирается на две существующие, но ныне заброшенные их авторами библиотеки на PyTorch:

Loss landscapes [llpg] — довольно удобная библиотека для построения ландшафтов функций потерь, клон-версия которой запускается «из коробки», но только для простых задач (что подразумевается под «простой задачей» — уточним ниже). Есть pip-версия, которая часто крашится в современных версиях PyTorch.

PyHessian [PHg] — библиотека, которая содержит реализацию ряда методов для оценки параметров гессиана нейросетей. Научная работа по данной теме описана в статье [PHa]. Тема того, что такое гессиан функции потерь нейросети и зачем его анализировать не входит в рамки данной статьи. Читателю рекомендуется ознакомиться с [PHa] и разделом 3 в [lla].

Эти библиотеки нами постепенно дорабатывались исходя из нужд наших проектов — фиксились баги, добавлялись дополнительные методики анализа и визуализации и т.п. Однако по ходу работы стало понятно, что мы представляем не новую версию существовавшей библиотеки, а создаём что-то новое, что и стало библиотекой Loss Landscape Analysis.

2.1  Использование LLA

Наша библиотека содержит одну директорию src_lla, которую необходимо и достаточно поместить в свой проект для импорта функций библиотеки. Также прилагаются скрипты на Python и Jupyter notebooks с примерами и тестами для типовых нейросетевых задач. В LLA все важные операции собраны в две функции viz_lla и viz_esd. viz_lla фокусируется на построении ландшафтов функции потерь с разнообразными параметрами и режимами, а  viz_esd — на анализе гессиана. Списки всех аргументов этих функций можно найти в гите.

Для работы viz_lla нужно несколько элементов, которые есть в любом проекте по нейросетям в PyTorch: объект нейросетевой модели model (torch.nn.Module), загрузчик батчей данных dataloader (torch.utils.data.DataLoader), функция потерь criterion и оптимизатор optimizer (torch.optim). Также нужно указать данные, по которым будет строиться ландшафт функции потерь (по умолчанию — один батч из dataloader’a), и один объект типа metric, который мы сразу назовём CustomLoss.

from src_lla import viz_lla, metrics
x_plot, y_plot = iter(data_loader).__next__()
metric = CustomLoss(x_plot, y_plot, device)

Цель metric — придать анализу гибкость для применения в самых разных нейросетевых задачах. Этот класс ввёл автор loss landscapes [llpg] с целью решить ряд проблем в сложных случаях, о которых пойдёт речь ниже. При этом, для простых случаев, когда функция потерь — это какая-то библиотечная функция со стандартной реализацией, вроде Cross Entropy или Mean Squared Error, metric может быть простой обёрткой вокруг функции потерь, например:

criterion = torch.nn.CrossEntropyLoss()
metric = metrics.Loss(criterion, x_plot, y_plot, device)
viz_lla(model,metric)

И всё! Этого достаточно для того, чтобы запустить lla для типовых задач типа LeNet на MNIST или ResNet на ImageNet. В библиотеку входит Jupyter notebook под названием import_example, где представлен пример того, как импортировать и использовать viz_lla и viz_esd своём проекте на примере LeNet на MNIST.

2.2 Многообразие задач и вытекающие из этого проблемы

Пока всё, что было описано выше, можно было сделать силами библиотеки loss landscapes [llpg]. Однако есть нюанс. Чтобы разобраться в нём, остановимся на минуту и ещё раз посмотрим на то, что и как мы вообще считаем. Для построения ландшафта функции потерь фактически нужно две вещи: задаться правилом, по которому будут меняться веса модели, и иметь возможность посчитать значения функции потерь, сделав инференс модели на заданных данных.

Проблема в том, что все необходимые операции, как они были реализованы в loss landscapes, предполагали максимально простой формат вычислений. Например, предполагалось, что инференс модели проводится в виде pred = model(x), из-за чего в функционале библиотеки не было поддержки флагов, поддержки вывода данных из промежуточных модулей (что бывает необходимо для вычисления значений некоторых функций потерь), были жёсткие требования к формату данных и т. д. Так что, на практике для каждой более-менее специфической задачи требовалось лезть и что-то менять в коде библиотеки, а это слишком непрактично — нельзя требовать от юзера лезть «с отвёрткой» внутрь библиотеки при каждом небольшом отклонении от самых базовых режимов работы.

Чтобы осознать «масштабы бедствия» ниже в таблице приведена сводка по ряду проектов и тому, является ли та или иная часть общего пайплайна стандартной. Если для задачи в Таблице 1 есть по крайней мере одно «нет», то нужно править что-то в коде (в случае loss landscapes). Хотя список моделей и задач далеко не полный, видно, что в стандартную схему ложатся только самые базовые задачи классификации, а что-то более сложное, вроде парсинга изображений [schp], уже требует модификаций. Также проблемы возникают с современными архитектурами типа Variational Autoencoder (VAE) или Visual Transformer (VIT).

*VAE часто применяется с функциями потерь, которые оперируют непосредственно над скрытым пространством, и могут быть как простыми (l1loss), так и сложносоставными.

Довольно быстро стало понятно, что хоть функции библиотеки и можно сделать более гибкими — была добавлена поддержка флагов, поддержка выводов промежуточных значений для некоторых модулей, поддержка вычислений на GPU — но в общем случае такое решение ненадёжно. Поэтому было принято решение сделать несколько объектов, в которых юзер может указать все нюансы своей модели, данных, и правил вычисления функции потерь, цель которых — свести все input/output к одной стандартной форме. Иными словами, накладываются жёсткие ограничения на формат вывода каждого элемента, и это гарантирует, что все вычисления в viz_lla и viz_esd отработают корректно.

2.3 Стандартные элементы loaders

Для выполнения описанной выше задачи «в сложном случае» предлагаются loaders. Юзеру требуется указать всё необходимое для реализации своего «сложного случая» в loader и потом просто импортировать его в скрипт. В src_lla.loaders есть template.py и заполненные лоадеры для LeNet и ResNet. Импортировать лоадер можно так:

from src_lla.loaders.{your_loader} import *

В LLA входят скрипты на Python для стандартных режимов тренировки (train) и инференса (eval). С eval всё просто и скрипт отработает после импорта в него корректно заполненного loader’a. В train прописан самый базовый train loop, и если требуется более сложный, то скрипт нужно модифицировать. Кроме .py версий прилагаются Jupyter notebooks с аналогичным функционалом.

3. Функционал библиотеки lla

Как было сказано выше, все функции LLA собраны в viz_lla и viz_esd. Функция viz_lla позволяет рассчитать ландшафт функции потерь по случайным осям, осям-моментам Адама, осям-собственным векторам гессиана для разных типов нормализации, разных уравнений модификации весов, с или без заморозки отдельных слоёв. Также можно посчитать спектр гессиана с дефолтными настройками. Функция viz_esd позволяет оценить собственные значения и собственные вектора гессиана, его след, построить его спектр. Также возможна оценка гессиана определённых блоков и слоёв нейросети.

3.1. Выбор осей для построения ландшафта функции потерь: аргумент axes

Данный аргумент позволяет выбрать тип векторов для осей, по которым будут строиться ландшафты функции потерь. Поддерживаются три типа осей: random, adam, hessian.

На Рис. 2 (а) показано, что картина ландшафта функции потерь, построенная по случайным осям, может быть довольно простой. Однако если для того же случая смотреть ландшафты по осям Адама или Гессиана, то картина резко меняется. Если в первом случае ландшафт уже выглядит довольно плоским, что говорит о хорошем качестве нейросети, то для других осей все неравномерности, которые помешают хорошему обобщению, становятся куда более очевидны.

3.2. Выбор уравнения апдейта весов: аргумент mode

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

где α и β — коэффициенты при векторах осей δ и η; θ* — начальное значение весов и L — функция потерь. Это уравнение используется для mode add. Коэффициенты меняются в зависимости от значения аргумента steps, который по умолчанию равен 40, при этом начальное значение весов (когда α и β равны нулю) соответствует центру построенного ландшафта, по умолчанию точке (20,20).

Однако если мы добавляем возможность использовать оси-моменты Adam, то и для апдейта весов логично будет использовать те же правила, что использует Adam [PyA]. Это возможно с mode adameq

где γ — это learning rate; ϵ — малое число для стабилизации вычислений, а δ и η будут первым и вторым моментами Adam. На Рис. 3 видно, что ландшафт, построенный по осям-моментам Адама в режиме adameq, обладает неровностью, которая не видна при построении в режиме add.

3.3. Выбор типа нормализации: аргумент norm

Несложно заметить, что эффект от прибавления векторов случайных величин к матрице весов будет зависеть от соотношения значений в векторах и весах. Чтобы минимизировать этот эффект часто используется нормализацию между векторами и весами, которую можно включить, используя аргумент norm. Допустимые значения norm это None, weight, layer, model, filter. Для нормализаций кроме None и weight есть аргумент order, который позволяет выбрать L1 (1) или L2 (2) нормализацию.

Выбор типа нормализации — тема сложная, поэтому читателю предлагается ознакомиться с разделами 2.2 и 2.3 в нашей статье на ArXiv [lla]. Здесь лишь укажем, что ландшафты потерь могут очень сильно меняться в зависимости от типа нормализации, как показано на Рис. 4  для ResNet, обученном на ImageNet. У viz_lla также есть аргумент viz_dev (all_modes для скриптов на Python), который позволяет построить ландшафты по ряду заранее заданных настроек, в том числе с разными типами нормализации для их сравнения.

Нужно отметить, что при построении ландшафтов диапазон значений функции потерь ResNet очень велик, что усложняет анализ в случае, если не задаться максимальным значением функции потерь для визуализации. В нашей библиотеке это можно сделать с помощью аргумента losscap, который для Рис 4 был задан равным 10. Такое поведение ранее было замечено авторами [LLOg].

3.4. «Заморозка» весов в некоторых слоях нейросети: аргумент freeze

Для определённых архитектур может быть интересно «заморозить» веса некоторых слоёв при построении ландшафта функции потерь, чтобы увидеть, как те или иные слои влияют на стабильность. Это отчасти мотивировано тем, что в нейросетях сложных архитектур могут быть целые модули с принципиально разными назначениями и свойствами, как, например, в задачах парсинга изображений [schp].

В данный момент freeze принимает положительный или отрицательный int и «замораживает» веса во всех слоях после указанного (положительный int) или до него (отрицательный int). На Рис. 5 приведён пример ландшафтов функции потерь с заморозкой для ResNet на ImageNet. Видно, что причиной всей неравномерности являются основные слои, в то время как классификационный слой совершенно гладкий. Это логично, так как классификатор — это простой полносвязный слой выход которого практически является линейной комбинацией параметров входа.  Как уже было сказано ранее, более интересные результаты можно получить, замораживая большие участки нейросетей со специфическими функциями.

3.5 Вычисление спектра Гессиана: флаг hessian

Флаг hessian позволяет рассчитать спектр гессиана. Это важно, так как анализ параметров гессиана позволяет получить информацию о состоянии нейросети, которую сложно извлечь, анализируя ландшафты функций потерь. На Рис. 6 приведён пример спектров гессианов LeNet на разных этапах обучения на MNIST.

На Рис. 6 видно, что отрицательная часть спектра постепенно убывает с ростом точности при обучении. Приятно, что этому даже есть более-менее строгое математическое обоснование [EP]. Важно отметить, что корреляция между видом спектра и точностью сохраняется не только для режима обучения и теста на одном датасете, но и при использовании обученной сети на других датасетах. Это позволило нам сформулировать критерии оценки гессиана и показать, что существует корреляция между изменениями в критериях и изменением точности при смене датасета для предобученных сетей, что позволяет оценить обобщающую способность нейросети. Подробности читатель может найти в разделах 3.3 и 4 в [lla].

Заключение

Говоря об анализе ландшафтов функции потерь у нейросетей, важно понимать, какие варианты можно получить и на что именно обращать внимание. Чтобы упростить жизнь себе и другим разработчикам, мы написали библиотеку LLA. Она распространяется с открытым исходным кодом и предоставляет возможность более гибкого подхода как к построению, так и к последующему анализу обобщающей способности нейросетей. С нашей библиотекой становится доступен выбор осей, методов нормализации, анализ спектра гессиана и другие продвинутые функции. Работа над библиотекой будет продолжаться. Мы надеемся, что она заинтересует и других исследователей в области машинного обучения, чей фидбек позволит улучшить и расширить функционал LLA в будущем.

Ссылки:

[LLA] https://arxiv.org/abs/2412.10146

[llagh] https://github.com/GabdullinN/loss-landscape-analysis

[llagc] https://gitflic.ru/project/kryptodpi/loss-landscape-analysis

[llpg] https://github.com/marcellodebernardi/loss-landscapes/tree/master

[PHg] https://github.com/amirgholami/PyHessian

[LLcom] https://losslandscape.com/

[PWC] https://paperswithcode.com/task/classification-1

[LLO] https://arxiv.org/pdf/1712.09913

[LLOg] https://github.com/tomgoldstein/loss-landscape

[Hllt] https://habr.com/ru/companies/skillfactory/articles/536606/

[PHa] https://arxiv.org/abs/1912.07145

[schp] https://arxiv.org/abs/1910.09777

[PyA] https://pytorch.org/docs/stable/generated/torch.optim.Adam.html

[EP] https://arxiv.org/abs/1910.05929