序言：當前基於 Transformer 架構的大語言模型人工智慧技術，由於投入大、成本高、人才需求苛刻，導致許多企業望而卻步。動輒幾千萬甚至上億的成本，現實中有幾家企業能夠承擔？真正具有競爭力的技術應當在成本上更低、效率上更高，因此，各大院校和商業公司已不再僅僅關注模型的引數規模，而是在積極探索創新方法，顯著降低大語言模型的研發與使用成本，使得大多數企業也能輕鬆採用。史丹佛大學的這一最新研究成果正是朝著這一目標邁出了一大步——將一個 80 億引數模型的訓練成本降至 20 美元。同時，已有中國企業基於這一研究成果推出了適用於企業的私有 AI 模型及伺服器、前端應用的垂直解決方案，為 AI 在企業中的私有化部署鋪平了道路。
史丹佛大學的一組研究人員推出了LoLCATs，一種線性化標準Transformer LLMs的新方法，大幅降低了計算需求，同時保留了大部分最先進的（SOTA）效能。而這一切只需要幾個小時的GPU計算，總成本不到20美元，最終使模型效能在訓練投入上的效率提高至35,500倍。這聽起來難以置信，但他們今天展示的三種技巧讓這一切成為可能。這項非凡的成就可能很快會成為AI工程團隊追求一流表現的基本技能。
那麼，他們是如何做到的？
標準LLM的問題: 大型語言模型（LLM）為AI行業帶來了巨大的激動與資金，似乎正在走向全球主導的直線路徑，作為一種能讓文明邁向新時代的優越技術（至少，這是他們想讓我們相信的）。令人驚訝的是，這一巨大的賭注迫使大科技公司投資數十億，集中於一種架構型別——Transformer，然而它們效率極低，所以成本巨高，提高Transformer的效率就必然會降低成本。
Transformer是我們所需的一切。 簡單來說，AI不過是一個資料壓縮演算法，輸入資料後學習其中的模式，進而利用這種已掌握的知識做出有用的預測。在現有的實現方法中，沒有一種比Transformer更接近所需的效果，原因有二：
-
它們完全可並行化，非常適合用大量資料訓練模型。
-
模型的規模越大，效果越好，這引發了人們在如何擴大模型模型和訓練預算上的投資和研究狂潮。
然而，以Transformer為基礎的模型雖然具有無與倫比的表現力，並不意味著它們是完美的，因為它們的優點帶來了自身的問題。
平方複雜度問題。 最大的問題在於它們無法壓縮狀態或記憶體。換句話說，如果你希望模型能記住某些資訊，那就必須將其以原始狀態儲存在記憶體中。人類並不是這樣記憶的。人類不會記住所有的事情或經歷的事件；人類僅保留被認為重要的部分。
想象你在讀一本書。如果你在第11章，主角可能會提到書中之前發生的事情，或許是第1章。如果你留心了，你可能會記得，否則可能需要回頭重讀那部分。然而，這對記憶體負擔並不是問題，因為如果第2章中沒有發生什麼特別的事情，你可以隨意忘記它。Transformer並不是這樣處理資訊的。當閱讀第11章時，如果需要回憶第1章的內容，它會立即記住，因為它依然能夠訪問整個11章內容。事實上，如果不是我們在LLM推理中構建的KV快取（KV Cache），Transformer實際上在閱讀每一個新詞時都在重讀之前的內容。
這是否顯得低效？希望是如此，因為確實如此。
但這意味著什麼呢？簡單來說，這種未壓縮的記憶在序列變長時會增長。如果你在第15章，之前的章節會積累，記憶體需求比讀第8章時（只有七個章節）大得多。
更糟的是，由於底層注意力機制的工作方式（為了節省篇幅這裡不展開討論），我們需要為每個章節的每個詞儲存兩種資訊（鍵和值）。這意味著Transformer的計算和記憶體複雜度是O(n²)，也就是說，每當序列長度加倍，計算和記憶體需求就增加四倍（而三倍則意味著增加九倍）。
透過FlashAttention技術（一部分今天的研究人員開發的），我們幫助降低記憶體需求至次平方複雜度，透過避免將整個注意力矩陣物化到記憶體中，也就是說，以上問題部分解決。然而，FlashAttention並未處理計算複雜度，它仍然是平方複雜度（特別是O(n²*d)，其中n是序列長度，d是模型維度，即每個向量嵌入的數字量）。原因是因為我們不壓縮記憶體，每個詞必須關注每一個先前的詞，而在Transformer中主要的操作是成對的乘法運算，因此計算在序列長度方面仍然是平方複雜度。
但現在，一組研究人員找到了將計算複雜度擴充套件至次平方領域的方法。
如何線性化注意力 正如我們之前提到的，Transformer依賴於一種稱為注意力機制的數學操作，它的複雜度在序列長度方面是平方的。簡單來說，較長的序列會急劇增加計算需求。然而，其他注意力機制（如線性注意力）則具有線性複雜度，雖然它們在紙面上表現不佳（即線性注意力模型表現較差）。
但如果我們透過訓練線性注意力層去模仿它們更高計算密度的對等層，會發生什麼呢？
因此，LoLCaTs的目標是建立線性化的LLM，它們在保留其平方複雜度對等體效能的同時具備成本效率。
為此，他們將問題分為三步。
步驟1：替換層 線性化模型的第一步是插入一組線性注意力層，並訓練它們模仿標準注意力層的輸出。如下所示，透過兩層輸出之間的均方誤差（MSE），我們可以訓練新層表現得像原始層。
但此時，雖然新層個體上在模仿對等層，已建立的模型表現異常。為了解決這個問題，我們需要重新訓練模型執行標準的下一個詞預測任務。
步驟2：LoRA微調 當然，重新訓練模型以提高效率似乎是一個難以證明合理性的資源使用方式。幸運的是，完全模型微調並非必要，因為有了LoRA介面卡。我在這個通訊的多個期刊中提到過它們，但概念是：對任何給定任務，模型本質上是低秩的。簡而言之，只有一小部分權重對任何給定預測是重要的。因此，我們可以新增一組小權重，稱為介面卡，然後將其新增到每層中，同時保持原始模型權重‘未觸動’。某種程度上，這些介面卡會根據它們所訓練的資料調整模型的行為。
重要的是，這些介面卡比實際模型小得多，因此它們可以快速且低成本地訓練。
經過LoRA訓練後，我們得到一個在下一個詞預測上表現與原模型相似的模型，儘管其具有線性注意力層。
該模型並非完全線性化。實際上，他們將標準注意力（平方）與線性注意力結合起來，因此如果一個序列有D=N+M個詞，最後的M個詞使用標準注意力計算，而前面的N個（數量更多）則使用線性注意力。
這在標準注意力的表現力和線性注意力的計算效率之間找到了一個良好的平衡。由於語言通常具有鄰近性偏向，這意味著最近的詞通常比過去的詞對預測下一個詞更重要，我們使用softmax注意力（標準注意力）處理更近的詞，並使用線性注意力處理其餘詞。
序列中最後的M個詞采用傳統計算，其他則採用線性形式。
而在此之上，研究人員新增了最後一步。
步驟3：分層最佳化 在注意力遷移過程中，我們將模型中的所有層一起更新。然而，研究人員觀察到，這意味著模型的最後幾層具有更大的均方誤差（MSE），對於像Llama 3.1 405B這樣更深層的模型（比小型版本擁有更多層）尤其如此。因此，他們分批訓練層，以確保不會發生這種情況。
那麼，最終結果是什麼？
最先進…且便宜
LoLCaTs透過在4000萬個tokens上最佳化0.2%的模型引數，大幅提升了線性化大型LLM的效率，關閉了與完全softmax（標準）Transformer效能差距的80%，儘管所用訓練tokens量級少了數百倍，比標準微調高達35,500倍的訓練效率。
更令人印象深刻的是，在所有三個Llama尺寸上，LoLCaT模型的表現類似於原始模型，但享有線性注意力層帶來的顯著計算效率，而不是完全的平方注意力。
這些成果僅花費了Llama 3.1 8B單GPU上的幾個訓練小時（成本低於20美元）。
當然，我們必須承認這種方法仍需一個預訓練的平方Transformer來執行這一層蒸餾，但像LoLCaTs這樣的基於推理的最佳化將成為任何希望採用生成式AI的企業的重要組成部分，在獲得優異效能的同時節省大量計算成本。
總結： LoLCaTs技術可以讓大型語言模型（LLM）的訓練成本大大降低。利用這個技術，80 億引數的模型訓練成本只需要不到 20 美元，同時還保持了原本的高效能。和傳統的 Transformer 方法相比，LoLCaTs 大幅提升了效率，讓模型訓練變得又快又省錢。這一突破讓更多企業可以低成本獲得強大的 AI 能力，為 AI 的普及鋪平了道路。