ICCV 2023 | 傅里葉算子高效Token Mixer：輕量級視覺網絡新主干

來自微軟亞洲研究院的研究人員發現對 Token 進行傅里葉變換數學上等價于用超大尺寸自適應卷積核進行 Token 融合，而前者的計算代價遠低于后者。基于該核心發現，研究員設計了一種名為 Adaptive Fourier Filter（AFF）的輕量高效 Token Mixer，將 Token 融合的計算復雜度從 O (N^2) 降低到 O (N log N)，并以此為基礎算子構建了輕量級視覺神經網絡主干 AFFNet。該主干網絡能夠以 5.5M 的模型參數量，在 ImageNet-1K 上實現 79.8% 的準確率。

1. 背景

近年來，基于 Transformer、Large-kernel CNN 和 MLP 三種視覺主干網絡在廣泛的 CV 任務中取得了顯著的成功，這要歸功于它們在全局范圍內的高效信息融合能力。

現有的三大主流神經網絡，即 Transformer、CNN 和 MLP，分別通過各自的方式實現全局范圍的 Token 融合。其中，Transformer 網絡中的自注意力機制將 Query-Key pairs 的相關性作為 Token 融合的權重。CNN 通過擴大 kernel 尺寸實現與 transformer 相近的性能。MLP 通過在所有令牌之間的全連接實現另一種強大的范式。所有這些方法都是有效的，但計算復雜度高 (O (N^2))，難以在存儲和計算能力有限的設備上部署，限制了很多模型的應用范圍。

2. AFF Token Mixer: 輕量、全局、自適應

為了解決計算昂貴的問題，研究人員構建了一種名為 Adaptive Fourier Filter（AFF）的高效全局 Token 融合算子。它通過傅里葉變換將 Token 集合變換到頻域，同時在頻域學習到一個內容自適應的濾波掩膜，對變換到頻域空間中的 Token 集合進行自適應濾波操作。

論文《Adaptive Frequency Filters As Efficient Global Token Mixers》：

鏈接：https://arxiv.org/abs/2307.14008

根據頻域卷積定理，原始域中的卷積操作在數學上等價于對應的傅里葉域中的 Hadamard 乘積操作。這使得該工作所提出的 AFF Token Mixer 在數學上等價于使用一個空間分辨率和 Token 集合一樣大小的動態卷積核在原始域中進行 Token 融合 (如下圖右子圖所示), 具有在全局范圍內進行內容自適應 Token 融合的作用。

眾所周知，動態卷積的計算開銷大，大空間分辨率的動態卷積核的使用開銷對于高效 / 輕量級網絡設計似乎就更加不可接受。但是本文所提出的 AFF Token Mixer 卻可以作為同時滿足以上優點的低功耗等效實現，將復雜性從 O (N^2) 降低到 O (N log N)，顯著提高了計算效率。

圖 1：AFF 模塊和 AFFNet 網絡示意圖。

3. AFFNet：輕量級視覺網絡新主干

研究人員將 AFF Token Mixer 作為主要神經網絡操作算子，構建了一個輕量級神經網絡，稱為 AFFNet。大量實驗表明，AFF Token Mixer 在廣泛的視覺任務中實現了優越的準確性和效率權衡，包括視覺語義識別和密集預測任務。

4. 實驗結果

研究人員在視覺語義識別、分割、檢測等多個視覺任務上對所提出的 AFF Token Mixer 和 AFFNet 進行評測，并將其和目前研究領域中最先進的輕量級視覺主干網絡進行對比。實驗結果表明，該工作提出的模型設計在廣泛的視覺任務上均表現出色，驗證了所提出的 AFF Token Mixer 作為新一代輕量高效的 Token 融合算子的潛力。

圖 2：ImageNet-1K 數據集上的 Acc-Param, Acc-FLOPs 曲線，與 SOTA 的對比。

表 1：ImageNet-1K 數據集上和 SOTA 的對比

表 2：下游任務（視覺檢測和分割）和SOTA的對比。

5. 結論

此項研究成果從數學上證明了隱空間中的頻域變換能起到全局自適應 Token 融合的作用，是神經網絡中實現全局自適應 Token 融合的一種高效能低功耗的等效實現。為神經網絡中 Token 融合算子的設計打開了新的研究思路，也為神經網絡模型在存儲和計算能力有限的邊緣設備上的部署帶來了新的發展空間。