基于 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 的應用領域中大放光彩，讓我們拭目以待。

本文轉自 機器之心 ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/3zdfuhNr_HjvcJAy8REayA??