Mengimplementasikan Model Difusi dari Nol: Perspektif Teoretis Baru

• Model difusi digunakan bukan hanya untuk pembuatan gambar, tetapi juga untuk masalah yang membutuhkan sampling dari distribusi multimodal seperti audio, video, 3D, desain protein, dan perencanaan jalur robot; tutorial ini menghubungkan pelatihan dan sampling dari sudut pandang optimisasi
• Proses pelatihan membuat (x_\sigma=x_0+\sigma\epsilon), yaitu data yang dicampur noise, lalu meminimalkan galat kuadrat rata-rata agar jaringan saraf (\epsilon_\theta(x,\sigma)) memprediksi arah noise
• Denoiser yang telah dilatih ditafsirkan sebagai proyeksi hampiran ke himpunan data (\mathcal{K}), dan denoiser ideal terkait dengan gradien fungsi jarak kuadrat yang di-smoothing oleh (\sigma)
• Sampling DDIM dapat dilihat sebagai gradient descent hampiran untuk (f(x)=\frac{1}{2}\mathrm{dist}_{\mathcal{K}}(x)^2), dan jadwal (\sigma_t) menentukan jumlah iterasi serta biaya evaluasi denoiser
• Dengan menggabungkan update estimasi gradien dan penambahan noise, DDIM, DDPM, serta sampler yang ditingkatkan oleh penulis dapat dibahas bersama melalui parameter gam dan mu, lalu berlanjut ke contoh toy model dan latent diffusion

---

Model Difusi dari Sudut Pandang Optimisasi

• Model difusi unggul dalam menghasilkan sampel dari distribusi multimodal, dan diterapkan bukan hanya pada alat text-to-image seperti Stable Diffusion, tetapi juga pada pembuatan audio, video, 3D, desain protein, dan perencanaan jalur robot
• Landasan teoretis tutorial ini adalah interpretasi optimisasi dari makalah ICML 2024 dan makalah terkait
• Implementasinya terutama merujuk pada smalldiffusion, dan kode dalam artikel disederhanakan untuk tujuan edukasi dibandingkan library aslinya

Pelatihan: Memprediksi Arah Noise

• Model difusi bertujuan mempelajari himpunan data (\mathcal{K}) dari contoh pelatihan, lalu menghasilkan sampel dari himpunan tersebut
  • Untuk gambar, (\mathcal{K} \subset \mathbb{R}^{c\times h \times w}) adalah himpunan nilai piksel yang sesuai dengan gambar realistis
  • Kerangka yang sama juga berlaku untuk domain diskret seperti audio, video, lintasan robot, dan teks
• Prosedur pelatihan dapat dilihat dalam tiga langkah
  • Melakukan sampling (x_0 \sim \mathcal{K}), (\sigma), dan (\epsilon \sim N(0,I))
  • Membuat data bercampur noise dengan (x_\sigma=x_0+\sigma\epsilon)
  • Meminimalkan loss kuadrat agar (\epsilon_\theta(x_\sigma,\sigma)) memprediksi (\epsilon)
• Dalam kode, training_loop membuat sigma dan eps melalui generate_train_sample untuk setiap batch x0, lalu mengoptimalkan MSE antara output model(x0 + sigma * eps, sigma) dan eps
• Alih-alih mengambil sampel (\sigma) secara seragam dari interval kontinu, nilainya diambil dari jadwal (\sigma) yang mendiskretkan (\sigma) menjadi (N) nilai
  • Kelas Schedule membungkus daftar sigmas yang mungkin dan mengambil sampel nilai per batch selama pelatihan
  • Contoh dalam artikel menggunakan ScheduleLogLinear(N, sigma_min=0.02, sigma_max=10)
  • ScheduleDDPM adalah jadwal untuk model difusi ruang piksel, sedangkan ScheduleLDM adalah jadwal untuk model latent diffusion seperti Stable Diffusion

Contoh Toy Swissroll

• Dataset toy adalah himpunan titik berbentuk spiral yang digunakan dalam salah satu makalah difusi awal, Sohl-Dickstein et al. 2015, dengan (\mathcal{K}\subset\mathbb{R}^2)
• Pada dataset sederhana, denoiser diimplementasikan sebagai MLP
  • Input adalah penggabungan (x\in\mathbb{R}^2) dan embedding 2D dari (\sigma)
  • Output adalah prediksi noise (\epsilon\in\mathbb{R}^2)
  • Banyak model difusi memakai sinusoidal positional embedding untuk (\sigma), tetapi embedding 2D sederhana juga bekerja dengan baik dalam contoh ini
• Konfigurasi pelatihan contoh menggunakan ScheduleLogLinear(N=200, sigma_min=0.005, sigma_max=10) dan epochs=15000
• Denoiser yang telah dilatih dapat divisualisasikan sebagai medan vektor dengan menggambar (x-\sigma\epsilon_\theta(x,\sigma))
  • Saat (\sigma) besar, denoiser cenderung memprediksi rata-rata data
  • Saat (\sigma) rendah dan input (x) dekat dengan data, denoiser memprediksi titik data sebenarnya

Menafsirkan Denoising sebagai Proyeksi

• Fungsi jarak terhadap himpunan data (\mathcal{K}) didefinisikan sebagai (\mathrm{dist}_{\mathcal{K}}(x)=\min{|x-x_0|:x_0\in\mathcal{K}})
• Proyeksi (x), yaitu (\mathrm{proj}_{\mathcal{K}}(x)), adalah himpunan titik di dalam (\mathcal{K}) yang mencapai jarak tersebut
• Jika (\mathcal{K}) adalah himpunan tertutup, (x\notin\mathcal{K}), dan proyeksinya unik, maka gradien fungsi jarak kuadrat menjadi (x-\mathrm{proj}_{\mathcal{K}}(x))
• Karena fungsi jarak (\mathrm{dist}_{\mathcal{K}}) tidak terdiferensiasi di semua tempat, diperkenalkan fungsi jarak kuadrat yang di-smoothing oleh (\sigma) dengan menggunakan softmin sebagai pengganti min
• Gradien fungsi jarak yang di-smoothing mengarah ke rata-rata berbobot titik-titik dalam (\mathcal{K}), sesuai bobot yang ditentukan oleh (x)

Denoiser Ideal dan Model Galat Relatif

• Denoiser ideal (\epsilon^*) adalah denoiser yang secara tepat meminimalkan loss pelatihan pada (\sigma) tertentu
• Jika data merupakan distribusi seragam diskret pada himpunan terbatas (\mathcal{K}), denoiser ideal dapat dinyatakan dalam bentuk tertutup
  • Bobot setiap titik data ditentukan oleh jarak antara (x_\sigma) dan titik tersebut
  • Pada dataset kecil, ini dapat dihitung langsung dengan IdealDenoiser
• Pada data toy, denoiser ideal mengarah ke rata-rata data saat (\sigma) besar, dan ke titik data terdekat saat (\sigma) kecil
• Teorema inti menyatakan hubungan (\frac{1}{2}\nabla_x \mathrm{dist}^2_{\mathcal{K}}(x,\sigma)=\sigma\epsilon^*(x,\sigma)) untuk semua (\sigma>0), (x\in\mathbb{R}^n)
• Model galat relatif menggunakan kondisi bahwa (x-\sigma\epsilon_\theta(x,\sigma)) mengaproksimasi (\mathrm{proj}_{\mathcal{K}}(x)) dengan baik
  • Berlaku ketika (\sqrt{n}\sigma) mengestimasi (\mathrm{dist}_{\mathcal{K}}(x)) dengan baik dalam faktor konstanta
  • Galat diasumsikan dibatasi hingga tidak lebih dari (\eta\mathrm{dist}_{\mathcal{K}}(x))
  • Pada noise rendah, di bawah manifold hypothesis, sebagian besar noise tambahan ortogonal terhadap manifold data sehingga denoising mengaproksimasi proyeksi
  • Pada noise tinggi, jika (\sigma) lebih besar daripada diameter (\mathcal{K}), denoiser yang memprediksi rata-rata berbobot data pun memiliki galat relatif kecil
• CIFAR-10 berukuran cukup untuk memungkinkan perhitungan denoiser ideal, dan dalam eksperimen, galat relatif antara proyeksi tepat di atas lintasan sampling dan output denoiser ideal terlihat kecil

Sampling: Denoising Berulang dan DDIM

• Jika denoiser terlatih tersedia, dari (x_t) yang bercampur noise dan level noise (\sigma_t), (x_0) diprediksi sebagai (\hat{x}0^t=x_t-\sigma_t\epsilon\theta(x_t,\sigma_t))
• Titik awal menetapkan (\sigma_T) besar dibandingkan diameter (\mathcal{K}), lalu mengambil sampel (x_T) secara independen dari (N(0,\sigma_T)) agar berada jauh dari (\mathcal{K})
• Pada noise tinggi, satu pemanggilan denoiser dapat memiliki galat absolut besar meskipun galat relatifnya kecil, dan prediksi denoiser ideal dekat dengan rata-rata data
• Karena itu, sampling memanggil denoiser berulang kali mengikuti jadwal (\sigma_t) untuk membuat sekuens (x_T,\ldots,x_0)
• Update (x_{t-1}=x_t-(\sigma_t-\sigma_{t-1})\epsilon_\theta(x_t,\sigma_t)) sama dengan algoritme sampling DDIM deterministik setelah transformasi koordinat
  • Bukti kesetaraannya dengan DDIM ada di Appendix A makalah

DDIM sebagai Minimalisasi Jarak

• DDIM ditafsirkan sebagai gradient descent hampiran untuk (f(x)=\frac{1}{2}\mathrm{dist}_{\mathcal{K}}(x)^2)
  • Ukuran langkahnya adalah (1-\sigma_{t-1}/\sigma_t)
  • (\nabla f(x_t)) diestimasi dengan (\epsilon_\theta(x_t,\sigma_t))
• Jadwal (\sigma_t) menentukan jumlah dan ukuran langkah gradien selama sampling
  • Jika langkah terlalu sedikit, (\mathrm{dist}_{\mathcal{K}}(x_t)) tidak berkurang sehingga mungkin tidak konvergen
  • Jika menggunakan banyak langkah kecil, jumlah evaluasi denoiser meningkat dan biaya komputasi membesar
• Admissible schedule adalah jadwal yang pada setiap iterasi membuat (\sqrt{n}\sigma_t) sesuai dengan (\mathrm{dist}_{\mathcal{K}}(x_t)) dalam faktor konstanta
  • Sekuens (\sigma_t) log-linear yang menurun secara geometris adalah admissible schedule
• Menurut teorema, jika (\nabla\mathrm{dist}{\mathcal{K}}(x)) ada pada (x_t) yang dihasilkan DDIM dan (\mathrm{dist}{\mathcal{K}}(x_T)=\sqrt{n}\sigma_T), maka (x_t) dihasilkan oleh gradient descent fungsi jarak kuadrat dan (\mathrm{dist}_{\mathcal{K}}(x_t)/\sqrt{n}\approx\sigma_t) tetap terjaga
• Pada contoh toy, sampler DDIM 20 langkah diimplementasikan dengan subsampling dari jadwal log-linear asli; sebagian besar sampel dekat dengan data asli, tetapi masih ada ruang untuk perbaikan

Sampler yang Ditingkatkan Berbasis Estimasi Gradien

• Dengan memanfaatkan fakta bahwa (\nabla\mathrm{dist}{\mathcal{K}}(x)) tidak berubah di antara (x) dan (\mathrm{proj}{\mathcal{K}}(x)), digunakan update yang mencampur estimasi saat ini dan estimasi sebelumnya
• Update (\bar{\epsilon}t=\gamma\epsilon\theta(x_t,\sigma_t)+(1-\gamma)\epsilon_\theta(x_{t+1},\sigma_{t+1})) adalah cara mengoreksi galat langkah sebelumnya dengan estimasi saat ini
• Pada sampel toy model, metode ini konvergen lebih cepat daripada DDIM, dan sampelnya lebih dekat dengan data asli
• Dibandingkan DDIM, sampler ini dapat ditafsirkan sebagai penambahan momentum; lintasannya dapat mengalami overshoot, tetapi bisa konvergen lebih cepat
• Menambahkan noise selama proses generasi secara empiris meningkatkan kualitas sampling
  • Untuk mempertahankan jadwal (\sigma_t) asli, denoise dilakukan terlebih dahulu hingga (\sigma_{t'}) yang lebih kecil, lalu noise (w_t\sim N(0,I)) ditambahkan kembali
  • Saat (\mu=\frac{1}{2}), DDPM sampler direkonstruksi secara tepat
• Update keseluruhan (x_{t-1}=x_t-(\sigma_t-\sigma_{t'})\bar{\epsilon}_t+\eta w_t) menggeneralisasi tiga sampler
  • DDIM: gam=1, mu=0
  • DDPM: gam=1, mu=0.5
  • Sampler estimasi gradien: gam=2, mu=0

Model yang Lebih Besar dan Referensi

• Kode pelatihan di atas dapat digunakan bukan hanya untuk data toy, tetapi juga untuk melatih model difusi gambar dari nol
• Contoh FashionMNIST disediakan sebagai contoh pelatihan pada dataset FashionMNIST yang memperoleh skor peringkat ke-2 berdasarkan FID di leaderboard Papers with Code
• Kode sampling dapat digunakan tanpa modifikasi pada model latent diffusion pralatih
  • Contohnya menggunakan ScheduleLDM(1000) dan ModelLatentDiffusion('stabilityai/stable-diffusion-2-1-base')
  • Kondisi teks ditetapkan sebagai An astronaut riding a horse, lalu sampling dilakukan dengan 50 langkah (\sigma) sebelum latent didekode
• Efek term momentum (\gamma) divisualisasikan dalam perbandingan pada pembuatan text-to-image resolusi tinggi
• Materi tambahan yang layak dibaca
  • What are diffusion models: pengantar model difusi dari sudut pandang waktu diskret yang membalik Markov process
  • Generative modeling by estimating gradients of the data distribution: pengantar model difusi dari sudut pandang waktu kontinu yang membalik persamaan diferensial stokastik
  • The annotated diffusion model: penjelasan detail implementasi model difusi dengan PyTorch

[GN⁺]

Komentar Hacker News

• Saya penulisnya. Saat mencoba memahami model difusi, saya menyadari bahwa kode dan matematikanya bisa sangat disederhanakan, jadi saya membuat tulisan blog ini dan pustaka difusi ini.
  Kalau ada pertanyaan, saya bisa menjawabnya.
  • Sebagai peneliti, ada banyak blog tentang model difusi yang tidak saya sukai, tetapi tulisan ini benar-benar bagus. Tulisan ini langsung masuk ke inti sambil tetap menunjukkan bagian-bagian rumit yang sering menjebak, tanpa membuat pembaca tersesat atau melebar ke mana-mana.
    Saya terutama menyukai pembahasan tentang trajektori, karena itu memberi motivasi untuk memahami bagian yang menyulitkan banyak orang dalam topik seperti scheduler. Meski tidak selengkap tulisan Song atau Lilian, tulisan ini jauh lebih mudah diakses, jadi saya berencana merekomendasikannya kepada orang lain.
    Sebagai referensi, seorang teman pernah menulis implementasi difusi minimal; dari sudut pandang DDPM, ini sedikit lebih “lengkap” dan cukup berguna: https://github.com/VSehwag/minimal-diffusion/
  • Pada contoh gambar terakhir, term momentum tampaknya berdampak buruk pada lukisan digital rumah. Pada gambar gamma = 2.0, pintunya hilang, jadi untuk memahami secara intuitif efek sampler DDIM yang memakai informasi gradien, saya penasaran dengan detail lebih lanjut dari contoh itu.
    Dari pengalaman sedikit bereksperimen dengan prosedur sampling di Stable Diffusion, saya juga ingin melihat perbandingan waktu konvergensi dan jumlah langkah dibanding DDIM. Saya penasaran apakah ada hubungan antara momentum, konvergensi, dan error. Misalnya, akan menarik jika ada perbandingan seperti apakah sampler momentum 16 langkah kurang lebih setara dengan DDIM 20 langkah ± term error.
  • Sepertinya get_sigma_embeds(batches, sigma) tidak menggunakan input pertama. Saya penasaran apakah maksudnya untuk mem-broadcast sigma ke bentuk (batches, 1).
  • Saya penasaran apakah sebagian konsep ini berasal dari prinsip fisika. Apakah ini mirip dengan mengatakan bahwa neural network meniru jaringan saraf biologis, dan apakah ada wawasan dari sudut pandang seperti itu?
• Ada tulisan bagus lain yang juga berjudul Diffusion Models From Scratch: https://www.tonyduan.com/diffusion/index.html
  Tulisan itu membahas detail matematis jauh lebih dalam, sekaligus menyediakan implementasi minimal yang sangat mudah dipahami dengan kurang dari 500 baris kode.
• Bagus karena ada kodenya. Paper difusi terkenal penuh persamaan (https://twitter.com/cto_junior/status/1766518604395155830), tetapi bagi kita yang lain, kode jauh lebih mudah dibaca dan mungkin lebih akurat. Menurut saya semua paper teori harus disertai kode implementasi referensi.
  Akan menarik jika ini juga diperluas ke versi diffusion transformer yang menggerakkan Sora dan model generasi video lain. Tulisan ini dan https://jaykmody.com/blog/gpt-from-scratch/ sepertinya bisa digabung menjadi tulisan pengantar “diffusion transformer dari nol”.
  • Paper difusi memang terkenal penuh persamaan, tetapi sejujurnya kebanyakan peneliti difusi yang saya kenal juga bereaksi sama. Banyak orang menulis ulang persamaan yang sama, dan persamaan-persamaan itu pada dasarnya lebih mirip materi ulasan.
    Sebaliknya, kalau ingin benar-benar mendalami, saya sarankan membaca karya Kingma, Gao, Ricky Tian Qi Chen, para murid Max Welling (Tomczak adalah postdoc, Hoogeboom, dan lainnya), serta kontributor tersembunyi Aapo Hyvärinen. Contoh karya Kingma & Gao yang relatif lebih ringan dan juga terkait dengan paper SD3 ada di sini: https://arxiv.org/abs/2303.00848
    Kekurangannya adalah aksesibilitasnya rendah karena ada banyak ketergantungan pada pengetahuan dan pemahaman riset sebelumnya, tetapi sulit juga menyebut ini sebagai kritik yang bermakna. Bagaimanapun, itu riset, bukan materi edukasi untuk publik umum.
  • Cukup ganti U-net dengan encoder transformer. Hapus embedding dan proyeksikan patch gambar menjadi vektor berukuran n_embd; proses difusinya sendiri bisa tetap sama.
• Tulisan yang bagus, tetapi rasanya ada sifat penting yang terlewat: model difusi memodelkan fungsi skor (turunan dari log probabilitas)[1], dan sampling difusi mirip dengan dinamika Langevin[2]. Menurut saya perspektif ini menjelaskan dengan baik mengapa pelatihannya lebih mudah daripada GAN. Karena tujuan pemodelannya lebih mudah.
  [1] https://yang-song.net/blog/2021/score/
  [2] https://lilianweng.github.io/posts/2021-07-11-diffusion-mode...
  • Benar. Tulisan-tulisan blog ini memberikan interpretasi model difusi yang berbeda dari sudut pandang “proyeksi ke data” yang dijelaskan di artikel utama. Ini bisa dilihat sebagai beberapa cara berbeda untuk menafsirkan tujuan pelatihan dan proses sampling yang sama.
    Dalam sudut pandang kami, alasan model difusi mudah dilatih adalah karena tujuan pelatihannya memprediksi gradien fungsi jarak yang dihaluskan, bukan memprediksi gradien fungsi jarak yang eksak. Sampling model difusi mirip dengan melakukan beberapa langkah gradien aproksimasi.
    Untuk memahami model difusi lebih dalam, saya sarankan membaca semua tulisan blog seperti ini dan mempelajari interpretasi yang berbeda-beda.
• Sangat menarik. Iterative alpha-(de)Blending[1] langsung terlintas di kepala. Karya itu juga mencoba membangun model difusi yang secara konseptual lebih sederhana, dan sampai pada formulasi sebagai proses proyeksi iteratif aproksimatif.
  Namun pendekatan dalam tulisan ini tampaknya memungkinkan eksperimen yang lebih menarik, seperti analisis error denoiser.
  [1] https://arxiv.org/pdf/2305.03486.pdf
• Penjelasan teorinya bagus. Penjelasannya tampak independen dari dataset, tetapi saya penasaran dengan bagian konkret dari generasi gambar yang sebenarnya.
  Misalnya, mengapa generator gambar sulit membuat tuts piano? Untuk membuat struktur tuts hitam yang bergantian dua dan tiga, sepertinya perlu merepresentasikan constraint jarak menengah dengan lebih baik.
  • Ini sama seperti masalah jari. Jumlah, ukuran, sudut, posisi, dan seterusnya harus semuanya benar setiap kali; kalau satu saja salah, orang akan sangat cepat menyadarinya. Berbeda dengan objek seperti cabang pohon, di mana orang tidak terlalu menyadari jika posisi percabangannya “salah”.
• Apakah sebagian dari ide difusi adalah memperbanyak data pelatihan secara besar-besaran? Maksudnya, kita jadi bisa membandingkan gambar yang didifusikan secara acak dengan gambar aslinya yang belum didifusikan?
• Semua model machine learning adalah konvolusi. Lihat saja nanti.
  • Sepertinya Anda sudah beberapa kali mengatakan ini; bisa jelaskan lebih detail? Misalnya, rasanya sulit melihat reinforcement learning sebagai konvolusi.