來(lái)自微軟亞洲研究院的研究人員發(fā)現(xiàn)對(duì) Token 進(jìn)行傅里葉變換數(shù)學(xué)上等價(jià)于用超大尺寸自適應(yīng)卷積核進(jìn)行 Token 融合，而前者的計(jì)算代價(jià)遠(yuǎn)低于后者?；谠摵诵陌l(fā)現(xiàn)，研究員設(shè)計(jì)了一種名為 Adaptive Fourier Filter（AFF）的輕量高效 Token Mixer，將 Token 融合的計(jì)算復(fù)雜度從 O (N^2) 降低到 O (N log N)，并以此為基礎(chǔ)算子構(gòu)建了輕量級(jí)視覺(jué)神經(jīng)網(wǎng)絡(luò)主干 AFFNet。該主干網(wǎng)絡(luò)能夠以 5.5M 的模型參數(shù)量，在 ImageNet-1K 上實(shí)現(xiàn) 79.8% 的準(zhǔn)確率。

1. 背景

近年來(lái)，基于 Transformer、Large-kernel CNN 和 MLP 三種視覺(jué)主干網(wǎng)絡(luò)在廣泛的 CV 任務(wù)中取得了顯著的成功，這要?dú)w功于它們?cè)谌址秶鷥?nèi)的高效信息融合能力。

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

2. AFF Token Mixer: 輕量、全局、自適應(yīng)

為了解決計(jì)算昂貴的問(wèn)題，研究人員構(gòu)建了一種名為 Adaptive Fourier Filter（AFF）的高效全局 Token 融合算子。它通過(guò)傅里葉變換將 Token 集合變換到頻域，同時(shí)在頻域?qū)W習(xí)到一個(gè)內(nèi)容自適應(yīng)的濾波掩膜，對(duì)變換到頻域空間中的 Token 集合進(jìn)行自適應(yīng)濾波操作。

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

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

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

眾所周知，動(dòng)態(tài)卷積的計(jì)算開(kāi)銷(xiāo)大，大空間分辨率的動(dòng)態(tài)卷積核的使用開(kāi)銷(xiāo)對(duì)于高效 / 輕量級(jí)網(wǎng)絡(luò)設(shè)計(jì)似乎就更加不可接受。但是本文所提出的 AFF Token Mixer 卻可以作為同時(shí)滿足以上優(yōu)點(diǎn)的低功耗等效實(shí)現(xiàn)，將復(fù)雜性從 O (N^2) 降低到 O (N log N)，顯著提高了計(jì)算效率。

圖 1：AFF 模塊和 AFFNet 網(wǎng)絡(luò)示意圖。

3. AFFNet：輕量級(jí)視覺(jué)網(wǎng)絡(luò)新主干

研究人員將 AFF Token Mixer 作為主要神經(jīng)網(wǎng)絡(luò)操作算子，構(gòu)建了一個(gè)輕量級(jí)神經(jīng)網(wǎng)絡(luò)，稱為 AFFNet。大量實(shí)驗(yàn)表明，AFF Token Mixer 在廣泛的視覺(jué)任務(wù)中實(shí)現(xiàn)了優(yōu)越的準(zhǔn)確性和效率權(quán)衡，包括視覺(jué)語(yǔ)義識(shí)別和密集預(yù)測(cè)任務(wù)。

4. 實(shí)驗(yàn)結(jié)果

研究人員在視覺(jué)語(yǔ)義識(shí)別、分割、檢測(cè)等多個(gè)視覺(jué)任務(wù)上對(duì)所提出的 AFF Token Mixer 和 AFFNet 進(jìn)行評(píng)測(cè)，并將其和目前研究領(lǐng)域中最先進(jìn)的輕量級(jí)視覺(jué)主干網(wǎng)絡(luò)進(jìn)行對(duì)比。實(shí)驗(yàn)結(jié)果表明，該工作提出的模型設(shè)計(jì)在廣泛的視覺(jué)任務(wù)上均表現(xiàn)出色，驗(yàn)證了所提出的 AFF Token Mixer 作為新一代輕量高效的 Token 融合算子的潛力。

圖 2：ImageNet-1K 數(shù)據(jù)集上的 Acc-Param, Acc-FLOPs 曲線，與 SOTA 的對(duì)比。

表 1：ImageNet-1K 數(shù)據(jù)集上和 SOTA 的對(duì)比

表 2：下游任務(wù)（視覺(jué)檢測(cè)和分割）和SOTA的對(duì)比。

5. 結(jié)論

此項(xiàng)研究成果從數(shù)學(xué)上證明了隱空間中的頻域變換能起到全局自適應(yīng) Token 融合的作用，是神經(jīng)網(wǎng)絡(luò)中實(shí)現(xiàn)全局自適應(yīng) Token 融合的一種高效能低功耗的等效實(shí)現(xiàn)。為神經(jīng)網(wǎng)絡(luò)中 Token 融合算子的設(shè)計(jì)打開(kāi)了新的研究思路，也為神經(jīng)網(wǎng)絡(luò)模型在存儲(chǔ)和計(jì)算能力有限的邊緣設(shè)備上的部署帶來(lái)了新的發(fā)展空間。