FNet：在Transformer編碼器中使用傅立葉轉換混合詞元

1. 簡介與概述

自 Vaswani 等人於 2017 年提出以來，Transformer 架構已成為最先進自然語言處理（NLP）的事實標準。其核心創新——自注意力機制——允許模型在處理每個詞元時，動態地權衡序列中所有詞元的重要性。然而，此機制伴隨著顯著的計算成本，其複雜度隨序列長度（$N$）呈二次方增長（$O(N^2)$），這限制了其在處理長文件或高吞吐量應用中的效率。

本篇論文《FNet：使用傅立葉轉換混合詞元》提出了一種根本性的簡化方案。作者探討了計算昂貴的自注意力子層是否可以被更簡單的線性詞元混合機制完全取代。他們最令人驚訝的發現是，使用標準、無參數化的二維離散傅立葉轉換（DFT），在 GLUE 基準測試上能達到 BERT 模型 92-97% 的準確度，同時在 GPU 上訓練速度快 80%，在 TPU 上處理標準 512 詞元序列時快 70%。

2. 方法論與架構

2.1. 取代自注意力機制

核心假設是：自注意力所執行的複雜、依賴於數據的混合過程，或許可以透過固定的線性轉換來近似或取代。作者首先嘗試了參數化的線性混合層（稠密矩陣）。觀察到有希望的結果後，他們探索了更快速、結構化的線性轉換，最終選定了傅立葉轉換。

2.2. 傅立葉轉換子層

在 FNet 中，標準 Transformer 編碼器區塊中的自注意力子層被一個二維傅立葉轉換所取代。對於輸入表示 $X \in \mathbb{R}^{N \times d}$（其中 $N$ 為序列長度，$d$ 為隱藏維度），混合過程如下：

$\text{FNet}(X) = \mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{hidden}}(X))$

其中 $\mathcal{F}_{\text{hidden}}$ 沿隱藏維度（$d$）應用一維傅立葉轉換，而 $\mathcal{F}_{\text{seq}}$ 則沿序列維度（$N$）應用。僅保留轉換結果的實部。關鍵在於，此子層沒有可學習的參數。

2.3. FNet模型架構

一個 FNet 編碼器區塊保留了標準 Transformer 架構的其餘部分：一個帶有非線性函數（例如 GeLU）的前饋神經網路（FFN）子層、殘差連接以及層正規化。順序為：傅立葉混合子層 → 殘差連接與層正規化 → FFN 子層 → 殘差連接與層正規化。

3. 技術細節與數學公式

對於長度為 $N$ 的序列 $x$，其一維離散傅立葉轉換（DFT）定義為：

$X_k = \sum_{n=0}^{N-1} x_n \cdot e^{-i 2\pi k n / N}$

對於應用於輸入矩陣 $X$ 的二維轉換，其計算方式為兩個連續的一維轉換。使用快速傅立葉轉換（FFT）演算法，將此操作的複雜度降低至 $O(Nd \log N)$（針對序列維度轉換），這對於較大的 $N$ 值而言，顯著優於標準自注意力的 $O(N^2 d)$。

關鍵洞見在於，傅立葉轉換在頻域中對所有輸入詞元進行全域混合，這可能捕捉到與自注意力類似的全域依賴關係，但它是透過一個固定的數學基底，而非學習得來、依賴於數據的基底。

4. 實驗結果與效能表現

4.1. GLUE基準測試結果

FNet 模型（Base 和 Large 尺寸）與對應的 BERT 模型進行了比較評估。結果令人矚目：

FNet-Base 達到了 BERT-Base 平均 GLUE 分數的 92.2%。
FNet-Large 達到了 BERT-Large 平均 GLUE 分數的 97.3%。

這表明，精心調校的自注意力模型的大部分準確度，可以透過簡單的傅立葉混合機制來恢復。

4.2. Long Range Arena (LRA) 基準測試

在專為測試長序列（1k 至 4k 詞元）模型效能而設計的 LRA 基準測試上，FNet 的準確度與最準確的「高效能 Transformer」模型相當。更重要的是，在 GPU 上，它在所有序列長度上的速度都顯著快於最快的模型。

4.3. 速度與效率分析

效能提升相當顯著：

訓練速度： 在 512 序列長度下，比 BERT 在 GPU 上快 80%，在 TPU 上快 70%。
記憶體佔用： 比標準 Transformer 更輕量，對於較小的模型尺寸尤其有益。
擴展性： FFT 的 $O(N \log N)$ 擴展性，使 FNet 在 GPU 上處理長序列時，相較於即使是線性時間（$O(N)$）的注意力近似方法也具有決定性優勢，因為那些方法通常具有較大的隱藏常數因子。

5. 分析框架與案例範例

案例：長文件的文本分類
考慮像法律合約或科學文章分類這樣的任務，其中文件長度經常超過 2000 個詞元。標準 Transformer 模型會因二次方的記憶體和計算成本而難以處理。一個「高效能」的線性 Transformer 或許有幫助，但由於核心化（kernelization）的開銷，在實踐中可能速度緩慢。

FNet 應用： FNet 模型可以高效處理這些長序列。傅立葉子層以 $O(N \log N)$ 的時間全域混合詞元表示。隨後的 FFN 層則可以在這些混合後的表示上建立特徵。對於固定的延遲預算，可以部署比對應 Transformer 更大的 FNet 模型，有可能彌補在較短序列上觀察到的微小準確度差距。

框架要點： FNet 將歸納偏置從「數據驅動的關係權重分配」（注意力）轉變為「固定的全域頻譜混合」。FNet 的成功表明，對於許多 NLP 任務而言，全域組合資訊的能力比特定、學習得來的組合方法更為關鍵。

6. 核心洞見與批判性分析

核心洞見： 皇帝的新衣可能比我們想像的還要少。FNet 的成功對 NLP 的正統觀念提出了挑釁性的挑戰。它證明了自注意力這頭「神聖的牛」——通常被認為是 Transformer 強大能力不可或缺的來源——可以被一個無參數、擁有 150 年歷史的數學運算所取代，僅帶來微小的效能損失，卻獲得巨大的效率提升。這表明 Transformer 能力的相當一部分源於其整體架構（殘差連接、FFN、層正規化）及其全域資訊流動的能力，而非注意力機制本身複雜、學習得來的動態特性。

邏輯流程： 論文的邏輯引人入勝。從昂貴的問題（二次方注意力）出發。假設更簡單的混合可能有效。測試線性層（效果尚可）。意識到像 FFT 這樣的結構化轉換更快且擴展性極佳。測試它——令人驚訝地，效果幾乎一樣好。從問題到迭代解決方案再到驚人發現的流程清晰且科學嚴謹。

優點與缺點：
優點： 效率提升無可否認且具有實際意義。論文在標準基準測試（GLUE、LRA）上進行了嚴謹評估。這個想法極其簡潔優美，具有強烈的「我怎麼沒想到？」的吸引力。它為高效能架構開闢了新的設計空間。
缺點： 準確度差距雖然小，但確實存在，對於追求 SOTA 的應用可能至關重要。論文沒有深入分析傅立葉轉換為何效果如此之好，或者損失了哪些語言特性。有人懷疑其在需要非常細粒度句法推理或複雜多步驟推論（動態注意力至關重要）的任務上，效能可能達到瓶頸。其速度宣稱依賴於具有高度優化 FFT 核心的 GPU/TPU，這是一個隱藏的依賴條件。

可行動的洞見：
1. 對於實務工作者： 在吞吐量、延遲或成本是主要限制，且可接受 3-8% 準確度下降的生產部署中，強烈考慮使用 FNet。它是「足夠好」的大規模文本處理的首選方案。
2. 對於研究人員： 不要止步於傅立葉轉換。這篇論文為探索整個線性轉換（小波、哈特利、DCT）和結構化矩陣作為注意力替代品的領域開了綠燈。核心研究問題變成了：「什麼是最小、最快的混合機制，足以滿足語言理解的需求？」
3. 對於該領域： 這項工作，連同視覺領域的 MLP-Mixer 等同期研究，標誌著一場潛在的「回歸基礎」運動。在經歷了多年架構複雜度不斷增加之後，我們可能正進入一個根本性簡化的時代，質疑哪些組件是真正必要的。它提醒我們要定期挑戰基本假設。

7. 未來應用與研究方向

混合模型： 將 FNet 層與稀疏或局部注意力層結合，可以創造出既高效又能為關鍵推理步驟保留高精度的模型。
模態擴展： 將 FNet 的原理應用於多模態 Transformer（視覺、音訊）。透過傅立葉轉換早期混合跨模態訊號可能非常高效。
硬體-軟體協同設計： 設計針對 FFT 操作優化的專用 AI 加速器，可以使 FNet 類架構在效率至關重要的場景中更具主導地位。
理論理解： 深入分析傅立葉轉換執行了哪些語言功能，以及 FFN 層如何補償缺乏學習得來的注意力，是未來工作的豐富領域。
長上下文建模： FNet 是推動語言模型上下文長度界限的自然候選者，能夠以可管理的計算量處理整本書籍或長對話。

8. 參考文獻

Vaswani, A., et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems.
Devlin, J., et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL-HLT.
Tolstikhin, I., et al. (2021). MLP-Mixer: An all-MLP Architecture for Vision. Advances in Neural Information Processing Systems.
Tay, Y., et al. (2020). Efficient Transformers: A Survey. ACM Computing Surveys.
Wang, S., et al. (2020). Linformer: Self-Attention with Linear Complexity. arXiv preprint arXiv:2006.04768.
Katharopoulos, A., et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. International Conference on Machine Learning.
Google Research. FNet Official Code Repository. https://github.com/google-research/google-research/tree/master/f_net

目錄