T-LoRA: дообучить диффузионную модель на одной картинке и не переобучиться

Вы когда‑нибудь мечтали стать лучшей версией себя? Моложе, красивее, идеальнее… А вот LoRA уже стала!

Меня зовут Вера Соболева, я научный сотрудник лаборатории FusionBrain Института AIRI, а также стажер‑исследователь Центра глубинного обучения и байесовских методов НИУ ВШЭ. Cегодня я расскажу про наше свежее исследование T‑LoRA: Single Image Diffusion Model Customization Without Overfitting. Мы с коллегами придумали эффективный способ как файнтюнить диффузионные модели с помощью LoRA всего по одной картинке. 

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

Хорошая новость: эту задачу можно решить, копнув поглубже в свойства диффузии! В этой статье я расскажу, как это сделать.

Итак, вы решили дообучить диффузионную модель для генерации конкретного объекта, используя всего одно изображение. Вы принимаете решение обучить легковесную LoRA (Low‑Rank Adaptation), полагая, что этого будет достаточно, но сталкиваетесь с результатом, далёким от ваших ожиданий: монотонные позы, словно скопированные с оригинального снимка, размытые и скучные фоновые элементы, а сама генерация лишена яркости, гибкости и разнообразия.

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

Роль временных шагов диффузии в переобучении

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

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

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

Чтобы глубже понять влияние таймстепов на переобучение модели, мы провели следующий эксперимент. Во время дообучения SD‑XL с использованием LoRA мы изменили стандартный подход случайного семплирования таймстепов из полного диапазона (от 0 до 1000): вместо этого мы фиксировали узкие интервалы — по 100 шагов, семплируя таймстепы только из них.

Результаты показывают, что дообучение модели на наиболее шумных таймстепах (от 800 до 1000) приводит к быстрому переобучению, запоминанию поз и фона, что ограничивает разнообразие изображений и согласованность с промптом. В то же время эти временные шаги имеют решающее значение для определения формы и пропорций. На средних временных шагах (от 500 до 800) сгенерированный контекст становится богаче, и модель лучше воспроизводит мелкие детали концептов. Однако, мы теряем информацию, связанную с формой и пропорциями. Например, ботинок на рисунке сохраняет детали, соответствующие оригиналу, но становится короче. Наконец, дообучение на меньших временных шагах (от 0 до 500) демонстрирует наилучшее соответствие промпту и дает самую богатую генерацию. Однако этот подход не позволяет в точности воспроизвести объект, теряя как форму, так и мелкие детали объекта.

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

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

Vanilla T-LoRA

Исходя из этих выводов, мы предлагаем стратегию дообучения диффузионных моделей, зависящую от временного шага. Наш метод динамически корректирует ранги адаптеров LoRA в зависимости от временного шага диффузии, выделяя меньше параметров на более высоких временных шагах и больше — на более низких. Этот подход, которые мы назвали Vanilla T‑LoRA, реализован на основе маскирования столбцов LoRA.

Если стандартная LoRA описывается следующим выражением:

то Vanilla T-LoRA описывается другой формулой:

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

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

Ортогональность столбцов LoRA

Метод имеет важное ограничение: его эффективность зависит от ранга матриц  и . Если столбцы линейно зависимы, данные могут «просочиться» между маскированной и немаскированной частями, что нарушает надёжность исключения информации. Таким образом, для успешного применения стратегии маскирования требуется ортогональность матриц  и , а матрицы LoRA таким свойством не обладают.

Один из способов сделать матрицы ортогональными — применение регуляризации как в AdaLoRA:

где матрицы  и инициализируются случайно, а матрица  — нулями.

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

1. Применение этой регуляризации требует более 10 000 итераций обучения для того, чтобы матрицы сошлись к ортогональным, что значительно превосходит стандартные 1000–2000 итераций, необходимых для кастомизации диффузионной модели.

2. Инициализация матрицы очень маленькими значениями значительно замедляет обучение, как мы увидим далее.

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

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

Теперь мы можем выбрать их изначально ортогональными, используя SVD‑разложение некоторой матрицы:

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

T-LoRA

Объединяя вместе стратегию маскирования в зависимости от диффузионного шага и ортогональную инициализацию, позволяющую сохранять столбцы матриц линейно независимыми на протяжении всего обучения, мы приходим к финальному методу, который мы зовём T‑LoRA:

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

Vanilla T‑LoRA = LoRA + timestep dependent masking

T‑LoRA = Ortho‑LoRA + timestep dependent masking

Выбор инициализации

Осталось ответить на один важный вопрос: SVD какой матрицы использовать для инициализации?

Мы исследовали шесть вариантов инициализации Ortho‑LoRA, основанных на первых, средних и последних компонентах сингулярного разложения исходных весов и случайной матрицы . Метрики CLIP Image Similarity (IS) и CLIP Text Similarity (TS) для всех вариантов инициализации представлены на рисунке: 

Во‑первых, мы заметили, что для всех рангов и для обеих инициализаций (из и ) точки на графике по TS упорядочены в соответствии с величиной сингулярных значений при инициализации: использование первых компонентов приводит к самой низкой TS, за ним следуют средние компоненты, а инициализация последними компонентами даёт наивысшую TS. Это говорит о том, что более высокие сингулярные значения сильно коррелируют с переобучением. При рангах 4 и 8 инициализация с использованием последних компонент оказывается слишком близкой к нулю и замедляет процесс обучения. В то же время инициализация последними компонентами такого эффекта не вызывает. В итоге, использование последних компонентов сингулярного разложения случайной матрицы демонстрирует оптимальные результаты для большинства рангов, поэтому мы применяем этот подход во всех последующих экспериментах с Ortho‑LoRA и T‑LoRA.

Выбор минимального ранга

На рисунке ниже представлены примеры генерации для Ortho‑LoRA и T‑LoRA при значениях , установленных на уровне 25% и 50% от полного ранга. Оба варианта T‑LoRA значительно улучшают соответствие тексту и позволяют добиться большего разнообразия поз и фонов для концептов по сравнению с Ortho‑LoRA без маскирования. При меньших генерация становится более гибкой. Однако 25% иногда оказывается недостаточно для точной передачи концепта, в то время как 50% показывает оптимальные результаты и по передаче концепта, и по гибкости генерации.

Сравнение с LoRA

Мы сравнили предложенный метод со стандартной LoRA и Vanilla T‑LoRA на нескольких рангах. 

На всех рангах Vanilla T‑LoRA и T‑LoRA показывают более высокие значения TS по сравнению со стандартной LoRA, в то время как IS у них почти одинаковый. 

На низких рангах Vanilla T‑LoRA и T‑LoRA показывают схожую эффективность, так как LoRA приближается к полному рангу. Однако с увеличением ранга преимущества T‑LoRA становятся более очевидными: эффективность маскировки в Vanilla T‑LoRA снижается, в то время как T‑LoRA полностью раскрывает свой потенциал.

Сравнение с другими бейзлайнами

В дополнение к LoRA мы сравниваем T‑LoRA с другими легковесными методами настройки, такими как OFT, GSOFT и SVDiff. 

T‑LoRA демонстрирует наибольшее TS среди всех методов, за исключением SVDiff. Однако SVDiff показывает очень низкое IS и часто неудачно передает концепт. LoRA, напротив, обеспечивает наивысшее IS, но при этом показывает значительное переобучение. На рисунке ниже представлены примеры генерации для каждого метода.

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

User Study

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

Используя оригинальное изображение концепта, текстовый запрос и два сгенерированных изображения (одно создано с помощью T‑LoRA, а другое — альтернативным методом), мы попросили пользователей ответить на следующие вопросы: 

(1) «Какое изображение более точно представляет оригинальный концепт?» для оценки сходства с оригиналом; 

(2) «Какое изображение больше соответствует текстовому запросу?» для оценки текстового сходства; 

(3) «Какое изображение в целом лучше соответствует запросу и сохраняет идентичность концепта?» для оценки общего предпочтения. 

Для каждой пары методов мы случайным образом сгенерировали 60 уникальных пар концептов и запросов. Результаты User Study приведены в таблице

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

Заключение

Таким образом, наша работа делает важный шаг в сторону гибкого дообучения диффузионных моделей при ограниченных данных и вычислительных ресурсах. Сделано это ценой небольшого увеличения времени за счёт инициализации матриц через SVD, а также введения дополнительного гиперпараметра , за которым также нужно следить. В наших экспериментах значение 50% от полного ранга показало наилучшие результаты, однако оптимальный параметр может варьироваться в зависимости от конкретного концепта и задачи.

Если у вас остались вопросы, постараюсь на всё ответить в комментариях!

Также приглашаем вас подписаться на каналы авторов: if‑else, AI Brain, Complete AI.

Источник: habr.com