【論文筆記】Token Merging for Fast Stable Diffusion

論文連結：
https://arxiv.org/abs/2303.17604
發表於 CVPR 2023 ECV Workshop

相關研究：
Token Merging: Your ViT But Faster (ICLR 2023)
筆記：https://hackmd.io/@tsen159/token_merging
論文：https://arxiv.org/abs/2210.09461

Introduction

現今比較強大的 diffusion models 因為具有 transformer backbone，在進行 inference 的時候計算量都十分大，因此本篇研究利用 Token Merging (ToMe) 的方法，來嘗試加速 diffusion models，並同時減少記憶體的使用量。作者特別指出大部分的加速方法都需要重新訓練，但 ToMe 不需要額外的訓練就可以應用。

實驗結果顯示比起原始的 Stable Diffusion，加上 ToMe 可以加速兩倍，並且僅使用 5.6 倍少的記憶體。另外也可以和現有的其他加速方法一起使用，例如和 xFormers 一起使用的話，可以有 5.4 倍的速度。

Methods

Token Merging (ToMe)

ToMe 的基本想法是透過在每一個 transformer block 中合併數個 tokens，以逐步減少 token 的數量。要達成這件事情，首先會將 tokens 分成一個 source set 和另一個 destination set，並從 source set 當中找出和 destination set 裡最相似的幾個 tokens，將它們合併。更詳細的介紹參考【論文筆記】Token Merging: Your ViT But Faster。

Unmerging

先前 ToMe 是應用在分類問題上，只會需要一個 token 進行預測，但是 diffusion model 必須從每一個 token 上移除 noise，因此在這個任務上，我們會需要每一個 tokens，不能只是不斷地將 token 合併在一起，還需要進行 unmerging。

給定兩個 tokens

$x_1,x_2$，如果我們合併兩個 tokens 使得

$$x_{1,2}^{\ast }=(x_1+x_2)/2$$

定義 unmerging 的方式為

$$x_1^{\prime }=x_{1,2}^{\ast },x_2^{\prime }=x_{1,2}^{\ast }$$

這樣也許會造成資訊的遺失，但作者解釋因為選擇合併的 tokens 已經足夠相似了，因此產生的誤差可以被忽略。

Naive Approach

要在 Stable Diffusion 的 U-Net transformer block 中加入 ToMe，一個最單純的方法就是將 ToMe 套用在每次進入某個 module 之前，並在計算結束之後、skip connection 之前進行 unmerging，如上圖所示。

另外還有一些和使用在 ViT 的原始 ToMe 設計不同的地方：

1. 每一次 merging 都合併一定比例 (r%) 的 tokens
2. 只在每個 block 的最開始計算 token similarity
3. 計算 token similarity 所使用的 feature 是 input
   $x$，而非 key

用 naive approach 就可以看到 token merging 比起 token pruning，輸出的圖片品質好非常多：

然而我們可以看到，用 naive 的方法進行 token merging，當合併的 token 比例變高的時候，會出現圖片內容改變的問題：

因此作者接著提出了一些解決方法。

New Partitioning Method

ToMe 預設將 tokens 分成兩個 sets 的方法是 alternating，也就是交替著分配，這個方式雖然在 ViT 上可以有很好的效果，但在現在這個任務上，會產生 alternating columns，見下圖 (a)。假設我們現在合併 tokens 的比例是 50％，則經過 merging 和 unmerging 之後，相當於是在 row 的方向將解析度降低了一半（因為新的 token 會等於兩者平均）。

第一個可能的解決方法是可以利用 2D stride 來選擇要分配到 destination set 的 tokens，例如選擇 strided 2X2，如上圖 (b)。這個方式大幅改善了輸出圖片的品質。然而 stride 的方法會造成每一次分配到 destination set 的 tokens 都是同一個位置，因此為了解決這個問題，他們嘗試隨機抽取 tokens 進行分配。另外也可以合併兩個方法，變成每 2X2 的區域就隨機選取 token，這樣可以讓選取的 tokens 分佈比較平均。

下表為各種隨機方法的 FID 比較。因為使用 classifier-free guidance，他們發現 conditional 和 unconditional 的 samples 必須要用同樣的方式分配 tokens，才不會讓 FID 過大，因此針對同一個 batch，使用到的 randomness 會是相同的 (fix randomness)。

Experiments and Results

下表是各種不同設計的實驗結果。作者發現將 ToMe 只用在 self-attention 上可以有最好的輸出品質，此外其實不需要在每一個 block 都加上 ToMe，只需要在 token 數量比較多的 block 使用，就可以有加速的作用。另外作者也嘗試在不同階段合併不同比例的 tokens，發現在品質上沒有太大的影響。

下圖以及下表顯示 ToMe for Stable Diffusion 的 qualitative 和 quantitative results。使用 ToMe 合併 60% 的 tokens 時，可以維持差不多的圖片品質，但速度可以快 2 倍，記憶體使用可以少 5.6 倍。

ToMe 也可以和其他 off-the-shelf fast transformer implementations 一起合用，產生加乘的效果：

Conclusion

這篇研究將 Token Merging 應用在 Stable Diffusion 上，可以用更快的時間生成品質好的圖片，而且這個方法不需要額外的訓練就可以使用。作者指出既然 ToMe 應用在 Stable Diffusion 可以有很好的表現，應該也很有機會可以用在其他 dense prediction 任務上。