Token Fusion：連接Token修剪與合併，實現高效視覺Transformer

1. 簡介與概述

視覺Transformer (ViTs) 已經徹底改變咗電腦視覺領域，但由於自注意力機制嘅計算複雜度同輸入token數量呈二次方關係，導致運算成本極高。呢篇論文《Token Fusion: Bridging the Gap between Token Pruning and Token Merging》介紹咗Token Fusion (ToFu)，呢種混合方法會根據模型行為，動態選擇修剪或合併token，以優化效率同準確度之間嘅取捨。

核心見解係，無論係修剪（丟棄token）定合併（平均token）都唔係萬能嘅最佳方案。論文提出咗一種原則性方法，為每一層選擇合適嘅操作，並配合一種稱為MLERP（多token線性插值）嘅新穎合併技術，以解決標準平均合併中嘅分佈偏移問題。

2. 核心方法：Token Fusion (ToFu)

ToFu建基於分析模型對插值輸入嘅反應，從而判斷佢係適合合併定修剪。

2.1. 修剪與合併嘅兩難

作者確定咗一個關鍵標準：模型線性度。如果模型層對插值輸入嘅反應接近線性（例如，$f(\alpha x_1 + (1-\alpha)x_2) \approx \alpha f(x_1) + (1-\alpha)f(x_2)$），咁通過平均化嚟合併相似token係有效嘅，並且可以保留信息。然而，喺早期或較深嘅非線性層中（如佢哋圖1所示），輸入空間嘅線性插值會導致高度非線性嘅輸出，令平均合併出現問題，並可能導致分佈偏移。喺呢啲情況下，修剪較唔重要嘅token係一個更安全（雖然會損失信息）嘅替代方案。

2.2. ToFu框架

ToFu喺每個transformer區塊中運作：

Token重要性評分： 為每個token分配一個重要性分數（例如，基於注意力範數或梯度）。
線性度評估： 評估該層嘅近似線性度，通常通過經驗或輕量級探測器得出。
自適應操作： 對於目標token縮減比率：
- 喺高線性度區域：將最唔重要嘅token同佢哋最相似嘅重要鄰居合併。
- 喺低線性度區域：直接修剪最唔重要嘅token。

咁樣就創建咗一個動態、具情境感知嘅壓縮流程。

2.3. MLERP：保持範數嘅合併技術

為咗改進簡單平均法，作者提出咗MLERP，一種適用於合併$K$個token嘅球面線性插值 (SLERP) 變體。對於範數為$n_i = ||t_i||$嘅token $t_1, t_2, ..., t_K$，MLERP首先喺單位球面上插值方向，然後按原始範數嘅加權平均值進行縮放：

$t_{\text{merged}} = \left( \frac{\sum_{i=1}^K w_i n_i}{\| \sum_{i=1}^K w_i \frac{t_i}{n_i} \|} \right) \left( \sum_{i=1}^K w_i \frac{t_i}{n_i} \right)$

其中$w_i$係基於重要性嘅權重。咁樣可以保留特徵嘅統計範數分佈，減輕由簡單平均法引起嘅分佈偏移，從而帶來更穩定嘅效能，特別係喺非線性區域。

3. 技術細節與數學公式

論文將token縮減問題形式化。假設某一層有$N$個輸入token $T = \{t_1, ..., t_N\}$。目標係產生一個縮減後嘅集合$T'$，包含$M < N$個token。

關鍵方程式：

重要性分數： $I(t_i) = ||\text{Attn}(t_i)||_1$ 或基於梯度嘅度量。
相似度度量： 通常係餘弦相似度 $S(t_i, t_j) = \frac{t_i \cdot t_j}{||t_i|| \, ||t_j||}$。
線性度度量 ($\mathcal{L}$)： 通過層輸出偏離輸入線性插值嘅程度嚟測量。低$\mathcal{L}$值利於合併；高$\mathcal{L}$值利於修剪。

ToFu算法可以應用於預訓練模型而無需微調（零樣本），或者通過輕量訓練嚟增強。

4. 實驗結果與效能表現

作者喺圖像分類（使用ViT-B/16、DeiT嘅ImageNet）同圖像生成（潛在擴散模型）任務上評估ToFu。

關鍵效能亮點

分類： 與獨立嘅修剪（例如DynamicViT）或合併（ToMe）方法相比，ToFu喺準確度與FLOPs之間取得更好嘅取捨。例如，喺減少40% FLOPs嘅情況下，ToFu喺ImageNet上嘅top-1準確度損失<0.5%，比ToMe優越約0.3%。
圖像生成： 喺Stable Diffusion中，與ToMe相比，ToFu喺降低運算成本嘅同時保持更高嘅視覺保真度（以FID測量），特別係當縮減大量token時。MLERP合併喺輸出分佈至關重要嘅生成任務中顯示出更明顯嘅優勢。
消融研究： 自適應策略（選擇合併/修剪）被證明優於喺所有層中單獨使用任何一種操作。MLERP持續優於平均合併。

圖表描述（基於論文圖1）： 該圖說明咗ViT層嘅非線性。兩個輸入特徵點 (x1, x2) 被線性插值（彩色線）。繪製咗ViT內部四個唔同MLP層嘅輸出 (f1-f4)。早期同後期嘅MLP輸出 (f1, f4) 顯示出與直線嘅顯著偏差，表明咗強烈嘅非線性。兩個輸入嘅平均值（紫色星號）映射到一個遠離輸出平均值嘅輸出點，直觀地展示咗點解平均合併喺非線性層中會失敗。

5. 分析框架與案例示例

案例：將ToFu應用於預訓練ViT以進行邊緣部署

場景： 開發人員需要喺移動設備上運行ViT-B模型進行實時圖像分類。完整模型太慢。

框架應用：

性能分析： 通過模型運行一個小型校準數據集。對於每個transformer區塊，通過採樣token對並檢查輸出插值誤差嚟計算線性度度量$\mathcal{L}$。
策略圖： 創建一個分析檔案：區塊1-3（低線性度）→ 傾向修剪。區塊4-8（高線性度）→ 傾向MLERP合併。最後幾個區塊（低線性度）→ 傾向修剪。
配置： 設定一個全局token縮減預算（例如35%）。喺低線性度區塊應用修剪，喺高線性度區塊應用MLERP合併，並根據重要性分數得出嘅每區塊預算進行調整。
評估： 部署壓縮後嘅模型。與一刀切嘅方法相比，自適應方法確保咗最小嘅準確度下降，因為佢避免咗喺敏感嘅非線性層中進行激進合併。

呢個例子展示咗ToFu作為一個結構化壓縮框架嘅實用性，而不僅僅係一個單一嘅算法。

6. 未來應用與研究方向

多模態Transformer： 將ToFu擴展到視頻、音頻或多模態（例如CLIP、Flamingo）Transformer，呢啲領域嘅token動態更複雜。
硬件感知協同設計： 針對特定AI加速器（NPU、GPU）優化ToFu決策算法（修剪/合併）同MLERP實現，以最大化實際加速效果。
與其他技術整合： 將ToFu與量化、知識蒸餾或高效注意力機制（如Linformer）結合，以獲得複合嘅效率提升。
自動超參數搜索： 使用神經架構搜索 (NAS) 或強化學習自動確定最佳嘅每層修剪/合併比率同線性度閾值。
超越視覺領域： 探索其喺大型語言模型 (LLM) 中用於序列壓縮嘅功效，儘管token語義有顯著差異。

7. 參考文獻

Dosovitskiy, A., et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR 2021.
Bolya, D., et al. "Token Merging: Your ViT But Faster." ICLR 2023 (ToMe).
Wang, Y., et al. "DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification." NeurIPS 2021.
Rombach, R., et al. "High-Resolution Image Synthesis with Latent Diffusion Models." CVPR 2022.
Zhu, J.Y., et al. "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks." ICCV 2017 (CycleGAN).
Vaswani, A., et al. "Attention Is All You Need." NeurIPS 2017.

8. 專家分析與關鍵見解

核心見解： ToFu唔只係另一個壓縮工具；佢正式承認咗transformer層係異質嘅。用相同嘅壓縮原語處理所有層係天真嘅。論文嘅精妙之處在於其診斷方法——測量層線性度以開出正確嘅「手術」處方（修剪或合併）。呢種做法令人聯想到現代編譯器如何分析代碼以應用優化，呢種複雜程度喺ML效率研究中經常缺失。

邏輯流程： 論證非常有力：1) 展示平均合併喺非線性層中失敗（圖1）。2) 提出一個度量標準嚟檢測呢種失敗模式（線性度）。3) 使用該度量標準嚟路由token。4) 用MLERP修復失敗嘅操作（平均合併）。從問題識別到多組件解決方案嘅流程清晰且合乎邏輯。

優點與缺點：
優點： 混合方法喺理論上合理，並喺多個任務上得到實證驗證。MLERP係對一個實際問題（範數崩潰）嘅簡單而巧妙嘅修復。零樣本適用性對於部署現有模型係一個主要嘅實際優勢。
缺點： 論文稍微低估咗「線性度評估」嘅開銷。佢係預先計算嘅分析檔案（靜態）定係實時計算（動態開銷）？MLERP嘅好處雖然明顯，但喺分類任務中似乎適中；佢嘅真正價值似乎更明顯地體現喺生成任務中，呢點與擴散模型文獻中嘅發現一致，即輸出分佈至關重要。雖然比較公平，但可以更積極地對比最先進嘅訓練後量化方法，後者提供正交嘅好處。

可行見解： 對於從業者：立即採用ToFu/MLERP作為你嘅ViT首選token縮減方法，特別係對於生成任務。 佢取代咗ToMe作為默認嘅合併策略。對於研究人員：「層感知壓縮」範式係關鍵要點。未來嘅工作應該集中於自動檢測模型嘅壓縮友好區域與壓縮敏感區域，或許可以從CNN中嘅網絡修剪工作或CycleGAN等GAN中嘅模式崩潰分析中汲取靈感。下一個前沿係利用呢類診斷研究嘅見解來指導架構搜索，從而設計出本質上高效嘅模型，超越單純嘅事後壓縮。