Розуміння дифузійних моделей: Поглиблений аналіз генеративного ШІ

Дифузійні моделі стали потужним підходом у генеративному ШІ, демонструючи найсучасніші результати в генерації зображень, аудіо та відео. У цій детальній технічній статті ми дослідимо, як працюють дифузійні моделі, їх ключові інновації та чому вони стали такими успішними. Ми розглянемо математичні основи, процес навчання, алгоритми вибірки та передові застосування цієї захоплюючої нової технології.

Вступ до дифузійних моделей

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

Цей підхід був натхненний нерівноважною термодинамікою — зокрема, процесом зворотної дифузії для відновлення структури. У контексті машинного навчання ми можемо розглядати це як навчання зворотному процесу поступового додавання шуму до даних.

Деякі ключові переваги дифузійних моделей включають:

• Найсучасніша якість зображень, що перевершує GAN у багатьох випадках
• Стабільне навчання без адверсарної динаміки
• Високий рівень паралелізації
• Гнучка архітектура — може бути використана будь-яка модель, яка відображає входи на виходи тієї ж розмірності
• Міцне теоретичне підґрунтя

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

Source: Song et al.

Стохастичні диференціальні рівняння керують прямим і зворотним процесами в дифузійних моделях. Пряме СДР додає шум до даних, поступово перетворюючи їх на шумовий розподіл. Зворотне СДР, кероване вивченою функцією оцінки (score function), прогресивно видаляє шум, що призводить до генерації реалістичних зображень з випадкового шуму. Цей підхід є ключовим для досягнення високоякісної генеративної продуктивності в неперервних просторах станів.

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

Прямий дифузійний процес починається з точки даних x₀, вибірки з реального розподілу даних, і поступово додає гаусів шум протягом T кроків часу, створюючи все більш зашумлені версії x₁, x₂, …, xT.

На кожному кроці часу t ми додаємо невелику кількість шуму відповідно до:

x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε

Де:

• β_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₀ відповідно до прямого процесу

Іншими словами, ми навчаємо модель передбачати шум, який був доданий на кожному кроці часу.

Архітектура моделі

Source: Ronneberger et al.

Архітектура U-Net є центральною для кроку видалення шуму в дифузійній моделі. Вона має структуру кодувальника-декодувальника з перехресними зв’язками (skip connections), які допомагають зберегти деталі високої роздільної здатності під час процесу реконструкції. Кодувальник прогресивно зменшує роздільну здатність вхідного зображення, захоплюючи високорівневі ознаки, а декодувальник підвищує роздільну здатність закодованих ознак для реконструкції зображення. Ця архітектура особливо ефективна в задачах, що вимагають точної локалізації, таких як сегментація зображень.

Мережа передбачення шуму ε_θ може використовувати будь-яку архітектуру, яка відображає входи на виходи тієї ж розмірності. Архітектури у стилі 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)