ICML 2025 | 又快又強的獅虎獸！僅20M詞元微調，Transformer秒切線性RNN

近日上海 AI Lab 聯合華南理工大學、香港科技大學（廣州）、南京大學和香港中文大學發佈了一項研究成果：Liger（獅虎獸），即 Linearizing LLMs to gated recurrent structures，這是一項能夠高效地將預訓練大型語言模型架構轉換為帶有門控循環結構的線性模型的線性化技術。

目前這項研究工作已被 ICML 2025 接收，代碼與模型已全部開源。

論文標題：

Liger: Linearizing Large Language Models to Gated Recurrent Structures

論文連結：

https://arxiv.org/abs/2503.01496

代碼連結：

https://github.com/OpenSparseLLMs/Linearization

模型連結：

https://huggingface.co/collections/linear-moe-hub/liger-67d904bffd7f9b77ade7747d

大型語言模型（Llama、Mistral等）在各種序列建模任務上取得了優異的表現，特別是基於 Transformer 架構的大型語言模型已被廣泛驗證其序列建模任務的有效性，然而這種架構也面臨著其固有的缺陷：

1. 注意力機制關於序列長度呈現二次方的計算複雜度，每次為了生成下一個詞都需要回顧歷史序列數據進行注意力計算，導致 Transformer 架構模型在長序列情境下效率低下；

2. KV-Cache 機制需要保存歷史序列數據用於後續計算，導致顯示記憶體壓力隨著序列長度增長而升高。

傳統 Transformer 架構模型的效率瓶頸日益凸顯，如何在不犧牲性能的前提下實現高效推論，成為學術界與工業界共同關注的焦點。

在這一背景下，基於線性序列建模方法的模型架構嶄露頭角，因為線性循環模型在架構上具有明顯優勢：

1. 線性注意力機制關於序列長度呈現線性計算複雜度，每次生成下一個詞僅需要訪問固定大小的記憶體/狀態，計算效率高；

2. 無需 KV-Cache，推論生成階段的顯示記憶體佔用恆定，無論生成序列有多長，顯示記憶體佔用一直維持恆定不變。

得益於其高效性質能夠完美解決 Transformer 架構的固有缺陷，線性循環架構模型展現出作為 LLM 基礎架構的優勢。

然而驗證一種新興的模型架構的有效性並不容易，這是因為訓練一個參數量非常龐大的模型動輒需要上千乃至上萬億的高品質數據量，並且算力需求極其高昂，需要在大規模顯示卡叢集上對隨機初始化的大模型進行從頭預訓練。

因此從頭訓練此類線性循環模型成本高昂，並且通常難以匹配現有 Transformer LLM 的性能，使得絕大多數研究人員望而卻步，難以實際投入如此高成本訓練出一個效果可能並不理想的線性 LLM。

既然我們已經有了預訓練好的 Transformer 大型模型（Llama、Mistral 等），將現有模型架構調整為線性循環模型架構並在此基礎上進一步訓練或許是成本更低的方案，我們稱之為模型架構線性化。

然而現有的線性模型為了擬合 Transformer 架構的 Softmax Attention 效果，需要在原有線性注意力上增加各種模組，包括 Feature Mapping、門控機制等，能夠一定程度上提升原有線性模型的表現。

然而現有線性化方法仍未探究如何更好地將 Transformer 線性化為門控循環結構的線性模型，並且在線性化情境下這些額外模組需要初始化進行訓練，這增加了架構複雜性和差異性，額外增加了線性化成本。

在此背景下，Liger 應運而生，這是一個極其高效、簡潔且通用的線性化技術，僅需要極少量的微調成本即可將預訓練的 Transformer LLM 線性化為門控循環結構，成功恢復原模型 93% 以上的性能，同時實現高效線性計算複雜度的序列建模。

方法描述

Liger 的核心目標是透過簡潔且低成本的訓練進行模型結構轉換，將預訓練 LLM 的權重直接遷移至門控循環架構中，避免從頭預訓練的高昂成本。

化繁為簡：巧用模型參數冗餘

基於門控機制的線性循環模型需要獨立設計門控模組，導致引入了額外的可訓練參數且增加了模型複雜度。Liger 巧妙運用 LLM 中固有的參數冗餘特性，透過將鍵矩陣（Key Matrix）進行變換用於門控機制建構：

具體透過池化（Pooling）的無參操作，直接從鍵投影矩陣中提取門控資訊，從而無需新增可訓練參數。由於線性循環模型去除了 Softmax 操作，這可能導致 QK 的乘積未進行歸一化數值膨脹而無法擬合原始輸出的分佈，因此線性循環模型通常需要引入可學習的 Feature Mapping 函數以擬合 Softmax Attention。

具體實現中，我們將 Feature Mapping 簡化為 Softmax 函數分別作用在 Q 和 K 上，提供 QK 乘積數值歸一化的穩定性，確保與原始 LLM 注意力機制的相容性，同時無需引入任何可訓練參數，透過完全複用 LLM 權重減少了模型架構複雜度和差異性，而無需多階段訓練，從而進一步減少線性化成本並提高模型表現。

Liger 方法能夠相容各種帶有門控機制的線性循環模型架構，十分靈活高效。

輕量微調：LoRA 助力線性結構適應

在進行模型結構轉換後，Liger 採用低秩適應（LoRA）技術對模型進行微調以適應線性循環模型架構。

Liger 線性化僅僅是改變了注意力層 QKV 的運算順序，透過右乘核技巧達到線性高效計算，因此僅需要採用 LoRA 對注意力層的 QKV 投影矩陣進行低秩分解訓練，而無需對整個模型進行全參數微調，訓練目標採用自迴歸的下一詞元預測（Next Token Prediction），損失函數為交叉熵損失（Cross-Entropy Loss）：

LoRA 輕量微調使得 Liger 線性化過程中能夠充分保留 LLM 預訓練知識，降低線性化成本並快速恢復大部分性能。

混合機制：Liger Attention

為了進一步提升線性化表現，本文提出了 Liger Attention 混合注意力機制，透過結合視窗注意力機制（Sliding Window Attention，SWA）與門控循環建模（Gated Recurrent Modeling，GRM）達成層內線性序列建模方法與注意力機制的混合，同時保留線性計算複雜度的高效性。

Liger 也可用於層間混合架構的高效線性化，每 7 層門控循環模組後插入 1 層標準注意力模組，既能捕捉長程依賴，又透過局部注意力增強關鍵資訊的處理，進一步提升了模型適應性。

實驗分析

作者透過實驗對比了 Liger 與現有的各種模型架構線性化方法，結果表明 Liger 在訓練成本都小於其他方法的前提下，僅需要 20M 訓練詞元的成本就能夠恢復預訓練 Transformer 大型模型 93% 以上的性能，在各種語言建模任務中均接近或超過現有的 SOTA 線性化方法，非常接近 Llama、Mistral 等 Transformer 架構的 LLM 表現。

得益於線性模型的架構優勢，Liger 的推論時間隨序列長度線性增長，16K 的序列長度推論較 Flash Attention 加速 2 倍，處理 32K 長度序列時，Liger 的 GPU 記憶體佔用恆定在 16.37GB，而原始基於 Flash Attention 的 Llama-3 因記憶體不足（OOM）無法完成推論。