細觀(guān)Diffusion隱空間里UNet的訓練流程(一)

在上一期里，您已經(jīng)領(lǐng)會(huì )Encoder 和Decoder 兩者的涵意和功能了。請您看看圖1 里，在兩者之間有個(gè)空間，通稱(chēng)為：隱空間(latent space)。

圖1 （圖源：https://arxiv.org/pdf/1909.11573.pdf）

本期以SD(Stable Diffusion) 模型為例，來(lái)欣賞其隱空間里的架構之美。由于Diffusion 隱空間里的主角是：UNet模型。于是，先來(lái)細觀(guān)Diffusion 隱空間里UNet模型的訓練流程。

1 簡(jiǎn)介UNet模型

在本專(zhuān)欄里，曾經(jīng)詳細介紹過(guò)AE(autoencoder)，相信您對AE 已經(jīng)很熟悉了， 不再重復介紹AE模型。

UNet是AE模型的延伸型式。由于典型AE 模型的特性是： 其前段的Encoder( 編碼器) 是特征提取的過(guò)程， 會(huì )過(guò)濾掉一些信息。這些被過(guò)濾掉的信息， 就沒(méi)有辦法傳遞到后段的Decoder( 解碼器) 了。

于是，UNet 就來(lái)延伸AE 模型，將原有的Encoder與Decoder 之間增加了一些連結，例如Encoder 的每一層輸出都與對映層級的譯碼器連接。因此使得編碼器每一層的信息，額外輸入到譯碼器的對映層，讓Decoder在重建的過(guò)程中，比較不會(huì )遺失掉重要的信息了。而Unet 模型的架構如其名呈現一個(gè)U 字形( 圖2)。

圖2

（來(lái)源：https://www.researchgate.net/)

如此，UNet 使用對映層級的連接來(lái)保留小且精細尺度的特征。

2 Diffusion隱空間里的UNet角色

在A(yíng)IGC 潮流中，SD(Stable Diffusion) 產(chǎn)品的推出是AIGC 圖像生成發(fā)展歷程中的一個(gè)里程碑，提供了高性能模型，能快速生成創(chuàng )意十足的圖像( 圖3)。

圖3 （來(lái)源：https://jalammar.github.io/illustrated-stable-diffusion/)

SD 提供兩項主要功能：輸入文本(Text) 提示來(lái)生成圖像( 即text2img)；輸入圖像，然后根據文本描述來(lái)修改圖像( 即輸入text + img)。在SD 隱空間里，UNet 扮演關(guān)鍵性角色( 圖4)。

圖4 （來(lái)源：https://jalammar.github.io/illustrated-stable-diffusion/)

在SD 隱空間里， 使用了1 個(gè)UNet 模型， 并搭配1 個(gè)時(shí)間調度(scheduling) 器，來(lái)?yè)螆D像生成的核心任務(wù)。擴散(diffusion) 一詞描述了SD 隱空間里進(jìn)行的圖像生成情形，整個(gè)過(guò)程都是在隱空間里逐步推進(jìn)(stepby step) 的，每一步都會(huì )增加更多的噪音(noise)。例如，我們準備了一張圖像，接著(zhù)產(chǎn)生隨機噪音加入到圖像中，然后將這含有噪音的圖像輸入給UNet，讓它來(lái)預測噪音。

為了加速圖像的生成效率，SD 使用VAE 的Encoder 先將原圖像壓縮為4×64×64型式的隱空間數組(array)，然后才展開(kāi)上述的擴散、生成過(guò)程。換句話(huà)說(shuō)，整個(gè)擴散過(guò)程都是在隱空間里進(jìn)行的。

3 認識噪音(Noise)

剛才提到了，我們將隨機噪音添加到原圖像里，成為含噪音圖像(noisyimage)。然后把它輸入給UNet 模型，讓它預測出此圖像所含的噪音。在SD里，使用數學(xué)運算，根據時(shí)間步數(time step) 來(lái)決定將多少噪音量添加到原圖像里( 圖5)。

圖5 來(lái)源(https://jalammar.github.io/illustrated-stable-diffusion/)

圖4 里的step-1 是： 準備1 張原圖像， 并使用VAE的Encoder將原圖壓縮為4x64x64 型式的數組。接著(zhù)，step-2 是：產(chǎn)生1 張也是4x64x64 型式的隨機噪音數組；step-3 是：選取1 個(gè)隨機整數( 如圖里的0~3)，即是時(shí)間步數(time step)，并由此整數( 即時(shí)間步數)來(lái)決定欲添加的噪音量；step-4 是：將所決定的噪音量實(shí)際添加到原圖像的壓縮數組里，就成為含有噪音的圖像了。

4 訓練UNet模型

4.1 輸入1張訓練圖像

剛才提到：我們會(huì )根據時(shí)間步數(time step) 來(lái)計算出多少噪音量將添加到原圖像里。也就是，我們掌握了兩項數據：含噪音的圖像(noisy image) 數組；時(shí)間步數(time step)。我們就把這兩項數據輸入給UNet 模型，來(lái)訓練它，讓它學(xué)習預測所添加的噪音，所以它輸出是：預測噪音(predicted noise)。在訓練過(guò)程中，希望它輸出的預測噪音愈接近所添加的噪音就愈好，也就是預測誤差值(loss) 愈小愈好( 圖6)。

圖6 (來(lái)源：https://medium.com/@steinsfu/diffusion-model-clearlyexplained-cd331bd41166)

4.2 訓練很多回合

在實(shí)際訓練UNet模型時(shí)，把全部圖像的壓縮數組都輸入給UNet，讓它學(xué)習一遍，通稱(chēng)為：訓練1 回合(epoch)。在每一回合里，都會(huì )為每一張圖像搭配1 個(gè)隨機整數( 即時(shí)間步數) 值。并根據時(shí)間步數值來(lái)計算出噪音量，于是得到兩項數據：含噪音的圖像數組和時(shí)間步數。然后，把這兩項數據輸入給UNet 模型，來(lái)訓練它，讓它學(xué)習預測噪音( 圖7)。

圖7 （來(lái)源：https://medium.com/@steinsfu/diffusion-model-clearlyexplained-cd331bd41166)

上述SD 的UNet 訓練方法比較特殊之處是：這UNet的功能是輸出是噪音，而不是去除噪音。

5 結束語(yǔ)

本期說(shuō)明了SD 隱空間里UNet 模型的角色及其訓練步流程的起步。下一期，將繼續介紹及欣賞SD 隱空間的架構之美，也領(lǐng)會(huì )AIGC 生成藝術(shù)的源點(diǎn)。

（本文來(lái)源于《電子產(chǎn)品世界》雜志2023年8月期）