# 【論文筆記】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](https://hackmd.io/@tsen159/token_merging)。

### Unmerging

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

給定兩個 tokens $x_1, x_2$，如果我們合併兩個 tokens 使得
$$
x_{1, 2}^* = (x_1 + x_2) / 2
$$

定義 unmerging 的方式為
$$
x^\prime_1 = x^*_{1, 2}, x^\prime_2 = x^*_{1, 2}
$$

這樣也許會造成資訊的遺失，但作者解釋因為選擇合併的 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 任務上。