細(xì)觀Diffusion隱空間里UNet的訓(xùn)練流程(一)
在上一期里,您已經(jīng)領(lǐng)會Encoder 和Decoder 兩者的涵意和功能了。請您看看圖1 里,在兩者之間有個空間,通稱為:隱空間(latent space)。
本文引用地址:http://www.ex-cimer.com/article/202308/449711.htm圖1 (圖源:https://arxiv.org/pdf/1909.11573.pdf)
本期以SD(Stable Diffusion) 模型為例,來欣賞其隱空間里的架構(gòu)之美。由于Diffusion 隱空間里的主角是:UNet模型。于是,先來細(xì)觀Diffusion 隱空間里UNet模型的訓(xùn)練流程。
1 簡介UNet模型
在本專欄里,曾經(jīng)詳細(xì)介紹過AE(autoencoder),相信您對AE 已經(jīng)很熟悉了, 不再重復(fù)介紹AE模型。
UNet是AE模型的延伸型式。由于典型AE 模型的特性是: 其前段的Encoder( 編碼器) 是特征提取的過程, 會過濾掉一些信息。這些被過濾掉的信息, 就沒有辦法傳遞到后段的Decoder( 解碼器) 了。
于是,UNet 就來延伸AE 模型,將原有的Encoder與Decoder 之間增加了一些連結(jié),例如Encoder 的每一層輸出都與對映層級的譯碼器連接。因此使得編碼器每一層的信息,額外輸入到譯碼器的對映層,讓Decoder在重建的過程中,比較不會遺失掉重要的信息了。而Unet 模型的架構(gòu)如其名呈現(xiàn)一個U 字形( 圖2)。
圖2
(來源:https://www.researchgate.net/)
如此,UNet 使用對映層級的連接來保留小且精細(xì)尺度的特征。
2 Diffusion隱空間里的UNet角色
在AIGC 潮流中,SD(Stable Diffusion) 產(chǎn)品的推出是AIGC 圖像生成發(fā)展歷程中的一個里程碑,提供了高性能模型,能快速生成創(chuàng)意十足的圖像( 圖3)。
圖3 (來源:https://jalammar.github.io/illustrated-stable-diffusion/)
SD 提供兩項主要功能:輸入文本(Text) 提示來生成圖像( 即text2img);輸入圖像,然后根據(jù)文本描述來修改圖像( 即輸入text + img)。在SD 隱空間里,UNet 扮演關(guān)鍵性角色( 圖4)。
圖4 (來源:https://jalammar.github.io/illustrated-stable-diffusion/)
在SD 隱空間里, 使用了1 個UNet 模型, 并搭配1 個時間調(diào)度(scheduling) 器,來擔(dān)任圖像生成的核心任務(wù)。擴散(diffusion) 一詞描述了SD 隱空間里進(jìn)行的圖像生成情形,整個過程都是在隱空間里逐步推進(jìn)(stepby step) 的,每一步都會增加更多的噪音(noise)。例如,我們準(zhǔn)備了一張圖像,接著產(chǎn)生隨機噪音加入到圖像中,然后將這含有噪音的圖像輸入給UNet,讓它來預(yù)測噪音。
為了加速圖像的生成效率,SD 使用VAE 的Encoder 先將原圖像壓縮為4×64×64型式的隱空間數(shù)組(array),然后才展開上述的擴散、生成過程。換句話說,整個擴散過程都是在隱空間里進(jìn)行的。
3 認(rèn)識噪音(Noise)
剛才提到了,我們將隨機噪音添加到原圖像里,成為含噪音圖像(noisyimage)。然后把它輸入給UNet 模型,讓它預(yù)測出此圖像所含的噪音。在SD里,使用數(shù)學(xué)運算,根據(jù)時間步數(shù)(time step) 來決定將多少噪音量添加到原圖像里( 圖5)。
圖5 來源(https://jalammar.github.io/illustrated-stable-diffusion/)
圖4 里的step-1 是: 準(zhǔn)備1 張原圖像, 并使用VAE的Encoder將原圖壓縮為4x64x64 型式的數(shù)組。接著,step-2 是:產(chǎn)生1 張也是4x64x64 型式的隨機噪音數(shù)組;step-3 是:選取1 個隨機整數(shù)( 如圖里的0~3),即是時間步數(shù)(time step),并由此整數(shù)( 即時間步數(shù))來決定欲添加的噪音量;step-4 是:將所決定的噪音量實際添加到原圖像的壓縮數(shù)組里,就成為含有噪音的圖像了。
4 訓(xùn)練UNet模型
4.1 輸入1張訓(xùn)練圖像
剛才提到:我們會根據(jù)時間步數(shù)(time step) 來計算出多少噪音量將添加到原圖像里。也就是,我們掌握了兩項數(shù)據(jù):含噪音的圖像(noisy image) 數(shù)組;時間步數(shù)(time step)。我們就把這兩項數(shù)據(jù)輸入給UNet 模型,來訓(xùn)練它,讓它學(xué)習(xí)預(yù)測所添加的噪音,所以它輸出是:預(yù)測噪音(predicted noise)。在訓(xùn)練過程中,希望它輸出的預(yù)測噪音愈接近所添加的噪音就愈好,也就是預(yù)測誤差值(loss) 愈小愈好( 圖6)。
圖6 (來源:https://medium.com/@steinsfu/diffusion-model-clearlyexplained-cd331bd41166)
4.2 訓(xùn)練很多回合
在實際訓(xùn)練UNet模型時,把全部圖像的壓縮數(shù)組都輸入給UNet,讓它學(xué)習(xí)一遍,通稱為:訓(xùn)練1 回合(epoch)。在每一回合里,都會為每一張圖像搭配1 個隨機整數(shù)( 即時間步數(shù)) 值。并根據(jù)時間步數(shù)值來計算出噪音量,于是得到兩項數(shù)據(jù):含噪音的圖像數(shù)組和時間步數(shù)。然后,把這兩項數(shù)據(jù)輸入給UNet 模型,來訓(xùn)練它,讓它學(xué)習(xí)預(yù)測噪音( 圖7)。
圖7 (來源:https://medium.com/@steinsfu/diffusion-model-clearlyexplained-cd331bd41166)
上述SD 的UNet 訓(xùn)練方法比較特殊之處是:這UNet的功能是輸出是噪音,而不是去除噪音。
5 結(jié)束語
本期說明了SD 隱空間里UNet 模型的角色及其訓(xùn)練步流程的起步。下一期,將繼續(xù)介紹及欣賞SD 隱空間的架構(gòu)之美,也領(lǐng)會AIGC 生成藝術(shù)的源點。
(本文來源于《電子產(chǎn)品世界》雜志2023年8月期)
評論