大型語言模型（LLMs）在近年來取得了令人矚目的進展，但它們的高計算需求和龐大的參數(shù)量使得普通用戶難以輕松獲取和應(yīng)用這些模型。在這個背景下，來自MediaTek Research的研究團隊包括Yen-Chen Wu、Feng-Ting Liao、Meng-Hsi Chen、Pei-Chen Ho、Farhang Nabiei和Da-shan Shiu最近發(fā)表了一項突破性研究——Latent Flow Transformer（LFT）。這項研究于2025年5月20日發(fā)布在arXiv預(yù)印本平臺（arXiv:2505.14513v1），旨在解決大型語言模型結(jié)構(gòu)效率問題，為模型壓縮提供全新思路。

想象一下，如果你正在搭建一座多層大樓，但發(fā)現(xiàn)有些樓層其實并不必要，去掉它們后大樓依然穩(wěn)固且功能齊全，你會怎么做？這正是LFT嘗試解決的問題。當(dāng)代的大型語言模型就像一座由數(shù)十上百層組成的高樓，每一層（transformer層）都會對信息進行處理和傳遞。研究人員發(fā)現(xiàn)，這些層中有許多可能是冗余的，但傳統(tǒng)的壓縮方法如剪枝和知識蒸餾并未能提供足夠規(guī)模的減少。

與此同時，在圖像生成領(lǐng)域，擴散模型和流匹配（Flow Matching）模型已經(jīng)證明了其令人印象深刻的效率，最新的方法甚至能夠在僅一次采樣過程中生成高質(zhì)量圖像。受此啟發(fā)，研究團隊提出了Latent Flow Transformer，這是一種創(chuàng)新的transformer變體，它通過流匹配（Flow Matching）來學(xué)習(xí)潛在空間的傳輸映射，實現(xiàn)了將多個transformer層壓縮為單一層的目標(biāo)。

簡單來說，LFT就像是在語言模型的多層結(jié)構(gòu)中建立了一條"高速公路"，讓信息可以直接從起點跳躍到終點，而不必經(jīng)過中間的每一個"收費站"（層）。這種方法不僅保持了與原始架構(gòu)的兼容性，還大大減少了計算需求和參數(shù)量。

為了確定哪些層適合被替換，研究團隊引入了一個名為"重耦合比率"（Recoupling Ratio）的指標(biāo)，它能準(zhǔn)確預(yù)測流匹配的質(zhì)量。有趣的是，他們的分析發(fā)現(xiàn)中間層特別適合壓縮，這支持了一個直覺——早期層和后期層在功能上存在質(zhì)的差異。

在Pythia-410M模型上的實驗中，研究團隊證明了通過流匹配訓(xùn)練的LFT可以將24層中的6層壓縮，并且性能優(yōu)于直接跳過2層的方法（LM logits的KL散度為0.407 vs. 0.529），這證明了該設(shè)計的可行性。更令人興奮的是，當(dāng)使用他們新提出的Flow Walking（FW）算法訓(xùn)練時，LFT進一步將12層蒸餾為一層，同時將KL散度降低到0.736，超過了跳過3層的效果（0.932），這顯著縮小了自回歸和基于流的生成范式之間的差距。

讓我們深入了解這項研究的細節(jié)，看看MediaTek Research團隊是如何實現(xiàn)這一令人印象深刻的突破的。

一、Transformer層的連續(xù)時間視角

要理解LFT的工作原理，我們首先需要轉(zhuǎn)變看待Transformer層的方式。想象一下，如果把神經(jīng)網(wǎng)絡(luò)的每一層比作一個加工站，那么傳統(tǒng)的觀點是數(shù)據(jù)包裹必須在每個站點都停留并接受處理。但研究者提出了一個更流暢的視角：將這些離散的處理站點看作一條連續(xù)的流水線。

在這種連續(xù)時間的視角下，一個有L個離散層的神經(jīng)網(wǎng)絡(luò)，特別是帶有殘差連接的網(wǎng)絡(luò)，可以被視為一個將隱藏狀態(tài)從初始狀態(tài)逐步改進到最終狀態(tài)的過程。用數(shù)學(xué)語言來說，這可以表示為一個常微分方程（ODE）：

``` dht/dt = uθ(ht, t) ```

這里，`θ`是參數(shù)化這個方程的參數(shù)，當(dāng)時間`t`與層數(shù)`l/L`匹配時，`ht`就對應(yīng)于第`l`層的隱藏狀態(tài)。想象成一條河流，水（數(shù)據(jù)）從起點流向終點，在每一個時刻都有一個微小的變化，這些微小變化的累積形成了整體的轉(zhuǎn)換效果。

然而，學(xué)習(xí)這種參數(shù)θ涉及到昂貴的模擬計算：

``` ht2 = ht1 + ∫(t1到t2) uθ(ht, t) dt ```

這就像是要預(yù)測河水從一點流到另一點需要解決復(fù)雜的水流方程，計算成本很高。

二、無需模擬的流匹配訓(xùn)練

這就是流匹配（Flow Matching）方法發(fā)揮作用的地方。流匹配提供了一種無需模擬的解決方案，讓我們可以直接學(xué)習(xí)從起點到終點的"流動"。

想象你有一個小球從位置x0要移動到位置x1，流匹配就是學(xué)習(xí)這個球在單位時間內(nèi)通過預(yù)定義路徑移動所需的速度場。它通過最小化以下?lián)p失函數(shù)來學(xué)習(xí)：

``` LFlowMatching = Et[||uθ(xt, t) - vt||?] ```

其中，xt和vt分別是粒子在時間t的位置和速度。

在實踐中，一條直線、勻速軌跡是流匹配的常見選擇。在這種情況下，xt和vt可以簡單地表示為： ``` xt = (1-t)·x0 + t·x1 vt = x1 - x0 ```

這就像是小球沿著直線勻速移動，從起點到終點。

在離散時間推理過程中，數(shù)據(jù)點沿著流動軌跡以離散步驟移動。當(dāng)從時間點t移動到t+d時，一個步驟是：

``` xt+d = xt + d·uθ(xt, t) ```

為了提高流軌跡的穩(wěn)定性和準(zhǔn)確性，可以使用中點估計：

``` xt+d = xt + d·uθ(xt + d/2·uθ(xt, t), t + d/2) ```

這就像是在預(yù)測未來位置時，不僅考慮當(dāng)前位置，還考慮中間過程的狀態(tài)，從而獲得更精確的估計。

三、配對數(shù)據(jù)流匹配的挑戰(zhàn)

在配對樣本的流匹配中，至關(guān)重要的是保持源分布和目標(biāo)分布之間的確定性對應(yīng)關(guān)系。這種設(shè)置下，當(dāng)插值軌跡相交時，會對標(biāo)準(zhǔn)流匹配方法帶來挑戰(zhàn)，因為這些方法在交叉點附近會平均沖突的速度信號，產(chǎn)生有偏估計，無法準(zhǔn)確地將源點傳輸?shù)狡渑鋵Φ哪繕?biāo)。

想象兩條河流交匯，如果只看交匯點的水流方向，你無法確定哪股水來自哪條上游河道。為了克服這個問題，有研究引入了輔助加速度場來調(diào)節(jié)每條軌跡上的變化率，從而使對真實配對映射的對齊更加清晰。另一種方法是學(xué)習(xí)潛在空間投影，在流估計之前解開交叉對，有效地防止軌跡交叉。

四、Latent Flow Transformer架構(gòu)

有了這些背景知識，我們現(xiàn)在可以深入了解Latent Flow Transformer（LFT）的具體工作原理。LFT是一種新型的、為語言建模設(shè)計的Transformer變體，它試圖利用流相關(guān)概念的優(yōu)勢，正如這些概念在圖像生成中所展示的那樣。

LFT通過將教師模型中的連續(xù)Transformer層塊替換為單個學(xué)習(xí)的傳輸運算符（稱為潛在流層）來減少模型大小。該運算符使用流概念進行訓(xùn)練，以準(zhǔn)確地將該塊輸入的潛在表示映射到其對應(yīng)的輸出。

對于LFT的速度場估計網(wǎng)絡(luò)，研究團隊遵循了前人的方法，通過附加的縮放和移位運算符增強了一個教師Transformer層，以及預(yù)測這些因子的MLP網(wǎng)絡(luò)。他們通過從這個增強網(wǎng)絡(luò)的輸出中減去輸入潛在表示來獲得速度估計。

五、重耦合比率：預(yù)測流匹配質(zhì)量

選擇最佳的層塊進行替換對LFT的性能至關(guān)重要。通過流匹配學(xué)習(xí)的可行性受到流路徑交叉的限制，而在LFT中這一挑戰(zhàn)更為嚴重，因為保留原始輸入-輸出配對排除了使用重流方法的可能。

為了指導(dǎo)層選擇，研究團隊引入了一個稱為"重耦合比率"的指標(biāo)。給定來自層m到n的樣本潛在對，這個估計器量化了它們原始配對與由最優(yōu)傳輸（OT）指定的配對之間的偏差。最優(yōu)傳輸確定了層m和層n的潛在表示之間的最小成本映射。

重耦合比率R被定義為與最優(yōu)傳輸矩陣M不一致的配對關(guān)系的百分比：

``` R := 1 - E[Tr(M)/OM] ```

這里OM是矩陣M的階。因為重耦合比率量化了與原始配對的不一致程度，所以較低的R表示更好的一致性，從而預(yù)測較少的流交叉問題，學(xué)習(xí)LFT的可行性更高。

六、學(xué)習(xí)速度場

要學(xué)習(xí)速度場，研究團隊對一系列x0、x1對應(yīng)用了標(biāo)準(zhǔn)流匹配算法。對于替換從層m到層n（包括n）的所有層的LFT，他們將給定令牌的層m的輸入潛在表示作為x0，將層n的輸出潛在表示作為相應(yīng)的x1。

學(xué)習(xí)過程如下： 1. 從訓(xùn)練數(shù)據(jù)集中抽取一個樣本d 2. 獲取(x0, x1)對，分別是LLM教師模型在層m和層n上處理數(shù)據(jù)d的潛在表示 3. 隨機抽取時間點t 4. 計算插值位置xt = (1-t)x0 + tx1 5. 更新參數(shù)θ以最小化uθ(xt, t)和真實速度(x1-x0)之間的平方誤差

七、LFT的展開結(jié)構(gòu)與Transformer的相似性

在推理時，將潛在流層展開成一組固定的時間點t0=0 < t1 < t2... < 1，會將潛在演化過程硬化為靜態(tài)處理圖，這有利于可視化數(shù)據(jù)流和優(yōu)化硬件實現(xiàn)。對于單步流匹配結(jié)合簡單重建近似規(guī)則的特定情況，潛在流層相當(dāng)于單個標(biāo)準(zhǔn)Transformer層。如果使用多個步驟，潛在流層相當(dāng)于具有跨層注意力的Transformer層堆棧。

這種展開的LFT與標(biāo)準(zhǔn)Transformer之間的結(jié)構(gòu)相似性具有重要的實際意義，使研究人員和實踐者能夠利用為基于Transformer的LLM開發(fā)的廣泛生態(tài)系統(tǒng)和高度優(yōu)化的基礎(chǔ)設(shè)施。

八、Flow Walking：提高配對數(shù)據(jù)流匹配

為了解決交叉軌跡的挑戰(zhàn)，同時支持將潛在流層展開回Transformer結(jié)構(gòu)，研究團隊引入了Flow Walking (FW)算法來訓(xùn)練和推理LFT的潛在流層。

FW使用數(shù)值積分來近似從x0到x1的路徑，定義一個步驟為sθ(xt, t, t')=xt+d，其中d=t'-t，xt+d可以由前面介紹的方程確定。關(guān)鍵直覺是通過在交叉點附近輕微分離軌跡來學(xué)習(xí)非交叉軌跡。學(xué)習(xí)速度場的目標(biāo)函數(shù)定義為：

``` LFlowWalking(k) = Et1,...,tk-1[||x0 + ∑Δθ,ti - x1||?] ```

其中Δθ,ti = sθ(xti-1, ti-1, ti) - xti-1，t0=0，tk=1，ti從[0,1]中隨機采樣。

研究團隊在實驗中發(fā)現(xiàn)，選擇k=3并隨機采樣t可以提供一個高效且泛化的訓(xùn)練損失。此外，他們還探索了添加直線性正則化器的可能性，通過標(biāo)準(zhǔn)流匹配來實現(xiàn)：

``` L = LFlowWalking + αLFlowMatching ```

這產(chǎn)生了一個非交叉、直線插值的配對數(shù)據(jù)，顯示出在配對數(shù)據(jù)應(yīng)用中的潛力。

九、實驗結(jié)果與分析

研究團隊在Pythia-410M模型上進行了廣泛的實驗來驗證所提出的框架。他們的主要發(fā)現(xiàn)包括：

1. **層選擇至關(guān)重要**：通過重耦合比率選擇的層構(gòu)建的LFT始終優(yōu)于使用任意層選擇的LFT。這驗證了該指標(biāo)在預(yù)測流匹配質(zhì)量方面的有效性。

2. **重耦合比率分析**：分析顯示，Transformer模型的中間層特別適合壓縮，支持早期層和后期層在質(zhì)量上有所不同的假設(shè)。特別是通過計算層0到層6之間的最優(yōu)傳輸矩陣，他們發(fā)現(xiàn)重耦合比率為0.53，表明可能會遇到大量與流交叉相關(guān)的問題。相比之下，層6到層18之間的最優(yōu)傳輸矩陣顯示重耦合比率為零，表明對于這批數(shù)據(jù)，現(xiàn)有的配對已經(jīng)是最優(yōu)的。

3. **LFT的蒸餾質(zhì)量**：在The Pile上訓(xùn)練時，標(biāo)準(zhǔn)流匹配（LFT-SFM）和Flow Walking（LFT-FW）都顯示出快速收斂，并且明顯優(yōu)于簡單的層跳過方法。當(dāng)替換Pythia-410m的層6-12（25%的參數(shù)）或?qū)?-18（50%的參數(shù)）時，LFT-FW在多個離散時間點的推理上始終優(yōu)于基線。值得注意的是，LFT-FW與k=1達到了與回歸模型相當(dāng)?shù)男阅?，并且?yōu)于具有k=8的LFT-SFM，表明Flow Walking的隱式速度估計更準(zhǔn)確地引導(dǎo)模型朝向其目標(biāo)隱藏狀態(tài)。

4. **離散時間點對推理性能的影響**：離散時間點的數(shù)量k是LFT在推理時的關(guān)鍵超參數(shù)。對于LFT-SFM，KL散度和NMSE隨著k的減少而減少，這意味著早期速度估計更正確地引導(dǎo)隱藏狀態(tài)朝向其目標(biāo)。對于LFT-FW，它在k=3時達到最佳性能，與方程中的三步積分相匹配。當(dāng)k接近1時，KL散度急劇上升，表明t=0時的隱式速度估計不準(zhǔn)確，必須通過多步校正進行改進。有趣的是，盡管從未使用超過三個積分步驟進行訓(xùn)練，LFT-FW對k=8表現(xiàn)出強大的泛化能力，僅顯示輕微的性能下降。

5. **性能對比**：在替換層6-18的設(shè)置下，LFT-FW（k=3）達到了0.736的KL散度，明顯優(yōu)于跳過三層的基線（0.932），驗證了該方法的有效性。這表明Flow Walking成功地克服了主要障礙，即在遠距離Transformer層之間對齊潛在傳輸，有效地橋接了自回歸和基于流的建模范式之間的差距。

十、討論與未來方向

研究團隊指出了幾個有趣的應(yīng)用和未來研究方向：

1. **推測解碼**：推測解碼指的是使用小型草稿模型來加速全尺寸模型的推理。由于直線流Transformer的大小可能比其教師模型小一個數(shù)量級，將直線流Transformer用作推測解碼中的草稿模型是很自然的?？紤]到研究中展示的KL距離結(jié)果，這樣的配置非常有前景。

2. **動態(tài)計算**：傳統(tǒng)Transformer具有固定的計算成本和固定的生成質(zhì)量。相比之下，直線流Transformer允許動態(tài)地改變步驟數(shù)，即在逐句或甚至逐令牌的基礎(chǔ)上動態(tài)改變計算投入。這讓我們可以考慮如何在使用的步驟數(shù)和相應(yīng)的質(zhì)量之間達到最佳平衡。

3. **應(yīng)用于其他模型**：最近，狀態(tài)空間模型和循環(huán)模型因其處理上下文的恒定計算復(fù)雜度而成為Transformer的替代候選。研究團隊推測，如果應(yīng)用他們的方法，可以創(chuàng)建直線流RWKV、直線流MAMBA、直線流xLSTM等。

4. **流解纏**：在當(dāng)前方法中，他們保留了源Transformer的輸入和輸出層不變。已經(jīng)確立的是，安排流匹配對以最小化流交叉顯著提高了性能。一個提議的優(yōu)化方向是優(yōu)化輸入和輸出層以最小化流交叉，條件是輸入嵌入和輸出logit。

5. **從頭訓(xùn)練流替換的Transformer**：一個開放問題是，一旦流替換并結(jié)構(gòu)簡化，是否可以從頭訓(xùn)練所得到的淺層Transformer，而無需首先預(yù)訓(xùn)練完整深度模型。擴散文獻中的先前觀察表明，從頭訓(xùn)練的一步生成器很少能匹配蒸餾的性能。語言建模中可能也存在類似的模式。

這項研究為提高Transformer模型的效率提供了新的視角和方法，有望在降低計算需求的同時保持性能，使大型語言模型更加普及和易于訪問。