“又西三百五十里曰天山，多金玉，有青雄黃，英水出焉，而西南流注於湯谷。有神鳥，其狀如黃囊，赤如丹火，六足四翼，渾敦無面目，是識歌舞，實惟帝江也。”——《山海經》
基於 Transformer 架構的大語言模型在 NLP 領域取得了令人驚豔的效果，然而，Transformer 中自注意力帶來的二次複雜度使得大模型的推理成本和記憶體佔用十分巨大，特別是在長序列的場景中。
此前，研究者們提出了線性 Transformer、Mamba、RetNet 等。這些方案可以大幅降低 Transformer 計算成本，並且取得媲美原有模型的精度，但是由於架構更換，模型重訓練帶來的巨大成本令人望而卻步。
為了解決這一問題，最近的一篇論文提出了一種基於頻域的大語言模型架構 — 帝江（源於山海經的一種神話生物，以跑得快而聞名），同時解決了現有大模型的兩大痛點：推理成本和訓練成本。
論文地址：https://arxiv.org/abs/2403.19928
開源連結：https://github.com/YuchuanTian/DiJiang
該論文基於頻域自注意力變換核，尋找到一種原始自注意力的線性逼近，使得原有的 Transformer 模型可以經過少量資料（1/10-1/50）的微調，可以近乎無損地變形為論文提出的帝江模型。具體來說，在 LLaMA2-7B 上僅僅需要使用 40B 左右的訓練資料，就可以取得最多 5 倍的推理加速，且在各個評測集上取得相當的精度。
DiJIang-7B 模型和 LLaMA-7B 的精度對比
DiJIang-7B 模型和 LLaMA-7B 的速度對比
研究背景
Transformer 架構自從推出以來，徹底革新了自然語言處理（NLP）領域，並在多種任務中取得了傑出成果。這一成功導致了大型語言模型（LLMs）主導的時代的到來，在這個時代中，Transformer 結構被放大以處理越來越複雜的任務。然而，這種規模的擴大也帶來了巨大的計算需求，特別是由於需要每個 token 之間的計算的自注意力機制。
面對更高效 Transformer 模型的迫切需求，研究者們提出了線性 Transformer、Mamba、RetNet 等方案，雖然這些方案可以大幅降低 Transformer 計算成本，並且取得媲美原有模型的精度，但是由於架構更換，模型重訓練帶來的巨大成本令人望而卻步。
然而，大多數現有的最佳化 Transformers 方法，特別是與最佳化注意力機制有關的，需要對模型從頭重新訓練。這一重新訓練過程是一個巨大的挑戰，特別是對於引數龐大的模型，需要大量的計算資源和時間投入。例如，像 LLaMA-7B 這樣的大型模型的訓練需要大約 8 萬多 GPU hours。儘管有部分研究如 Performer 努力尋找注意力機制的快速近似方法，但這些方法在大型語言模型中還沒有得到徹底的驗證。
為了解決大型語言模型中快速注意力近似的問題，論文對現有的線性注意力方案和自注意力近似方案進行了徹底的分析。論文發現，這些方法中近似誤差的主要來源是基於蒙特卡洛方法的取樣。因此，論文提出採用加權擬蒙特卡洛取樣來代替蒙特卡洛取樣進行對映，論文進一步引入頻域離散餘弦變換（DCT）來作為擬蒙特卡洛取樣的值，從而高效且準確地將 Transformer 的 query 和 key 對映到頻域。使得注意力機制中的 softmax 操作可以被去除，達到線性的計算複雜度。論文還從理論上證明了，這種頻域對映是與原始注意力機制的一個近似等效，從而使得帝江模型可以不需要從頭開始訓練，只需要少量資料就可以從 Transformer 的引數中進行微調繼承。論文的實驗表明，論文的方法達到了與原始 Transformer 相當的效能，但訓練成本大大減少（<1/10），同時也受益於更快的推理速度（在不同模型上最高約 10 倍）。
方法介紹
論文首先回顧了 Attention 的計算方式：
其中是一句話中 token 的數目，d 是隱藏層的維度，傳統的 Attention 計算複雜度是。
為了減少 Attention 的計算複雜度，線性 Attention 方案希望將 softmax 函式去掉，這樣 K 和 V 的計算可以提前進行，從而使得計算複雜度變為，由於 n 通常要遠大於 d，因此在變化後計算複雜度可以被大幅減小。例如，Performer 採用了 PRF 核來逼近原始 Attention 的計算，具體為：
然而，由於蒙特卡洛方案存在的近似誤差，Performer 等方案常常要將隱藏層從維度對映為更大的維度，這導致了線性注意力帶來的計算複雜度變為，使得計算加速的收益減少。
為了解決這個問題，論文首先提出一種基於加權擬蒙特卡洛的方案，具體的，論文提出了一種新的 WPFF 核對映：
和 PRF 對映不同，WPFF 核對映在兩點上進行了改進：1. 將原有的隨機對映 w 變為給定的均勻正交變換 v 和其模長部分 t ，即使用擬蒙特卡洛變換來代替蒙特卡洛變換，減少逼近誤差從到
。2. 使用加權矩陣 D 來對對映進行加權求和，減少蒙特卡洛對映的誤差。
論文提供了理論證明，來表明提出的 WPFF 對映核是一種更優的對映方式，具體的證明內容詳見論文附錄：
基於 WPFF 核，論文又進一步對其進行改進，由於給定的均勻正交變換 v 可以使用任意的均勻正交變換，論文提出使用頻域 DCT 變換來進行計算，由於 DCT 變換具有特殊的形式，其計算複雜度僅為，相比其他的正交變換
要來的更低，最終，論文使用的 WDCF 對映為：
最終，帝江模型的自注意力計算被代替為：
帝江模型和傳統自注意力計算的區別
上圖展示了帝江模型和傳統自注意力計算的區別，在 Transformer 的注意力機制中，key 和 value 的計算透過快速離散餘弦變換（DCT）高效地對映到頻域。這種對映有效地消除了 softmax 操作，從而顯著降低了 Transformer 的計算複雜度。
實驗結果
不同模型大小的對比
上表展示了提出的帝江模型在不同大小的 scale 上的結果，可以看到，提出的帝江模型可以取得和原始模型基本相同的精度，並且擁有更快的推理速度和更低的訓練成本，顯著解決了現有 LLM 遇到的訓推成本過大的問題。此外，模型在 1B 的模型量級上超越了 1.3B 大小的 Mamba 模型。需要注意的是，儘管傳統 Transformer 可以透過 Flash Attention 的方式進行進一步加速，但由於針對帝江模型的加速框架尚未開發，為了公平對比模型本身的速度，推理速度的測試都是在模型都不使用加速框架的前提下進行的。
與不同 Transformer 改進方案精度對比
論文還展示了帝江和其他 Transformer 模型的改進方案進行了進一步的對比，可以發現，帝江模型具有比其他模型更好的效果，這得益於其透過更好的核對映近似了原始的 Transformer 模型計算。
論文還同時提供了帝江 - 7B 模型的續寫樣例展示，可以看到，帝江 - 7B 的續寫結果，和 LLaMA2-7B 相比毫不遜色，甚至條理性上要略勝一籌。
總結
論文提出了一種新的 LLM 架構：帝江，在 7B 以下的模型量級，所提出的模型可以大幅降低 LLM 所需的訓練和計算成本，為未來 LLM 的高效部署提出了一種新的思路。帝江架構是否會在更大的模型與多模態 VLM 等其他 Transformer 的應用領域中大放光彩，讓我們拭目以待。