Диффузионные модели научили качественно дообучаться на одной картинке

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

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

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

На эту проблему обратили внимание исследователи из AIRI, ВШЭ и МГУ. Они предположили, что переобучение возникает главным образом при дообучении самых шумных шагов диффузионного процесса. В этот момент модель учится восстанавливать изображения из сильно искажённых входных данных, что непреднамеренно ограничивает её способность генерировать разнообразные и структурно гибкие сцены. В то же время, эти шумные этапы критически важны для сохранения формы и пропорций целевого объекта.

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

Для проверки нового метода, исследователи экспериментировали с объектами в широком спектре категорий, включая домашних животных, игрушки, предметы интерьера, аксессуары и многое другое. Используя T-LoRA, они дообучали модель Stable Diffusion-XL на одном изображении каждого объекта, сравнивая результат с LoRA, а также с различными методами ортогонального дообучения. Эксперименты показали, что с T-LoRA модель лучше следует промпту и демонстрирует большее разнообразие положений объектов.

Сравнение T-LoRA с его неортогональной версией Vanilla T-LoRA, базовой LoRA и другими легковесными методами дообучения