擴散模型（Diffusion Model）最近在圖像生成領域大火。而在擴散模型中，帶有U-Net的卷積神經網絡居于統治地位。U-ViT網絡是将在圖像領域熱門的Vision Transformer結合U-Net，應用在了Diffision Model中。本文将從Vision Transformer出發，分析U-ViT這篇CVPR2023的Paper并記錄一些感想。
Paper：All are Worth Words: A ViT Backbone for Diffusion Models
Code：https://github.com/baofff/U-ViT
一、Vision Transformer（ViT）
ViT是第一個将标準的transformer block應用在了視覺領域中的網絡。在視覺領域中應用transformer最大的難點在于，如果将圖片的每個像素點看作一個單詞，那麽整張圖像的像素點太多，會導緻序列太長，訓練起來十分昂貴或者說根本無法訓練。以224x224x3的圖像爲例，寬W，高H分别爲224，可見光有RGB三個通道（channel）。那麽一張圖片的長度L即爲15w+。顯然這是不可接受的。
ViT對此的解決方法呢是借鑒了Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self attention and convolutional layers. In ICLR, 2020.這篇論文的思想。将一幅224x224的圖像進行Patch化操作，也就是将一副圖片分爲若幹個小塊。如圖1中左下角所示。
圖1 ViT的結構
舉個例子更方便理解。假設圖像X的shape爲[224, 224, 3]，Patch Size爲16，shape爲[16, 16, 3]，那麽小Patch塊的數量N= HW/PP =（224x224）/（16x16）= 196，也就是有196個小Patch塊，輸入的圖片（單詞）長度L爲197（加1因爲還要加上一個可學習的class embedding，用于分類）。每個Patch塊的像素總量爲16x16x3=768，768即爲d_model。然後再經過一個線性投影層Linear Projection，獲得圖片的特征，最後獲得的Patch Embedding的shape爲[197, 768]。這樣就可以輸入transformer進行訓練了。
ViT的數據流總的來說如圖2所示。在此就不贅述了。如果有誤，可以在評論區，歡迎指正！
圖2 ViT的數據流
其中，E是Linear Project的shape，可以根據d_model變換，主要用于提取圖像的特征以及将圖片映射（map）到Patch Embedding中。注意，位置爲0的embedding就是class token，用于分類。多頭自注意力最後進行的是concat操作而不是add，将12個k，q，v組合在一起恢複原來的shape。
二、U-ViT
U-ViT的創新點是在于将transformer替換掉了Difussion Model中原來的帶有U-Net設計的CNN，并且在transformer中也應用了U-Net的long skip結構。實現了利用transformer進行圖像生成的任務。U-ViT的結構圖如圖3所示。
圖3 U-ViT的結構
同樣的，也是和ViT類似，對圖片進行一個Patch化的操作。不同的是引入了time和condition作爲新的token進行輸入。并且在淺層的transformer和深層的transformer中引入long skip，爲的是将低水平的特征在深層也可以得到應用，這是有利于Difussion Model的像素級預測任務。
作者對long skip的設置方式、time的加入方式、放置3x3Conv的位置、Patch Embedding和Positional Embedding進行了消融實驗（也就是通過對比，看看怎樣設置更好）。實驗結果圖我就不放上來了，直接說結論。
設hm，hs分别爲主支和側支的Embeddings，有五種情況。
- Linear[ Concat(hm，hs) ]，先組合再線性層。
- Add(hm，hs)，直接相加。
- Add[ hm，Linear(hs) ]，hm加上linear後的hs。
- Linear[ Add(hm，hs) ]，先相加再線性層。
- 無long skip，也就是隻有hm。
其中，方法1效果最好，方法2最差。對于方法5，方法3和4均有提升。也就是說，long skip是必要的，但是對于hs需要獲得它線性投影後的信息，因爲在transformer中的加法器具有殘差結構，已經可以獲得低水平的線性信息了。
對于time token的加入方式嘗試了兩種。一種是直接作爲token，另一種是自适應層歸一化AdaLN。自适應層歸一化是将time先經過一個Linear Proj得到ys和yb。再利用這兩個參數和其他embedding做相關計算。第一種效果更好。
對于模型最後輸出層的3x3卷積，也有兩種位置。第一種是放在Linear Proj之後，第二種是之前。放在前面效果好。對于Patch Embedding，通過Linear Proj将Patch映射到Embedding的效果比用3x3 、1x1的Conv的方式要好。對于Positional Embedding，1d的位置編碼比2d的效果更好，并且對比不用位置編碼，使用位置編碼生成的圖像更合理。這也證明了Positional Information的重要性。
後續再閱讀一下代碼，看看數據流是如何實現的并利用MNIST數據集複現一個小型的U-ViT。
研0不能提前去學校，先用一下自己的1650學習先...
