Искусственный интеллект
Понимание моделей диффузии: глубокое погружение в генеративный ИИ
Модели диффузии стали мощным подходом в генеративном ИИ, демонстрируя передовые результаты при создании изображений, аудио- и видеоконтента. В этой подробной технической статье мы рассмотрим принципы работы моделей диффузии, их ключевые инновации и причины их успеха. Мы рассмотрим математические основы, процесс обучения, алгоритмы сэмплирования и передовые приложения этой захватывающей новой технологии.
Введение в модели диффузии
Диффузионные модели являются классом генеративных моделей, которые учатся постепенно очищать данные от шума, обращая процесс диффузии. Основная идея заключается в том, чтобы начать с чистого шума и итеративно очищать его до высококачественной выборки из целевого распределения.
Этот подход был вдохновлен неравновесной термодинамикой – в частности, процессом обращения диффузии для восстановления структуры. В контексте машинного обучения мы можем думать об этом как об обучении обращению постепенного добавления шума к данным.
Некоторые ключевые преимущества моделей диффузии включают в себя:
- Современное качество изображения, во многих случаях превосходящее GAN
- Стабильное обучение без противоборствующей динамики
- Высокая степень параллелизации
- Гибкая архитектура – можно использовать любую модель, которая сопоставляет входы с выходами той же размерности.
- Прочная теоретическая подготовка
Давайте подробнее рассмотрим, как работают модели диффузии.
Источник: Сонг и др.
Стохастические дифференциальные уравнения управляют прямыми и обратными процессами в моделях диффузии. Прямое SDE добавляет шум к данным, постепенно преобразуя их в распределение шума. Обратное SDE, управляемое функцией выученных оценок, постепенно удаляет шум, что приводит к генерации реалистичных изображений из случайного шума. Этот подход является ключом к достижению высококачественной генеративной производительности в непрерывных пространствах состояний
Процесс прямой диффузии
Процесс прямой диффузии начинается с точки данных x₀, выбранной из реального распределения данных, и постепенно добавляет гауссовский шум в течение T временных шагов для получения все более шумных версий x₁, x₂, …, xT.
На каждом временном шаге t мы добавляем небольшое количество шума в соответствии с:
x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε
Где:
- β_т это график дисперсии, который контролирует, сколько шума добавляется на каждом шаге
- ε это случайный гауссовский шум
Этот процесс продолжается до тех пор, пока xT не станет практически чистым гауссовым шумом.
Математически это можно описать как цепь Маркова:
q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)
Где N обозначает гауссово распределение.
График β_t обычно выбирается небольшим для ранних временных шагов и увеличивается со временем. Обычные варианты включают линейный, косинусный или сигмоидальный графики.
Процесс обратной диффузии
Цель диффузионной модели — изучить обратный процесс — начать с чистого шума xT и постепенно удалять из него шум, чтобы получить чистый образец x₀.
Мы моделируем этот обратный процесс следующим образом:
p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))
Где μ_θ и σ_θ^2 — обученные функции (обычно нейронные сети), параметризованные θ.
Ключевое нововведение заключается в том, что нам не нужно явно моделировать полное обратное распределение. Вместо этого мы можем параметризовать его в терминах прямого процесса, который нам известен.
В частности, мы можем показать, что оптимальное среднее значение обратного процесса μ* равно:
μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))
Где:
- α_t = 1 – β_t
- ε_θ — это обученная сеть прогнозирования шума
Это дает нам простую цель — обучить нейронную сеть ε_θ предсказывать шум, добавляемый на каждом шаге.
Цель обучения
Цель обучения для диффузионных моделей может быть получена из вариационного вывода. После некоторого упрощения мы приходим к простой потере L2:
L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]
Где:
- t выбирается равномерно от 1 до T
- x₀ выбирается из обучающих данных
- ε — это выборочный гауссовский шум
- x_t создается путем добавления шума к x₀ в соответствии с прямым процессом
Другими словами, мы обучаем модель прогнозировать шум, добавляемый на каждом временном шаге.
Модельная архитектура
Источник: Роннебергер и др.
Архитектура U-Net является центральной для этапа шумоподавления в модели диффузии. Она имеет структуру кодер-декодер с пропусками, которые помогают сохранить мелкие детали в процессе реконструкции. Кодер постепенно понижает разрешение входного изображения, захватывая высокоуровневые признаки, а декодер повышает разрешение закодированных признаков для реконструкции изображения. Эта архитектура особенно эффективна в задачах, требующих точной локализации, таких как сегментация изображения.
Сеть прогнозирования шума ε_θ может использовать любую архитектуру, которая отображает входы в выходы той же размерности. Архитектуры в стиле U-Net являются популярным выбором, особенно для задач генерации изображений.
Типичная архитектура может выглядеть так:
class DiffusionUNet(nn.Module):
def __init__(self):
super().__init__()
# Downsampling
self.down1 = UNetBlock(3, 64)
self.down2 = UNetBlock(64, 128)
self.down3 = UNetBlock(128, 256)
# Bottleneck
self.bottleneck = UNetBlock(256, 512)
# Upsampling
self.up3 = UNetBlock(512, 256)
self.up2 = UNetBlock(256, 128)
self.up1 = UNetBlock(128, 64)
# Output
self.out = nn.Conv2d(64, 3, 1)
def forward(self, x, t):
# Embed timestep
t_emb = self.time_embedding(t)
# Downsample
d1 = self.down1(x, t_emb)
d2 = self.down2(d1, t_emb)
d3 = self.down3(d2, t_emb)
# Bottleneck
bottleneck = self.bottleneck(d3, t_emb)
# Upsample
u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb)
u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb)
u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb)
# Output
return self.out(u1)
Ключевые компоненты:
- Архитектура в стиле U-Net с пропусками соединений
- Внедрение времени в условие на временном шаге
- Гибкая глубина и ширина
Алгоритм выборки
После обучения нашей сети прогнозирования шума ε_θ мы можем использовать её для генерации новых выборок. Базовый алгоритм выборки:
- Начнем с чистого гауссовского шума xT
- Для t = T до 1:
- Прогноз шума:
ε_θ(x_t, t) - Вычислить среднее:
μ = 1/√(1-β_t) * (x_t - β_t/√(1-α_t) * ε_θ(x_t, t)) - Пример:
x_{t-1} ~ N(μ, σ_t^2 * I)
- Прогноз шума:
- Возврат x₀
Этот процесс постепенно очищает образец от шума, руководствуясь нашей обученной сетью прогнозирования шума.
На практике существуют различные методы отбора проб которые могут улучшить качество или скорость:
- Выборка DDIM: Детерминированный вариант, который допускает меньше шагов выборки.
- Выборка предков: Включает изученную дисперсию σ_θ^2
- Усеченная выборка: Останавливается раньше для более быстрой генерации
Вот базовая реализация алгоритма выборки:
def sample(model, n_samples, device):
# Start with pure noise
x = torch.randn(n_samples, 3, 32, 32).to(device)
for t in reversed(range(1000)):
# Add noise to create x_t
t_batch = torch.full((n_samples,), t, device=device)
noise = torch.randn_like(x)
x_t = add_noise(x, noise, t)
# Predict and remove noise
pred_noise = model(x_t, t_batch)
x = remove_noise(x_t, pred_noise, t)
# Add noise for next step (except at t=0)
if t > 0:
noise = torch.randn_like(x)
x = add_noise(x, noise, t-1)
return x
Математика, лежащая в основе моделей диффузии
Чтобы по-настоящему понять модели диффузии, крайне важно глубже разобраться в математической основе, лежащей в их основе. Давайте рассмотрим некоторые ключевые концепции подробнее:
Цепь Маркова и стохастические дифференциальные уравнения
Процесс прямой диффузии в диффузионных моделях можно рассматривать как цепь Маркова или, в непрерывном пределе, как стохастическое дифференциальное уравнение (SDE). Формулировка SDE обеспечивает мощную теоретическую основу для анализа и расширения диффузионных моделей.
Прямое SDE можно записать как:
dx = f(x,t)dt + g(t)dw
Где:
- f(x,t) — дрейфовый член
- g(t) — коэффициент диффузии
- dw — винеровский процесс (броуновское движение)
Различные варианты f и g приводят к различным типам диффузионных процессов. Например:
- Взрыв дисперсии (VE)
SDE: dx = √(d/dt σ²(t)) dw - Сохранение дисперсии (VP)
SDE: dx = -0.5 β(t)xdt + √(β(t)) dw
Понимание этих SDE позволяет нам выводить оптимальные стратегии выборки и распространять модели диффузии на новые области.
Сопоставление оценок и шумоподавление
Связь между моделями диффузии и сопоставлением оценок дает еще одну ценную перспективу. Функция оценки определяется как градиент логарифмической плотности вероятности:
s(x) = ∇x log p(x)
Denoising score matching нацелен на оценку этой функции оценки путем обучения модели шумоподавлению слегка возмущенных точек данных. Эта цель оказывается эквивалентной цели обучения диффузионной модели в непрерывном пределе.
Эта связь позволяет нам использовать методы генеративного моделирования на основе оценок, такие как отожженная динамика Ланжевена для выборки.
Расширенные методы обучения
Выборка по важности
Стандартная модель диффузии обучает выборки временных шагов равномерно. Однако не все временные шаги одинаково важны для обучения. Методы выборки по важности могут использоваться для фокусировки обучения на наиболее информативных временных шагах.
Один из подходов заключается в использовании неравномерного распределения по временным шагам, взвешенного по ожидаемой норме L2 оценки:
p(t) ∝ E[||s(x_t, t)||²]
Это может привести к более быстрому обучению и улучшению качества выборки.
Прогрессивная дистилляция
Прогрессивная дистилляция — это метод создания более быстрых моделей отбора проб без ущерба качеству. Процесс работает следующим образом:
- Обучите базовую модель диффузии с большим количеством временных шагов (например, 1000)
- Создайте модель студента с меньшим количеством временных шагов (например, 100)
- Обучите студента применять процесс шумоподавления базовой модели.
- Повторяйте шаги 2-3, постепенно уменьшая время выполнения шагов.
Это позволяет добиться высококачественной генерации со значительно меньшим количеством шагов шумоподавления.
Архитектурные инновации
Модели диффузии на основе трансформатора
В то время как архитектуры U-Net были популярны для моделей диффузии изображений, недавние работы исследовали использование архитектур трансформаторов. Трансформаторы предлагают несколько потенциальных преимуществ:
- Лучшая обработка долгосрочных зависимостей
- Более гибкие механизмы кондиционирования
- Более простое масштабирование до более крупных размеров моделей
Модели как DiT (диффузионные трансформаторы) показали многообещающие результаты, потенциально открывающие путь к еще более качественной генерации.
Иерархические модели диффузии
Иерархические модели диффузии генерируют данные в различных масштабах, что позволяет обеспечить как глобальную согласованность, так и мелкие детали. Процесс обычно включает:
- Генерация выходных данных с низким разрешением
- Постепенное повышение частоты дискретизации и уточнение
Такой подход может быть особенно эффективен для создания изображений высокого разрешения или создания контента большого формата.
Продвинутые темы
Руководство без классификаторов
Руководство без классификаторов это метод улучшения качества и управляемости образцов. Основная идея заключается в обучении двух моделей диффузии:
- Безусловная модель p(x_t)
- Условная модель p(x_t | y), где y — некоторая условная информация (например, текстовая подсказка)
Во время выборки мы интерполируем между этими моделями:
ε_θ = (1 + w) * ε_θ(x_t | y) - w * ε_θ(x_t)
Где w > 0 — это шкала ориентира, которая контролирует, насколько следует подчеркнуть условную модель.
Это позволяет добиться более сильной обусловленности без необходимости переобучения модели. Это имело решающее значение для успеха моделей преобразования текста в изображение, таких как DALL-E 2 и Stable Diffusion.
Скрытая диффузия
Источник: Ромбах и др.
Модель скрытой диффузии (LDM) процесс включает кодирование входных данных в скрытое пространство, где происходит процесс диффузии. Модель постепенно добавляет шум к скрытому представлению изображения, что приводит к созданию шумной версии, которая затем шумоподавляется с использованием архитектуры U-Net. U-Net, управляемая механизмами перекрестного внимания, интегрирует информацию из различных источников обусловливания, таких как семантические карты, текст и представления изображений, в конечном итоге реконструируя изображение в пиксельном пространстве. Этот процесс имеет решающее значение для создания высококачественных изображений с контролируемой структурой и желаемыми атрибутами.
Это дает ряд преимуществ:
- Более быстрое обучение и выборка
- Лучшая обработка изображений высокого разрешения
- Легче включить кондиционирование
Процесс работает следующим образом:
- Обучить автоэнкодер сжимать изображения в скрытое пространство
- Обучить модель диффузии в этом скрытом пространстве
- Для генерации выполните выборку в скрытом пространстве и декодируйте в пиксели
Этот подход оказался весьма успешным и лег в основу таких моделей, как Stable Diffusion.
Модели согласованности
Модели согласованности — это недавнее новшество, направленное на повышение скорости и качества моделей диффузии. Основная идея заключается в обучении одной модели, которая может отображать любой уровень шума напрямую в конечный результат, а не требовать итеративного шумоподавления.
Это достигается с помощью тщательно разработанной функции потерь, которая обеспечивает согласованность между прогнозами при разных уровнях шума. Результатом является модель, которая может генерировать высококачественные образцы за один прямой проход, что значительно ускоряет вывод.
Практические советы по обучению диффузионных моделей
Обучение высококачественных моделей диффузии может быть сложным. Вот несколько практических советов по улучшению стабильности обучения и результатов:
- Отсечение градиента: Используйте ограничение градиента, чтобы предотвратить взрывное увеличение градиентов, особенно на ранних этапах обучения.
- EMA весов моделей: Сохраняйте экспоненциальное скользящее среднее (EMA) весов модели для выборки, что может привести к более стабильной и качественной генерации.
- Увеличение данных: Для моделей изображений простые дополнения, такие как случайные горизонтальные перевороты, могут улучшить обобщение.
- Планирование шума: Поэкспериментируйте с различными графиками шума (линейным, косинусным, сигмоидальным), чтобы найти тот, который лучше всего подходит для ваших данных.
- Смешанная точная тренировка: Используйте смешанную точность обучения, чтобы сократить использование памяти и ускорить обучение, особенно для больших моделей.
- Условное поколение: Даже если вашей конечной целью является безусловная генерация, обучение с кондиционированием (например, на классах изображений) может улучшить общее качество выборки.
Оценка моделей диффузии
Правильная оценка генеративных моделей имеет решающее значение, но является сложной задачей. Вот некоторые общие метрики и подходы:
Начальное расстояние Фреше (FID)
FID широко используемая метрика для оценки качества и разнообразия сгенерированных изображений. Она сравнивает статистику сгенерированных образцов с реальными данными в пространстве признаков предварительно обученного классификатора (обычно InceptionV3).
Более низкие значения FID указывают на более высокое качество и более реалистичные распределения. Однако FID имеет ограничения и не должен быть единственной используемой метрикой.
Начальный балл (IS)
Начальная оценка измеряет как качество, так и разнообразие сгенерированных изображений. Он использует предварительно обученную сеть Inception для вычисления:
IS = exp(E[KL(p(y|x) || p(y))])
Где p(y|x) — условное распределение классов для сгенерированного изображения x.
Более высокий показатель IS указывает на лучшее качество и разнообразие, но он имеет известные ограничения, особенно для наборов данных, сильно отличающихся от ImageNet.
Отрицательное логарифмическое правдоподобие (НЛЛ)
Для моделей диффузии мы можем вычислить отрицательную логарифмическую вероятность удерживаемых данных. Это дает прямую меру того, насколько хорошо модель соответствует истинному распределению данных.
Однако точная оценка NLL для многомерных данных может оказаться затратной с точки зрения вычислений.
Человеческая оценка
Для многих приложений, особенно творческих, человеческая оценка остается решающей. Это может включать:
- Параллельное сравнение с другими моделями
- Оценки в стиле теста Тьюринга
- Оценки, специфичные для конкретных задач (например, субтитры изображений для моделей преобразования текста в изображение)
Несмотря на свою субъективность, человеческая оценка может охватывать аспекты качества, которые упускают из виду автоматизированные показатели.
Модели диффузии в производстве
Развертывание моделей диффузии в производственных средах представляет собой уникальные проблемы. Вот некоторые соображения и передовые практики:
Оптимизация для вывода
- ONNX экспорт: Преобразование моделей в формат ONNX для более быстрого вывода на различном оборудовании.
- квантование: Используйте такие методы, как квантование INT8, чтобы уменьшить размер модели и повысить скорость вывода.
- Кэширование: Для условных моделей кэшируйте промежуточные результаты для безусловной модели, чтобы ускорить руководство без классификатора.
- Пакетная обработка: Используйте пакетную обработку для эффективного использования ресурсов графического процессора.
Масштабирование
- Распределенный вывод: Для приложений с высокой пропускной способностью реализуйте распределенный вывод на нескольких графических процессорах или машинах.
- Адаптивная выборка: Динамически регулируйте количество шагов выборки на основе желаемого соотношения качества и скорости.
- Прогрессивное поколение: Для больших выходных данных (например, изображений с высоким разрешением) сгенерируйте их постепенно от низкого к высокому разрешению, чтобы быстрее получить начальные результаты.
Безопасность и фильтрация
- Фильтрация контента: Внедрите надежные системы фильтрации контента, чтобы предотвратить создание вредоносного или неприемлемого контента.
- водяные знаки: Рассмотрите возможность включения невидимых водяных знаков в сгенерированный контент для обеспечения возможности отслеживания.
Области применения
Модели диффузии успешно применяются в широком спектре генеративных задач:
Генерация изображения
Генерация изображений — это то, где модели диффузии впервые получили известность. Вот некоторые известные примеры:
- ДАЛЛ-Э 3: Модель преобразования текста в изображение OpenAI, объединяющая кодировщик текста CLIP с декодером диффузионного изображения
- Стабильная диффузия: Модель скрытой диффузии с открытым исходным кодом для преобразования текста в изображение
- Изображение: Модель распространения текста в изображение от Google
Эти модели способны генерировать высокореалистичные и креативные изображения на основе текстовых описаний, превосходя предыдущие подходы на основе GAN.
Генерация видео
Модели диффузии также применялись при создании видео:
- Модели распространения видео: Создание видео путем рассмотрения времени как дополнительного измерения в процессе распространения
- Сделать видео: Модель распространения текста в видео Meta
- Изображение Видео: Модель распространения текста в видео от Google
Эти модели могут генерировать короткие видеоклипы из текстовых описаний, открывая новые возможности для создания контента.
3D-генерация
Недавние работы расширили возможности диффузионных моделей до трехмерной генерации:
- ДримФьюжн: Генерация текста в 3D с использованием 2D-моделей диффузии
- Точка: Модель диффузии облака точек OpenAI для создания 3D-объектов
Эти подходы позволяют создавать 3D-ресурсы из текстовых описаний, которые можно применять в играх, виртуальной и дополненной реальности, а также в дизайне продуктов.
Проблемы и будущие направления
Хотя модели диффузии продемонстрировали значительный успех, все еще существует ряд проблем и направлений для будущих исследований:
Вычислительная эффективность
Итеративный процесс выборки диффузионных моделей может быть медленным, особенно для выходов с высоким разрешением. Такие подходы, как модели скрытой диффузии и согласованности, направлены на решение этой проблемы, но дальнейшее повышение эффективности является активной областью исследований.
контролируемость
Хотя такие методы, как управление без классификаторов, улучшили управляемость, ещё предстоит проделать большую работу по обеспечению более точного контроля над генерируемыми результатами. Это особенно важно для творческих приложений.
Мультимодальная генерация
Текущие модели диффузии превосходны в генерации одной модальности (например, изображений или аудио). Разработка действительно многомодальных моделей диффузии, которые могут беспрепятственно генерировать между модальностями, является захватывающим направлением для будущей работы.
Теоретическое понимание
Хотя модели диффузии демонстрируют убедительные эмпирические результаты, ещё многое предстоит понять, почему они так эффективны. Более глубокое теоретическое понимание может привести к дальнейшим улучшениям и новым приложениям.
Заключение
Модели диффузии представляют собой шаг вперед в генеративном ИИ, предлагая высококачественные результаты в различных модальностях. Обучаясь обращению вспять процесса добавления шума, они обеспечивают гибкий и теоретически обоснованный подход к генерации.
От творческих инструментов до научного моделирования — возможность генерировать сложные, многомерные данные может трансформировать многие области. Однако важно подходить к этим мощным технологиям обдуманно, учитывая как их огромный потенциал, так и связанные с ними этические проблемы.
