Від моделювання дифузії до відповідності потоку для генеративного ШІ — скорочена версія

pic

Що таке дифузія і яке її відношення до генеративного ШІ?

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

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

Маючи набір даних для навчання x, генеративні моделі мають на меті оцінити невідому розподільчу функцію (1)

pic

у вигляді моделюючого розподілу (2)

pic

параметризованого θ, а потім вибірки з цієї генеративної моделі (2) для генерування нових даних, ніби вони були вибрані з оригінального, невідомого розподілу (1).

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

Глибинні генеративні моделі, тобто генеративні моделі, що використовують багатошарові нейронні мережі, вирішують цю проблему, явно моделюючи невідому, оригінальну функцію щільності, але обмежуючи її певним чином, щоб зробити її прораховуваною, наприклад, варіаційні автоенкодери (VAE) і моделі дифузії, або використовуючи апроксимацію щільності з прораховуваною функцією, наприклад, авторегресивні та нормалізуючі потоки. Альтернативно, генеративні змагальні мережі (GAN) моделюють функцію щільності опосередковано, використовуючи стохастичний процес генерації даних (Foster, 2023).

Явне моделювання функції щільності у вигляді моделей дифузії революціонізувало і домінує в таких галузях, як комп'ютерна графіка та комп'ютерне зорове сприйняття, зокрема. Вони є прикладами генеративних моделей сімейства Flow Matching (Lipman et al., 2023). Flow Matching, або, загалом, Generator Matching (Holderrieth et al., 2024), фактично можна показати як уніфікувати всі генеративні моделі, згадані вище, по всіх модальностях, і це є новим досягненням у генерації зображень, відео та аудіо (Po et al., 2023). Моделі дифузії (Ho et al., 2020; Song et al., 2021) та Flow Matching (Esser et al., 2024) є основою продуктів для генерації зображень та відео, таких як OpenAI DALL.E 3, Google ImageGen, Black Forest Labs Flux, Runway Gen3-Alpha та OpenAI Sora відповідно.

Цей блог присвячений переважно моделям дифузії. Другий блог розширить обговорення до Flow, відповідно Generative, Matching. Їх математичне виведення не є простим. Обговорення орієнтоване на аудиторію, яка зацікавлена в розумінні основних сил, що стоять за цими революційними генеративними досягненнями, і їх реалізацією без необхідності проходити через їх покрокове математичне виведення (для чого я рекомендую ознайомитися з науковими працями, на які посилаюся нижче).
Він передбачає середній рівень знань з теорії ймовірностей та Python.

Що саме роблять моделі дифузії?

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

Моделі дифузії можна насправді обґрунтувати і вивести більш детально кількома різними способами. Вони були спочатку представлені як дискретно-часові гауссові процеси (Sohl-Dickstein et al., 2015). Згодом Song та Ermon (2019) обґрунтували моделі дифузії через безперервно-часові стохастичні диференціальні рівняння (SDE). Схожі, але нещодавно, Flow Matching розглядає дифузію як моделювання ймовірнісного шляху розподілу від джерела до цілі через процеси додавання шуму, де генератор цього процесу параметризується за допомогою функції оцінки (Lipman et al., 2023; Lipman et al., 2024). Давайте спершу розглянемо саме це виведення, яке називається score matching.

Ідея score matching полягає в тому, щоб параметризувати генеративну модель (2) таким чином, щоб градієнт, і тільки градієнт, її логарифмічної функції правдоподібності нагадував градієнт логарифмічної функції правдоподібності невідомого розподілу даних (1). Тобто, θ має бути обраний таким чином, щоб градієнти обох розподілів максимально нагадували один одного, а не самі функції правдоподібності. Це завдання обмежує інформацію, яку модель отримує про підлеглу, невідому структуру функції щільності, що генерує дані, лише до градієнта для певного значення θ. Однак, позитивним аспектом є те, що апроксимація функції щільності стає прораховуваною.

Градієнт логарифмічної функції правдоподібності генеративної моделі відносно вхідних даних x представлений функцією оцінки (Stein) (3):

pic

Інтуїтивно це представляє собою векторне поле, яке вказує в напрямку найбільшого зростання логарифмічної функції правдоподібності. Ідея score matching полягає в мінімізації дивергенції між цією функцією оцінки моделі та функцією оцінки розподілу даних (4),

pic

Тобто, після оптимізації цієї цільової функції модель вивчає ймовірнісний шлях невідомого розподілу, що генерує дані. Результати Hyvärinen (2005) та Song і Ermon (2019) дозволяють переписати (4) таким чином, щоб вона залежала тільки від генеративної моделі (2), що дозволяє оптимізувати її за допомогою стандартних методів оптимізації, таких як градієнтний спуск з нейронними мережами.

Для більшості реальних наборів даних, навіть дуже високорозмірних, можна спостерігати, що навчальні дані перебувають на двовимірній маніфольді, вбудованій в більш високорозмірний простір, так звану гіпотезу маніфольди (Whiteley et al., 2022). Оскільки функція оцінки є градієнтом, обчисленим в амбієнтному просторі, вона стає невизначеною, якщо вхідні дані перебувають на низькорозмірній маніфольді. Song та Ermon (2019) показали, що додавання невеликих пертурбацій різної величини до вхідних даних, наприклад, у вигляді гауссового шуму, вирішує цю проблему, і нейронна мережа для оцінки градієнта сходиться, тобто дифузія шуму робить оптимізацію (4) прораховуваною.

Дифузія і score matching насправді є еквівалентними, і основний принцип також працює в дискретно-часовій моделі Denoising Diffusion Probabilistic Model, запропонованій Ho et al.
(2020), який вважається проривною статтею в цій сфері, використовує зв'язок між дифузією та score matching для навчання моделі дифузії на момент її публікації.

Модель Denoising Diffusion Probabilistic Model

Моделі денойзінгової дифузії генерують зразки з невідомого розподілу даних (1), послідовно прибираючи шум з випадкових зразків шуму до зразків з обчислювально зручної моделі (1). Режим навчання — це двоетапний процес, з прямим процесом, званим енкодером і зворотнім (денойзінговим) процесом, званим декодером. Завдання енкодера — навчити репрезентацію, латентну змінну, вхідних навчальних даних. Ця латентна змінна покликана вловити суть навчальних даних, тому вона може бути втраточною, але зберігає основні характеристики вхідних даних. Вона не є частиною навчальних даних, але є частиною моделі і слугує як простір для семплінгу, з якого декодер може згенерувати нові, не спостережувані вихідні дані, наприклад, нове згенероване зображення. Тепер, як це працює детальніше?

Припустимо, що ми розглядаємо сценарій генерації зображень і позначаємо (5)

pic

як оригінальне зображення, а (6)

pic

як латентну змінну, при цьому (7)

pic

представляє проміжні латентні стани. Спільний розподіл послідовності

pic

позначається як

pic

для прямого процесу і

pic

для зворотного процесу.

Sohl-Dickstein et al. (2015) пропонують спочатку зробити кожний перехідний розподіл залежним тільки від його безпосереднього попередника, тобто створити «безпам'ятну» марківську ланцюг, тим самим розбиваючи процес на окремі етапи з простими(шими) розподілами, що робить обчислення зручними, тобто (8) та (9):

pic

По-друге, вони пропонують встановити (6) як гауссівський розподіл зі середнім значенням нуль та дисперсією одиниця, тобто,

pic

що дозволяє також розглядати кожний перехідний розподіл як гауссівський. В результаті, обчислення перехідних та (не нульових, одиничних дисперсій) спільних розподілів під час прямого проходу залежить лише від двох моментів гауссівського розподілу і стає зручним.

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

Зверніть увагу, що прямий процес не вимагає жодного параметричного моделювання: після того, як (5) задано у вигляді чітко визначеного навчального зображення, проміжні зашумлені зображення (7) можна просто отримати, додаючи гауссівський шум відповідно до графіка дисперсії шуму (так званого графіка дифузії)

pic

що представляє зашумлене зображення на етапі t. Зверніть увагу, що завдяки так званому трикстеру перепараметризації, зашумлене зображення (7) можна отримати з (5) безпосередньо за один обчислювальний крок, а не t обчислювальних кроків.

Денойзер тоді кондиціонується на денойзингу саме цього специфічного латентного зображення на етапі t і лише цього зашумленого зображення.
Тут приходить цікава частина: з відомими навчальними зразками з навчального набору зображень та прямим процесом дифузії, декодер може бути навчений за допомогою навчання з учителем (за допомогою тієї ж нейронної мережі на кожному етапі). Більш конкретно, денойзер (11)

pic

навчається для пари вихідних та зашумлених зображень, причому зашумлене зображення чітко визначається за допомогою прямого процесу, шляхом зворотного поширення градієнта втрат між зображенням, яке передбачив денойзер на етапі t, та справжнім зображенням на етапі t-1, що відомо з непараметричного прямого процесу. Після того, як денойзер (11) буде навчено, його можна послідовно застосовувати T разів до зображення білого шуму (6), щоб отримати нове згенероване зображення (12)

pic

Ці прямі та зворотні процеси дифузії ілюструються на прикладі зображення квітки (ліве зображення) нижче.

pic

Автор: Андраш Бере, пост в блозі Denoising Diffusion Implicit Models, червень 2022

Підсумовуючи, алгоритм навчання для моделі денойзінгової дифузії для кожного зображення (5) та одного денойзінгового нейронного мережевого елемента (11) виглядає так (Chan, 2024)¹:

  1. Вибрати випадкову мітку часу t,

pic

  1. Згенерувати зразок

pic

  1. Зробити крок градієнтного спуску за квадратною помилкою між денойзінговою латентною репрезентацією та справжнім зображенням для всіх M навчальних зразків,

pic

  1. Повторювати з пункту 1 до досягнення збіжності.

Цей навчений денойзер (11) можна потім використовувати для інференсу, тобто для генерації зображень з вектору білого шуму шляхом запуску його T разів від зображення білого шуму (6) до нового згенерованого зображення (12). Тобто, маючи зображення білого шуму,

pic

повторювати для мітки часу t=T, T-1, …, 1:

pic

це радіус гауссівського розподілу, який є великим, тобто нагадує білий шум, для t=T і наближається до нуля для t=1. В кінці цього послідовного процесу отримується нове згенероване зображення у вигляді (12).

Зверніть увагу, що термін гауссівського шуму ε спрощує похідні, але також робить процес повільним. Насправді для конвергенції зворотного процесу дифузії може знадобитися велика кількість ітерацій. Починаючи з Denoising Diffusion Implicit Models (Song et al., 2021), було здійснено численні покращення ефективності оригінального алгоритму, включаючи покращення ефективності прямого процесу. Рекомендую ознайомитись з Po et al. (2023) для обговорення цих покращень.

Реалізація моделі дифузії

Чудові реалізації Python процесів моделі дифузії були доступні від Андраша Бере та Девіда Фостера відповідно.² Я рекомендую використовувати Docker-образ Фостера для вивчення деталей реалізації та експериментів з моделями дифузії. Якщо ви плануєте запускати реалізацію моделі дифузії на вашому локальному комп'ютері, то він має бути обладнаний спеціалізованим GPU, оскільки часи тренувань можуть займати кілька годин в іншому випадку.

Навчальний набір даних — це набір даних "Oxford Flowers 102", наданий Visual Geometry Group Оксфордського університету. Він складається з зображень 102 категорій квітів, які зазвичай зустрічаються по всьому Сполученому Королівству, з 40 до 258 зображень на категорію.
Цей набір даних зручно доступний як TensorFlow dataset, вже розділений на навчальний, валідаційний та тестовий набори. Декілька прикладів навчальних зразків, масштабованих до 64x64 пікселів, показано нижче.

pic

Хоча можна отримати непогані результати за меншу кількість епох, для досягнення кращих результатів я рекомендую запускати імплементації моделі дифузії, розроблені Андрашем Берешем та Девідом Фостером, мінімум на 50 епох. Усі інші налаштування гіперпараметрів можна залишити незмінними.

pic

У моїх експериментах, поступово збільшуваний косинусний графік дифузії дав найпривабливіші візуальні результати.

pic

Імплементації підтримують TensorBoard для відслідковування прогресу навчання та збіжності, а також для побудови проміжних результатів навчання, приклади яких показано нижче для епох 1, 5, 15, 25 та 40 навчання.

pic

pic

pic

pic

pic

Навчена модель може бути використана для генерації нових зразків синтетичних зображень квітів:

pic

Це завершує обговорення концепції моделювання дифузії та її імплементації. Моделювання дифузії, як і будь-яка інша значуща концепція мульти-модальних генеративних AI, може бути розглянуте як особливий випадок Flow/Generative Matching, що вже згадувалося вище, і я планую детальніше розглянути його в другому і фінальному блозі.

Примітки

1 Ho et al. (2020) фактично формулюють функцію мети у термінах передбачення шуму. Формулювання еквівалентні, і для інтуїтивного розуміння я дотримуюсь формулювання функції мети у термінах передбачення зображень.

2 Андраш Берець надає імплементацію Denoising Diffusion Implicit Model (DDIM) (Song et al. 2021) замість Denoising Diffusion Probabilistic Model (DDPM) (Ho et al. 2020). Різниця полягає у детерміністичному (DDIM) проти стохастичного (DDPM) вибірці з навченої моделі дифузії "лише".

Література

S. H. Chan, Tutorial on Diffusion Models for Imaging and Vision, вересень 2024
https://arxiv.org/abs/2403.18103

P. Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, D. Podell, T. Dockhorn, Z. English, K. Lacey, A. Goodwin, Y. Marek, R. Rombach, Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, Proc. ICML, 2024
https://export.arxiv.org/pdf/2403.03206

D. Foster, Generative Deep Learning, 2-е видання, O’Reilly, Себастополь, США, 2023

J. Ho, A. Jain, і P. Abbeel, Denoising diffusion probabilistic models, Advances in Neural Information Processing Systems (NeurIPS), 2020
https://arxiv.org/abs/2006.11239

P. Holderrieth, M. Havasi, J. Yim, N. Shaul, I. Gat, T. Jaakkola, B. Karrer, R.T.Q. Chen, Y. Lipman, Generator Matching: Generative Modeling with Arbitrary Markov Processes, Computing Research Repository, 2024
https://arxiv.org/pdf/2410.20587

A. Hyvärinen, Estimation of non-normalized statistical models by score matching, Journal of Machine Learning Research (JMLR), Vol.
6, Випуск 24, стор. 695–709, 2005
https://jmlr.org/papers/volume6/hyvarinen05a/hyvarinen05a.pdf

Y. Lipman, R.T.Q. Chen, H. Ben-Hamu, M. Nickel та M. Lee, Flow Matching для генеративного моделювання, Proc. ICLR, 2023
https://arxiv.org/pdf/2210.02747

Y. Lipman, M. Havasi, P. Holderrieth, N. Shaul, M. Le, B. Karrer, R.T.Q. Chen, D. Lopez-Paz, H. Ben-Hamu, I. Gat, Flow Matching Guide and Code, грудень 2024
https://arxiv.org/pdf/2412.06264

R. Po, W. Yifan, V. Golyanik, K. Aberman, J.T. Barron, A. Bermano, E. Chan, T. Dekel, A. Holynski, A. Kanazawa, C.K. Liu, L. Liu, B. Mildenhall, M. Nießner, B. Ommer, C. Theobalt, P. Wonka, G. Wetzstein, Стан розвитку моделей дифузії для візуальних обчислень, жовтень 2023
https://arxiv.org/pdf/2310.07204

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan та S. Ganguli, Глибоке ненаглядне навчання за допомогою термодинаміки нерівноваги, Proc. of Int. Conf. on Machine Learning (ICML), Vol. 27, стор. 2256–2265, 2015
https://arxiv.org/pdf/1503.03585

Y. Song та S. Ermon, Генеративне моделювання шляхом оцінки градієнтів розподілу даних, у: Advances in Neural Information Processing Systems, стор. 11895–11907, 2019

J. Song, C. Meng та S. Ermon, Denoising diffusion implicit models, Int. Conf. on Learning Representations (ICLR), 2021
https://openreview.net/forum?id=St1giarCHLP

N. Whiteley, A. Gray та P. Rubin-Delanchy, Статистичне дослідження Маніфесту гіпотези, 2022
https://arxiv.org/pdf/2208.11665

D. Foster, Generative Deep Learning, 2-е видання, O’Reilly, Себастополь, США, 2023

Y. Song та S. Ermon, Генеративне моделювання шляхом оцінки градієнтів розподілу даних, у: Advances in Neural Information Processing Systems, стор. 11895–11907, 2019

Оригінально опубліковано на https://technoids.substack.com.

Перекладено з: From Diffusion Modeling to Flow Matching for Generative AI - Abridged

Leave a Reply

Your email address will not be published. Required fields are marked *