Назад
573

Введение в файн-тюнинг Stable Diffusion

573

Генеративные модели

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

Под капотом такие модели оценивают апостериорную вероятность наблюдения данных \( P(x) \). Это знание они используют для генерации новых сэмплов.

Часто генеративные модели строятся таким образом, чтобы иметь возможность включать некоторое условие в сам процесс генерации. Следовательно, они оценивают условную апостериорную вероятность \( P(x \mid y) \), где \( y \) — репрезентация некоторого условного сигнала.

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

Давайте остановимся на моделях, генерирующих изображения. Их классификация следующая:

  • VAE, или Variational Autoencoder — вариационный автокодировщик. Основан на энкодер-декодер архитектуре. Энкодер значительно сжимает входной сигнал в некоторое внутреннее представление; при сжатии часть информации теряется. Декодер разжимает это внутреннее представление в оригинальный сигнал. По такому принципу работает обычный автокодировщик. Чтобы его можно было назвать вариационным, в архитектуру добавляют еще более узкое горлышко — теперь выходом энкодера является не просто некоторое скрытое представление, а параметры некоторого распределения. Затем они используются декодером для сэмплирования. Самые известные представители этого семейства моделей: VQ-VAEVQ-VAE-2 и AutoEncoderKL.
  • GAN, или Generative Adversarial Network — генеративно-состязательная сеть. Состоит из двух частей: генератора и дискриминатора (иногда его называют критиком). Генератор порождает некоторый выходной пример, а дискриминатор пытается отличить его от реальных данных. Процесс тренировки таких моделей заключается в minimax состязании между этими двумя моделями. GAN’ы — очень красивый концепт, хоть и подверженный различным проблемам в сходимости. Но научное сообщество смогло разработать большое количество техник для борьбы с ними. Для обучения модели используют различные adversarial функции потерь (minimax loss / NSGANLSGANEBGANBEGANWGANWGAN-GP). Наиболее известные представители здесь: ProGANStyleGANStyleGAN2Alias-Free GAN / StyleGAN3BigGANpix2pixHD и SPADE.
  • NF, или Normalizing flow — нормализирующие потоки. Основаны на построении набора инвертируемых преобразований, которые трансформируют простое вероятностное распределение в более сложное. В отличие от других моделей, они в явном виде получают оценку распределения данных \( P(x) \), что позволяет их использовать как для генерации новых данных, так и для оценки реалистичности имеющихся примеров. Самые известные архитектуры данного семейства — RealNVPNICE и Glow.
  • Diffusion-based models — диффузионные модели. Обучаются так, чтобы итеративно удалить шум с сигнала со случайным шумом. Простота обучения и элегантность таких моделей сильно повлияла на всю генеративную сферу последних лет. Наиболее известные модели: ImageGenDALL-EDALL-E 2Kandinsky и, конечно, Stable-Diffusion.

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

DDPM

DDPM расшифровыватся как “Denoising Diffusion Probabilistic models”. Эти модели относятся к классу моделей со скрытыми латентными переменными. То есть мы предполагаем, что вместе с чистыми сэмплами \( x_0 \) существуют латенты \( x_1, \dots, x_T \) одной размерности с чистыми данными.

\( p_{\theta}\left( x_0 \right) := \int p_{\theta} \left( x_{0:T} \right) \mathrm{d} x_{1:T} \)

Для таких моделей определяются два процесса: прямой диффузный и обратный denoise процессы.

Рисунок 1. Denoising Diffusion Probabilistic Models [источник: arXiv:2006.11239]

Прямой диффузный процесс

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

\( q\left ( x_{1:T} \right ) := \prod_{t=1}^{T} q \left ( x_{t} \mid x_{t-1} \right ) \)

\( q \left ( x_{t} \mid x_{t-1} \right ) := N \left ( x_{t} ; \sqrt{1 — \beta_t} x_{t-1}, \beta_t L \right ) \)

где \( \beta_1, \dots, \beta_T \) — параметры дисперсионного расписания для гауссова шума.

Рисунок 2. Прямой диффузионный процесс

Одно из приятных свойств прямого процесса — возможность получения произвольного \( x_t \) на любом шаге \( t \) из \( x_0 \) с помощью простого переопределения:

\( \alpha_t := 1 — \beta_t \)

\( \bar{\alpha}_t := \prod_{s=1}^{t} \alpha_{s} \)

\( q \left( x_t \mid x_0 \right) = N\left( x_t; \sqrt{\bar{\alpha}_t} x_0, (1 — \bar{\alpha}_t) L\right) \)

Обратный denoise процесс

Обратный denoise процесс — буквально та часть, где происходит вся магия. Этот процесс параметризован нейронными сетями \( \mu_{\theta} \left( x_t, t \right) \) и \( \Sigma_{\theta} \left( x_t, t \right) \), чьи предсказания используются для постепенного удаления шума с входного сигнала. Это повторяется несколько раз до тех пор, пока весь шум не удалится. Следовательно, весь процесс работает как марковская цепь.

\( p_{\theta} \left( x_{0:T} \right) = p \left( x_T \right) \prod_{t=1}^{T} p_{\theta} \left( x_{t-1} \mid x_{t} \right) \)

\( p_{\theta} \left( x_{t-1} \mid {x_t} \right) = N \left( x_{t-1} ; \mu_{\theta} \left( x_t, t \right), \Sigma_{\theta} \left( x_t, t \right) \right) \)

Рисунок 3. Обратный диффузионный процесс

Обычно для большинства моделей \( \Sigma_{\theta} \left( x_t, t \right) \) не обучается, а фиксируется в \( \Sigma_{\theta} \left( x_t, t \right) = \beta_t L \). Таким образом, обучается и применяется только одна модель \( \mu_{\theta} \left( x_t, t \right) \).

Следовательно, она принимает на вход зашумленный латент (noisy latent), метку текущего временного шага, и выдает один из следующих выходов:

  1. \( \epsilon \)-предсказание — предсказанием модели является оценка всего шума на изображении.
  2. \( x_0 \)-предсказание — предсказанием модели является оценка чистого латента \( x_0 \) на текущем шаге.
  3. \( v \)-предсказние — предсказанием модели является \( v := \alpha_t \epsilon — \sigma_t x \). Такое представление позволяет существенно снизить количество шагов сэмплирования. Впервые использовались в этой статье как метод дистилляции моделей.

Эти предсказания применяются как для подсчета лосса, так и при эксплуатации модели.

На практике чаще всего используется \( \epsilon \)-предсказание. Обычно такие модели строятся по примеру UNet-подобных архитектур, но в последнее время сложился тренд на их усложнение (разбиение на несколько стадий генерации, усложнение архитектур, обучение нескольких моделей в разных разрешениях, использование каскадных подходов и др.).

Stable Diffusion

Stable Diffusion — модель генерации изображений, созданная при сотрудничестве Runway ML и CompViz group в LMU Munich с использованием ресурсов Stability AI.

Рисунок 4. High-Resolution Image Synthesis with Latent Diffusion Models [источник: arXiv:2112.10752]

Эта модель относится к классу латентных диффузных моделей. Чтобы снизить вычислительные ресурсы как для тренировки, так и для инференса, она включает в себя предобученный Variational Auto-Encoder. Он способен значительно сжимать изображение в сущность под названием “латент” — уменьшенную версию изображения. Таким образом, вместо того, чтобы оперировать в пространстве пикселей, Stable Diffusion осуществляет генерацию в латентном пространстве.

Stable Diffusion включает в себя и другие предобученные модели, например, OpenAI CLIP для кодирования текстового условия (текстового промпта).

Самая важная часть модели — это, безусловно, denoising UNet. Она состоит из downsample блоков, mid-блока и upsample блоков. Каждый из таких блоков является смесью нескольких последовательных resnet и cross-attention блоков. Последние используются, чтобы модель начала реагировать на какой-то внешний сигнал, например, на эмбеддинги текстового промпта.

Как уже упоминалось выше, denoising UNet является time-aware моделью: при ее обучении и применении модель должна быть в курсе, от какого timestep’а ей пришел текущий noisy латент. Для кодирования временной метки могут использоваться разные подходы — Positional encoding или Fourier Encoding, а для подмешивания временного сигнала — простое прибавление проекции эмбеддинга, scale-shift перенормировка, AdaGroupNorm или SpatialNorm. Stable Diffusion применяет positional encoding и простое прибавление проекции эмбеддинга.

Версии Stable Diffusion

Важно отметить: Stable Diffusion разрабатывался и эволюционировал итеративно. Сейчас существует четыре мажорные версии:

  1. Stable Diffusion v1.1 — v1.5 (CompViz репозиторийRunwayML репозиторий). Обучение проходило на подвыборке из LAION-Aesthetics V2 в разрешении 512×512. Модель использует подход с \( \epsilon \)-предсказанием.
  2. Stable Diffusion v2.0 — v2.1 (StabilityAI репозиторий). Обучение осуществлялось на подвыборке LAION-5B в разрешении 768×768. Модель применяет подход с \( v \)-предсказанием. Она становится больше, обучается с нуля и использует более тяжелую версию OpenCLIP вместо OpenAI CLIP.
  3. Stable Diffusion XL (StabilityAI репозиторийarXiv:2307.01952). Модель имеет другую архитектуру и работает в две стадии: base и refiner. Включает в себя целых 2 CLIP text encoder’а. При обучении использовались дополнительные трюки — Micro-Conditioning и SDEdit. Обучение проходило в разрешении 1024×1024. Модель снова стала применять подход с \( \epsilon \)-предсказанием.
  4. Stable Diffusion 3 (блогпост, arxiv.org:2403.03206). Сейчас модель не находится в публичном доступе, но доступна по API и показывает потрясающие результаты. Из основных изменений: заметно изменена архитектура модели (используется DiT вместо Unet и 3 text encoder’а, в процессе инференса сам text condition обрабатывается вместе с латентами параллельно, обучение базируется на концепции Rectified Flow).

Хотя базовая модель Stable Diffusion v1.5 уже достаточно древняя, она до сих пор пользуется большой популярностью в комьюнити. Существуют сотни (если не тысячи) различных чекпоинтов, совместимых с v1.5. Чтобы убедиться в этом, достаточно посмотреть на отфильтрованный список моделей на портале CivitAI.

Реализации Stable Diffusion

Есть две широко используемых реализации Stable Diffusion:

  1. Реализация LDM. Оригинальная реализация, сделанная CompViz/RunwayML/StabilityAI. Много проектов и исследовательских статей основаны на ней или используют данную имплементацию (например, k-diffusion).
  2. Реализация от diffusers. Реимплементация оригинального кода, при этом максимально совместимая со всей остальной экосистемой HuggingFace (добавлена удобная возможность загружать модели из их хаба, сериализовать и десериализовать одной строчкой).

Думаю, что реализация от diffusers больше подходит для использования в production окружении, хоть и иногда может не хватать имплементации последних разработок в этой области. Зато это хорошая возможность добавить их самостоятельно! Для этого рекомендуем ознакомиться с их гайдом для контрибьюторов.

Экосистема Stable Diffusion

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

  • AUTOMATIC1111 / Stable Diffusion web UI -> 100 тысяч звезд на GitHub. webUI позволяет вам взаимодействовать со Stable Diffusion через простой, но эффективный web-интерфейс. На первый взгляд может показаться неприметным, но он очень мощный, и его функционал легко можно расширить через внушительный список расширений.
Рисунок 5. Stable Diffusion webUI
  • ComfyUI -> 20 тысяч звезд на GitHub. Это продвинутый интерфейс — он позволяет строить пайплайн генерации изображений в виде ациклического направленного графа. Такой подход более гибкий по сравнению с webUI. Но чтобы им овладеть, нужно обладать глубоким пониманием работы генеративных моделей и других пайплайнов.
Рисунок 6. ComfyUI

Файн-тюнинг моделей генерации изображений

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

Рисунок 7. Image Generation model fine-tuning

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

Есть разные подходы для решения этой задачи: например, Textual Inversion, HyperNetworks. А еще полное или частичное дотюнивание модели (их мы тоже рассмотрим в статье 🙂).

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

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

Textual Inversion / Textual Embedding

Рисунок 8. An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion [источник: arXiv:2208.01618]

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

Таким образом, мы внедряемся в модель кодирования текста и добавляем туда новые эмбеддинги. Дообучиваем только их, оставляя другие части модели нетронутыми. На этот процесс можно смотреть как на попытку найти одно / несколько “слов”, которые наилучшим образом опишут объект или стиль.

Процесс обучения для textual inversion выглядит следующим способом:

  1. Мы добавляем новые эмбеддинги в text encoder модель и инициализируем их каким-то стартовым значением (например, близким по смыслу концепту, который мы пытаемся дообучить).
  2. Берем картиночный пример, содержащий наш концепт, и добавляем на него шум в соответствии с расписанием для шума.
  3. Затем зашумленный пример вместе с text embedding’ом подаем в UNet. Он, в свою очередь, выдает предсказанный шум. Далее мы сравниваем его с изначально добавленным шумом (например, с использованием mean-squared-error) и осуществляем backpropagation градиентов до новых добавленных эмбеддингов.

Ниже приведены результаты textual embedding’а, обученного на моих фотографиях:

Рисунок 9. Textual Inversion trained on my personal photos

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

  1. Следует использовать хорошую стартовую инициализацию для обучаемых эмбеддингов. Например, если вы пытаетесь обучить embedding на персону — начните с embedding’а какой-либо знаменитости, похожей на целевую персону. Тогда можно снизить learning rate при обучении, которое пройдет гораздо быстрее, а результаты получатся более стабильными. Тем не менее некоторые реализации для тренировки textual embedding’а (как, например, скрипты diffusers textual_inversion.py) могут не позволять указывать несколько токенов для инициализации.
  2. Иногда даже хорошая стартовая инициализация не помогает. Бывает, что текущие эмбеддинги порождают нестабильные результаты или вовсе выдают объекты, непохожие на желаемые (как на рисунке выше). На мой взгляд, это может быть связано с влиянием диффузионного лосса на обучаемые эмбеддинги при отсутствии ограничений. Я проводил несколько экспериментов, которые показали: добавление регуляризации на обучаемые эмбеддинги может помочь побороть эту проблему. В экспериментах я использовал mean-squared-error между средней нормой обучаемых эмбеддингов и каждого эмбеддинга в text encoder модели. Возможно, есть и другие варианты 😊.

Достоинства

  • Легковесный подход. При использовании textual embedding’а дообучается единственный слой в модели. Обученные эмбеддинги могут храниться отдельно от модели и подгружаться в нее, практически не занимая места (их вес варьируется от нескольких килобайт до их сотен).
  • Plug and play с различными чекпоинтами. Если несколько чекпоинтов используют один text encoder, но разные UNet’ы — обученные эмбеддинги могут легко переиспользоваться между ними.

Недостатки

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

HyperNetwork

Рисунок 10. NovelAI Improvements on Stable Diffusion [источник: Blog]

Подход HyperNetwork может быть знаком с тем же именем, который представил Google в 2016 году.

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

HyperNetwork при к Stable Diffusion отличается от подхода Google. Его представил NovelAI в одном из блог-постов. Здесь ключевая идея — модификация keys и values в Cross-Attention блоках через Multi-Layer-Perceptron (MLP), которые называются HyperNetworks.

Процесс обучения аналогичен обучению Textual Inversion (кодирование изображения в латенты, их зашумление, обработка зашумленных латентов через UNet).

Ниже приведены результаты HyperNetwork, обученной на моих фотографиях:

Рисунок 11. HyperNetwork trained on my personal photos

Достоинства

  • Схожий с файн-тюнингом подход. Он модифицирует cross-attention слои. Поэтому он схож с полным дообучением модели сильнее, чем Textual Inversion.

Недостатки

  • Плохая сходимость. Эмпирические эксперименты показывают: у HyperNetwork в применении к Stable Diffusion плохая сходимость + их сложно обучать.
  • Модификация только Cross-Attention слоев. HyperNetwork не модифицируют другие слои, такие как Linear, Conv2d.
  • Checkpoint-зависимый подход. HyperNetwork, обученная на одном базовом чекпоинте, не может быть переиспользована в другом чекпоинте из-за разницы в распределениях активаций cross-attention.

Хотя HyperNetwork не всегда показывают хорошие результаты, думаю, что подход мог бы получить больше внимания, если бы его реализация была открыта сообществу open-source, а не получена в результате утечки кода NovelAI (пост на reddit, issue в webui). В каком-то смысле HyperNetwork для Stable Diffusion схож с представленным позднее IP-Adapter, который показывает очень достойные результаты.

Dreambooth

Рисунок 12. DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation [источник: arXiv:2208.12242]

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

“Сatastrophic forgetting” — ситуация, когда модель научилась порождать требуемый стиль или объект, но забыла все / часть того, на чем училась ранее.

Чтобы справиться с этим, Google предложил простую и эффективную методику: добавить в тренировочный датасет набор регуляризирующих изображений (они могут быть порождены и самой моделью). Такой подход получил название Dreambooth. Сегодня он считается стандартом в дообучении диффузионных моделей.

Ниже представлены результаты обучения Dreambooth на моих фотографиях:

Рисунок 13. Dreambooth trained on my personal photos

Достоинства

  • Лучшие результаты. Dreambooth выдает наилучшие результаты по сравнению с остальными методами (Textual Inversion, HyperNetwork).

Недостатки

  • Тяжеловесный подход. Dreambooth тяжеловат в плане потребления GPU памяти, поскольку дообучаются все слои модели (кроме VAE). Конечно, можно дообучать только часть слоев, но достаточно сложно сказать, какие слои влияют на качество сильнее всего (про это есть отдельные исследования, например вот это).
  • Большой размер выходного чекпоинта. После обучения Dreambooth имеем выходной чекпоинт размером, равным размеру полной модели. А еще трудно перенести результаты обучения на другой базовый чекпоинт — вероятно, придется тренировать модель заново на новом чекпоинте.

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

Файн-тюнинг

Давайте проанализируем объем памяти, необходимый для дообучения Stable Diffusion через Dreambooth.

По меркам 2024 года базовый Stable Diffusion — небольшая модель с чуть более, чем 1 миллиардом параметров. Она состоит из:

  • VAE — 84 миллиона параметров;
  • UNet — 860 миллионов параметров;
  • Text Encoder — 123 миллиона параметров;

Суммарно — 1 миллиард параметров.

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

Давайте посмотрим на использование памяти при дообучении модели на потребительской видеокарте. Предположим, у нас есть видеокарта с 16 гигабайтами VRAM. Для визуализации потребления памяти я вдохновлялся прекрасным туториалом от Predibase по теме: «Efficient Fine-Tuning for Llama-v2-7b on a Single GPU». Каждый квадрат на изображении ниже представляет собой 1 гигабайт VRAM.

Рисунок 14. Fine-tuning / Dreambooth — Memory footprint — FP32

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

Получается, все модели занимают около 5 гигабайт VRAM (часть из них является обучаемыми). Чтобы осуществлять оптимизацию, необходимо где-то хранить градиенты. Обычно для этого нужен такой же объем памяти, как и для обучаемых параметров (всего, кроме VAE). Следовательно, для Stable Diffusion это еще около 4.5 гигабайт VRAM.

Последнее по порядку, но не по значимости — оптимизатор и его параметры. Если мы используем Adam (сейчас он самый популярный) — мы должны помнить, что у него большое потребление памяти, равное почти x2 по количеству обучаемых параметров. Таким образом, это еще около 9 гигабайт VRAM.

И… к этому моменту у нас уже закончилась память на видеокарте. Для обучения нам нужно что-то предпринять, чтобы это исправить. Самый очевидный шаг — перевести модели с fp32 в fp16.

Рисунок 15. Fine-tuning / Dreambooth — Memory footprint — FP16

При переходе в fp16 потребление памяти моделями и градиентами уменьшается почти в 2 раза, однако для хранения состояний оптимизатора нам все еще приходится использовать fp32.

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

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

Parameter-Efficient fine-tuning

Рисунок 16. Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning [источник: arXiv:2303.15647]

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

Самый популярный метод, активно применяющийся сообществом Stable Diffusion, — LoRA. Давайте его и рассмотрим 🙂.

LoRA

LoRA расшифровывается как Low-Rank-Adaptation. Этот подход красив тем, что он прост.

Рисунок 17. LoRA architecture

\( h = W_0x + \Delta Wx = W_0x + BAx \)

В дополнение к базовым весам слоя \( W \) мы вводим новые матрицы \( B \) и \( A \). Они по размерности меньше базовой матрицы \( W \).

При этом их скалярное произведение \( \Delta W \) имеет тот же размер, что и размер матрицы \( W \) (например, \( W \in \mathbb{R}^{N \times M} \), \( B \in \mathbb{R}^{N \times r} \), \( A \in \mathbb{R}^{r \times M} \), \( r \) — гиперпараметр алгоритма; чем больше \( r \), тем больше LoRA).

Такой подход позволяет тренировать различные пары матриц \( B \) и \( A \) под разные задачи с использованием одной базовой модели. Эти пары можно применять как подключаемые адаптеры. А еще обучение только матриц \( B \) и \( A \)сильно сокращает потребление памяти. Плюс многие эмпирические исследования показывают: модели, дотренированные с помощью LoRA, способны потягаться по точности с моделями, полученными путем файн-тюна всех параметров.

Особенности LoRA

  1. Предобученная модель можеть использоваться, чтобы обучить разные LoRA под целевые задачи.
  2. LoRA делает обучение более эффективным с точки зрения потребления VRAM и предоставляет качество, сравнимое с полным дообучением модели.
  3. Матрицы LoRA объединяются с базовыми весами при деплое моделей, поэтому они не оказывают влияния на скорость их исполнения.
  4. LoRA может быть получена из двух чекпоинтов одной архитектуры. Для этого достаточно подсчитать \( \Delta W \) и разложить его на произведение матриц с использованием, например, алгоритма SVD (Singular Value Decomposition). Этот подход также применяется, чтобы уменьшить ранг натренированной LoRA с потерей в качестве, но с сильным уменьшением веса.
  5. Подход LoRA применим к разным слоям: Embedding, Linear, Convolutional, 1d, 2d, 3d и др. Поэтому он может использоваться для очень широкого спектра моделей (классификация, компьютерное зрение, генерация текста).

Потребление памяти LoRA

Рисунок 18. LoRA Memory Footprint

С точки зрения потребления памяти LoRA сильно уменьшает количество обучаемых параметров, что приводит и к уменьшению требуемой памяти для хранения градиентов и состояния оптимизатора. Это, в свою очередь, создает немаленький свободный объем памяти для хранения промежуточных активаций, больше памяти для увеличения размера батча; в нашем игрушечном примере удалось освободить целых 6 гигабайт VRAM из 16 на карте, что, на мой взгляд, является неплохим результатом.

Заметим также: выходной размер LoRA для Stable Diffuision значительно меньше полного чекпоинта, поскольку нам необходимо хранить только матрицы \( B \) и \( A \). Конечно, размер LoRA зависит от ранга, но в целом мы говорим о порядке от нескольких мегабайт до сотен мегабайт (в сравнении с несколькими гигабайтами для полного чекпоинта Stable Diffusion):

  • 11MiB — ранг 8;
  • 21MiB — ранг 16;
  • 41MiB — ранг 32;
  • 82MiB — ранг 64;
  • 163MiB — ранг 128.

Смешивание LoRA

Рисунок 19. Mixing LoRAs

Также в применении к Stable Diffusion мы получаем уникальную возможность — смешивать разные LoRA (или другие схожие адаптеры) и получать новые стилистики. Выше представлен пример, где я смешал LoRA, которая добавляет детали с LoRA в стилистике Ван Гога.

Реализации LoRA

Помимо оригинальной реализации LoRA для LLM от Microsoft есть большое количество других популярных реализаций (некоторые из них не ограничиваются Stable Diffusion):

  • kohya-ss — на мой взгляд, одна из самых продвинутых реализаций LoRA, созданная специально для Stable Diffusion и Stable Diffusion XL. В нее включено много вспомогательного кода вокруг подготовки тренировочных данных, хорошо отлаженные эвристики, тренировка на нескольких разрешениях. Эта реализация может использовать и другие адаптеры, такие как LyCORIS (LoHa/LoKr).
  • HuggingFace diffusers — официальная имплементация от HuggingFace. Применима ко всем диффузионным моделям, поддерживаемым их библиотеками. Хотя в ней могут отсутствовать последние наработки комьюнити, она достаточно проста в использовании.
  • HuggingFace peft — реализация от HuggingFace. Применима с разными типами моделей (не только с генеративными). У них есть хорошие туториалы по использованию адаптеров для классификации изображений, сегментации, классификации токенов, распознавания речи и др. А еще эта библиотека хорошо интегрирована в библиотеки diffusers и transformers и поддерживает большое количество других адаптеров.

LyCORIS: альтернативные реализации для low-rank-adaptation

Рисунок 20. Navigating Text-To-Image Customization: From LyCORIS Fine-Tuning to Model Evaluation [источник: arXiv:2309.14859]

Отметим также: сообщество разработало различные типы адаптеров, похожих на LoRA.

Есть популярная библиотека под названием LyCORIS, куда включены реализации таких адаптеров, как LoHa и LoKr.

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

LoHa (Low-rank Hadamard product Adaptation) базируется на идеях статьи FedPara. Одно из преимуществ LoHa — максимальный ранг результирующей матрицы больше, чем ранг результирующей матрицы в LoRA.

LoKr (Low-rank Kronecker product Adaptation) схож с идеями статьи Krona. Одно из преимуществ использования произведения Кронекера — его мультипликативная природа приводит к произведению рангов перемножаемых матриц. Поэтому теоретически при использовании данного адаптера ранг выходной матрицы может быть гораздо ниже, чем ранг LoRA.

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

Заключение

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

Тем не менее такие подходы, как Textual Inversion / Dreambooth / LoRA до сих пор используются в качестве ключевых компонентов во многих других, более продвинутых алгоритмах. Мне кажется, методы LoRA предлагают хороший компромисс между точностью и затратами на железо при обучении, и позволяют решать большой спектр задач. Поэтому если вы занимаетесь машинным обучением — очень рекомендуем обратить свое внимание на методы LoRA 😉

Телеграм-канал

DeepSchool

Короткие посты по теории ML/DL, полезные
библиотеки и фреймворки, вопросы с собеседований
и советы, которые помогут в работе

Открыть Телеграм

Увидели ошибку?

Напишите нам в Telegram!