華為諾亞頻域LLM「帝江」：僅需1/50訓(xùn)練成本，7B模型媲美LLaMA，推理加速5倍
基于 Transformer 架構(gòu)的大語言模型在 NLP 領(lǐng)域取得了令人驚艷的效果，然而，Transformer 中自注意力帶來的二次復(fù)雜度使得大模型的推理成本和內(nèi)存占用十分巨大，特別是在長(zhǎng)序列的場(chǎng)景中。
此前，研究者們提出了線性 Transformer、Mamba、RetNet 等。這些方案可以大幅降低 Transformer 計(jì)算成本，并且取得媲美原有模型的精度，但是由于架構(gòu)更換，模型重訓(xùn)練帶來的巨大成本令人望而卻步。
為了解決這一問題，最近的一篇論文提出了一種基于頻域的大語言模型架構(gòu) — 帝江（源于山海經(jīng)的一種神話生物，以跑得快而聞名），同時(shí)解決了現(xiàn)有大模型的兩大痛點(diǎn)：推理成本和訓(xùn)練成本。
- 論文地址：https://arxiv.org/abs/2403.19928
- 開源鏈接：https://github.com/YuchuanTian/DiJiang
該論文基于頻域自注意力變換核，尋找到一種原始自注意力的線性逼近，使得原有的 Transformer 模型可以經(jīng)過少量數(shù)據(jù)（1/10-1/50）的微調(diào)，可以近乎無損地變形為論文提出的帝江模型。具體來說，在 LLaMA2-7B 上僅僅需要使用 40B 左右的訓(xùn)練數(shù)據(jù)，就可以取得最多 5 倍的推理加速，且在各個(gè)評(píng)測(cè)集上取得相當(dāng)?shù)木取?/p>
DiJIang-7B 模型和 LLaMA-7B 的精度對(duì)比
DiJIang-7B 模型和 LLaMA-7B 的速度對(duì)比
研究背景
Transformer 架構(gòu)自從推出以來，徹底革新了自然語言處理（NLP）領(lǐng)域，并在多種任務(wù)中取得了杰出成果。這一成功導(dǎo)致了大型語言模型（LLMs）主導(dǎo)的時(shí)代的到來，在這個(gè)時(shí)代中，Transformer 結(jié)構(gòu)被放大以處理越來越復(fù)雜的任務(wù)。然而，這種規(guī)模的擴(kuò)大也帶來了巨大的計(jì)算需求，特別是由于需要每個(gè) token 之間的計(jì)算的自注意力機(jī)制。
面對(duì)更高效 Transformer 模型的迫切需求，研究者們提出了線性 Transformer、Mamba、RetNet 等方案，雖然這些方案可以大幅降低 Transformer 計(jì)算成本，并且取得媲美原有模型的精度，但是由于架構(gòu)更換，模型重訓(xùn)練帶來的巨大成本令人望而卻步。
然而，大多數(shù)現(xiàn)有的優(yōu)化 Transformers 方法，特別是與優(yōu)化注意力機(jī)制有關(guān)的，需要對(duì)模型從頭重新訓(xùn)練。這一重新訓(xùn)練過程是一個(gè)巨大的挑戰(zhàn)，特別是對(duì)于參數(shù)龐大的模型，需要大量的計(jì)算資源和時(shí)間投入。例如，像 LLaMA-7B 這樣的大型模型的訓(xùn)練需要大約 8 萬多 GPU hours。盡管有部分研究如 Performer 努力尋找注意力機(jī)制的快速近似方法，但這些方法在大型語言模型中還沒有得到徹底的驗(yàn)證。
為了解決大型語言模型中快速注意力近似的問題，論文對(duì)現(xiàn)有的線性注意力方案和自注意力近似方案進(jìn)行了徹底的分析。論文發(fā)現(xiàn)，這些方法中近似誤差的主要來源是基于蒙特卡洛方法的采樣。因此，論文提出采用加權(quán)擬蒙特卡洛采樣來代替蒙特卡洛采樣進(jìn)行映射，論文進(jìn)一步引入頻域離散余弦變換（DCT）來作為擬蒙特卡洛采樣的值，從而高效且準(zhǔn)確地將 Transformer 的 query 和 key 映射到頻域。使得注意力機(jī)制中的 softmax 操作可以被去除，達(dá)到線性的計(jì)算復(fù)雜度。論文還從理論上證明了，這種頻域映射是與原始注意力機(jī)制的一個(gè)近似等效，從而使得帝江模型可以不需要從頭開始訓(xùn)練，只需要少量數(shù)據(jù)就可以從 Transformer 的參數(shù)中進(jìn)行微調(diào)繼承。論文的實(shí)驗(yàn)表明，論文的方法達(dá)到了與原始 Transformer 相當(dāng)?shù)男阅?，但?xùn)練成本大大減少（<1/10），同時(shí)也受益于更快的推理速度（在不同模型上最高約 10 倍）。
方法介紹
論文首先回顧了 Attention 的計(jì)算方式：
其中是一句話中 token 的數(shù)目，d 是隱藏層的維度，傳統(tǒng)的 Attention 計(jì)算復(fù)雜度是
。
為了減少 Attention 的計(jì)算復(fù)雜度，線性 Attention 方案希望將 softmax 函數(shù)去掉，這樣 K 和 V 的計(jì)算可以提前進(jìn)行，從而使得計(jì)算復(fù)雜度變?yōu)?img src="https://s4.51cto.com/oss/202404/03/c1c129007ddfac04701036634eeae60375ec4b.webp" alt='華為諾亞頻域LLM「帝江」：僅需1/50訓(xùn)練成本，7B模型媲美LLaMA，推理加速5倍-AI.x社區(qū)' title='華為諾亞頻域LLM「帝江」：僅需1/50訓(xùn)練成本，7B模型媲美LLaMA，推理加速5倍-AI.x社區(qū)' style="width: 66px; visibility: visible;" data-type="inline">
，由于 n 通常要遠(yuǎn)大于 d，因此在變化后計(jì)算復(fù)雜度可以被大幅減小。例如，Performer 采用了 PRF 核來逼近原始 Attention 的計(jì)算，具體為：
然而，由于蒙特卡洛方案存在的近似誤差，Performer 等方案常常要將隱藏層從維度映射為更大的維度，這導(dǎo)致了線性注意力帶來的計(jì)算復(fù)雜度變?yōu)椋沟糜?jì)算加速的收益減少。
為了解決這個(gè)問題，論文首先提出一種基于加權(quán)擬蒙特卡洛的方案，具體的，論文提出了一種新的 WPFF 核映射：
和 PRF 映射不同，WPFF 核映射在兩點(diǎn)上進(jìn)行了改進(jìn)：1. 將原有的隨機(jī)映射 w 變?yōu)榻o定的均勻正交變換 v 和其模長(zhǎng)部分 t ，即使用擬蒙特卡洛變換來代替蒙特卡洛變換，減少逼近誤差從
到
。2. 使用加權(quán)矩陣 D 來對(duì)映射進(jìn)行加權(quán)求和，減少蒙特卡洛映射的誤差。
論文提供了理論證明，來表明提出的 WPFF 映射核是一種更優(yōu)的映射方式，具體的證明內(nèi)容詳見論文附錄：
基于 WPFF 核，論文又進(jìn)一步對(duì)其進(jìn)行改進(jìn)，由于給定的均勻正交變換 v 可以使用任意的均勻正交變換，論文提出使用頻域 DCT 變換來進(jìn)行計(jì)算，由于 DCT 變換具有特殊的形式，其計(jì)算復(fù)雜度僅為
，相比其他的正交變換
要來的更低，最終，論文使用的 WDCF 映射為：
最終，帝江模型的自注意力計(jì)算被代替為：
帝江模型和傳統(tǒng)自注意力計(jì)算的區(qū)別
上圖展示了帝江模型和傳統(tǒng)自注意力計(jì)算的區(qū)別，在 Transformer 的注意力機(jī)制中，key 和 value 的計(jì)算通過快速離散余弦變換（DCT）高效地映射到頻域。這種映射有效地消除了 softmax 操作，從而顯著降低了 Transformer 的計(jì)算復(fù)雜度。
實(shí)驗(yàn)結(jié)果
不同模型大小的對(duì)比
上表展示了提出的帝江模型在不同大小的 scale 上的結(jié)果，可以看到，提出的帝江模型可以取得和原始模型基本相同的精度，并且擁有更快的推理速度和更低的訓(xùn)練成本，顯著解決了現(xiàn)有 LLM 遇到的訓(xùn)推成本過大的問題。此外，模型在 1B 的模型量級(jí)上超越了 1.3B 大小的 Mamba 模型。需要注意的是，盡管傳統(tǒng) Transformer 可以通過 Flash Attention 的方式進(jìn)行進(jìn)一步加速，但由于針對(duì)帝江模型的加速框架尚未開發(fā)，為了公平對(duì)比模型本身的速度，推理速度的測(cè)試都是在模型都不使用加速框架的前提下進(jìn)行的。
與不同 Transformer 改進(jìn)方案精度對(duì)比
論文還展示了帝江和其他 Transformer 模型的改進(jìn)方案進(jìn)行了進(jìn)一步的對(duì)比，可以發(fā)現(xiàn)，帝江模型具有比其他模型更好的效果，這得益于其通過更好的核映射近似了原始的 Transformer 模型計(jì)算。
論文還同時(shí)提供了帝江 - 7B 模型的續(xù)寫樣例展示，可以看到，帝江 - 7B 的續(xù)寫結(jié)果，和 LLaMA2-7B 相比毫不遜色，甚至條理性上要略勝一籌。
總結(jié)
論文提出了一種新的 LLM 架構(gòu)：帝江，在 7B 以下的模型量級(jí)，所提出的模型可以大幅降低 LLM 所需的訓(xùn)練和計(jì)算成本，為未來 LLM 的高效部署提出了一種新的思路。帝江架構(gòu)是否會(huì)在更大的模型與多模態(tài) VLM 等其他 Transformer 的應(yīng)用領(lǐng)域中大放光彩，讓我們拭目以待。
本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心
