僅需幾十行 PyTorch 代碼即可大幅提高 GPU 利用率，在英偉達(dá) A100 上的 GPU 利用率高達(dá) 70%。這一能力由一款名為 LaCT 的新模型架構(gòu)實(shí)現(xiàn)，它由北京大學(xué)本科校友、美國麻省理工學(xué)院博士生張?zhí)爝h(yuǎn)和所在團(tuán)隊打造。

圖丨張?zhí)爝h(yuǎn)（來源：https://tianyuanzhang.com/）

研究團(tuán)隊通過在不同模態(tài)的任務(wù)中采用范圍從 2000 到 100 萬 tokens 的大塊更新打造了這種新模型架構(gòu)。該架構(gòu)集成了用于捕捉長上下文的大塊測試時訓(xùn)練，并集成了用于針對局部結(jié)構(gòu)進(jìn)行建模的窗口注意力機(jī)制。

由于每個大塊內(nèi)的 tokens 都被視為一個無序集，因此研究團(tuán)隊將窗口注意力集成到 LaCT 中，以便捕獲塊內(nèi)的局部依賴關(guān)系，這讓 LaCT 能夠顯著提高并行性。

這能夠擴(kuò)展非線性快速權(quán)重，從而能夠提高存儲容量。而且，這種簡單的實(shí)現(xiàn)方式能夠輕松地集成更有效的測試時優(yōu)化器（比如 Muon）。

此外，LaCT 的大塊設(shè)計也能夠很自然地用于針對各種 N 維數(shù)據(jù)進(jìn)行建模，原因在于它可以將塊大小與數(shù)據(jù)的內(nèi)部結(jié)構(gòu)對齊，例如將圖像或連續(xù)視頻幀內(nèi)的 tokens 分組為一個塊。

為了驗(yàn)證本次方法的有效性，研究團(tuán)隊在不同數(shù)據(jù)模態(tài)和任務(wù)中，包括從圖像集、語言模型和自回歸視頻擴(kuò)散模型中進(jìn)行了新視圖合成。

實(shí)驗(yàn)結(jié)果顯示：研究團(tuán)隊的模型能夠處理多達(dá) 128 張分辨率為 960×536 的輸入圖像，形成最多 100 萬 tokens 的序列，并且在此輸入規(guī)模下，在渲染質(zhì)量方面優(yōu)于 3D 高斯?jié)姙R（3D Gaussian Splatting）技術(shù)。

盡管語言數(shù)據(jù)本身并不顯式包含塊狀結(jié)構(gòu)，但是與 DeltaNet 等 SOTA 方法相比，研究團(tuán)隊的模型實(shí)現(xiàn)了大致相當(dāng)?shù)男阅鼙憩F(xiàn)。

研究團(tuán)隊還通過將 LaCT 與滑動窗口注意力相結(jié)合，將一個 140 億參數(shù)的雙向視頻擴(kuò)散 Transformer 適配為自回歸模型。這種適配后的模型可以生成包含多達(dá) 56000 個視覺 tokens 的連貫視頻。

與此同時，在最長的序列實(shí)驗(yàn)中，他們進(jìn)行了超過 100 萬個上下文長度的新視圖合成。

目前，研究團(tuán)隊已經(jīng)開源了代碼和權(quán)重，預(yù)計 LaCT 將能推動人們對于更高效長上下文建模架構(gòu)的探索（https://tianyuanzhang.com/projects/ttt-done-right/）。

吞吐量開銷小至 1%-3%

當(dāng)前，處理長上下文的需求正在迅速增長。雖然 softmax 注意力已成為建模各類數(shù)據(jù)的解決方案，但其計算成本隨序列長度呈二次方增長，這推動了人們對更高效長上下文建模的廣泛研究。

最近，測試時間訓(xùn)練（TTT，Test-Time Training）已成為一種富有前景的高效二次序列建模方法。測試時間訓(xùn)練能將循環(huán)神經(jīng)網(wǎng)絡(luò)中的循環(huán)狀態(tài)概念擴(kuò)展到一個小型的、在線自適應(yīng)的子網(wǎng)絡(luò)。這個子網(wǎng)絡(luò)的參數(shù)也被稱為快速權(quán)重，它們通過自監(jiān)督目標(biāo)在線快速適配，以便記憶上下文中的信息。

近期，多個團(tuán)隊均探索了快速權(quán)重網(wǎng)絡(luò)的各種在線目標(biāo)、優(yōu)化器和架構(gòu)。盡管如此，已有的測試時訓(xùn)練方法仍然難以有效擴(kuò)展到長上下文場景，根本原因在于測試時訓(xùn)練層的硬件利用率極低，在當(dāng)前 GPU 上硬件利用率通常低于峰值算力的 5%。

這種低效性是由于使用了小批量規(guī)模，即每隔一個 token 或每 16 個到 64 個 tokens 更新一次快速權(quán)重，之所以這樣做是因?yàn)閭鹘y(tǒng)觀點(diǎn)認(rèn)為這種方式對于上下文學(xué)習(xí)更加有效。

但是，這種小批量處理方式會導(dǎo)致并行效率低下以及計算密度不足，尤其在使用大型非線性快速權(quán)重時，會給硬件高效實(shí)現(xiàn)帶來重大挑戰(zhàn)，以至于實(shí)際算力利用率難以突破 10% 的有效閾值。基于此，本次研究團(tuán)隊采用相反的策略并引入了 LaCT。

如下圖所示，LaCT 塊由三種類型的層組成：窗口注意力層、大塊測試時訓(xùn)練層和前饋層。

（來源：arXiv）

每一層都配備了殘差連接，這一設(shè)計也遵循了 Transformer 架構(gòu)中的標(biāo)準(zhǔn)做法。窗口注意力層通過執(zhí)行局部自注意力，來捕捉局部依賴關(guān)系。而在測試時訓(xùn)練層，研究團(tuán)隊則將序列分割成了大塊。

研究團(tuán)隊表示，歷史上下文通過“更新”操作逐漸被壓縮到快速權(quán)重中，最新的權(quán)重被“應(yīng)用”到當(dāng)前的查詢向量（Q）上，以便計算其對應(yīng)的輸出。前饋層則執(zhí)行與 Transformer 中類似的通道混合操作。

由于測試時訓(xùn)練的“更新”操作和“應(yīng)用”操作是解耦的，因此可以自適應(yīng)地設(shè)置塊大小，并以不同的順序應(yīng)用這些操作，進(jìn)而能夠模擬不同類型的數(shù)據(jù)依賴關(guān)系。

當(dāng)分塊大小等于完整序列長度時，會先執(zhí)行“應(yīng)用”操作再執(zhí)行“更新”操作，這在概念上與全注意力機(jī)制相似。通過交替使用“更新”操作和“應(yīng)用”操作，能夠形成分塊因果掩碼，其中分塊大小與塊大小互相對應(yīng)。在兩個操作之間切換順序會導(dǎo)致掩碼發(fā)生偏移，偏移掩碼不會在塊內(nèi)泄露未來信息，這在語言建模中構(gòu)建完整因果掩碼時非常重要。

（來源：arXiv）

大塊測試時訓(xùn)練層會將數(shù)據(jù)視為集合序列，因?yàn)槠淇焖俚臋?quán)重更新會忽略每個塊內(nèi)的 tokens 順序和空間局部性。然而，許多數(shù)據(jù)模態(tài)比如視頻、圖像集合或文本，并不完全符合這種基于集合的視角。對于這些模態(tài)而言，塊內(nèi)結(jié)構(gòu)和局部性對于捕獲整體數(shù)據(jù)結(jié)構(gòu)至關(guān)重要。

因此，研究團(tuán)隊將局部窗口注意力層與測試時訓(xùn)練層集成在一起，以便處理塊內(nèi)的數(shù)據(jù)結(jié)構(gòu)。此外，窗口注意力機(jī)制能有效捕捉數(shù)據(jù)中的局部特征。對于測試時訓(xùn)練層來說，這讓它能夠?qū)⑵涔潭ù笮〉目焖贆?quán)重容量集中用于建模非局部依賴關(guān)系。

總的來說，LaCT 是一種混合架構(gòu)，它采用二次計算注意力機(jī)制來處理局部結(jié)構(gòu)，針對非局部上下文采用線性計算的測試時訓(xùn)練機(jī)制。上下文并行（CP，Context Parallelism）沿著上下文長度維度針對序列進(jìn)行分區(qū)，并將分片分布在多個設(shè)備上來進(jìn)行并行計算。

前饋層和窗口注意力均屬于局部操作算子，因此天然地支持上下文并行。對于測試時訓(xùn)練層，小塊難以支持上下文并行，因此更傾向于使用張量并行。

研究團(tuán)隊的大塊測試時訓(xùn)練層通過在塊內(nèi)分片 tokens 來實(shí)現(xiàn)上下文并行。在訓(xùn)練新視圖合成時，他們采用了這種并行方法，并觀察到 1% 至 3% 的極小吞吐量開銷。與此同時，LaCT 架構(gòu)可以與數(shù)據(jù)并行、流水線并行和張量并行等其他并行策略兼容。

實(shí)驗(yàn)涵蓋：新視圖合成、語言建模和自回歸視頻生成

如前所述，研究團(tuán)隊開展了關(guān)于新視圖合成、語言建模和自回歸視頻生成的實(shí)驗(yàn)。在與線性成本基線方法的對比實(shí)驗(yàn)中，研究團(tuán)隊為其增加了相同的窗口注意力模塊，以便確保能夠進(jìn)行公平的比較。

表丨對每個實(shí)驗(yàn)中關(guān)鍵因素的總結(jié)（來源：arXiv）

在新視圖合成上，研究團(tuán)隊在場景級和物體級數(shù)據(jù)集上對本次方法進(jìn)行評估。他們使用 Objaverse 數(shù)據(jù)集進(jìn)行物體級訓(xùn)練，并遵循 LVSM 和 GS - LRM 的設(shè)置。

訓(xùn)練完成之后，研究團(tuán)隊在 Google Scanned Objects（GSO）數(shù)據(jù)集上進(jìn)行評估，該數(shù)據(jù)集的分辨率分別為 256×256 和 512×512。每次評估涉及 4 到 48 個輸入視圖，且每個物體有 8 個新視圖。

對于場景級評估，研究團(tuán)隊采用挑戰(zhàn)性較高的 DL3DV 場景數(shù)據(jù)集，其中包含超過 11000 個訓(xùn)練場景和 140 個測試場景，每個場景大約有 300 個視圖，評估的分辨率為 960 × 536。

對于物體級評估，研究團(tuán)隊使用了如下兩個基線模型：全注意力模型和寄存器注意力模型。

全注意力基線模型將測試時訓(xùn)練層替換為逐塊因果注意力層，實(shí)現(xiàn)了輸入 tokens 之間的雙向交互和來自新視圖的交叉注意力。

寄存器注意力模型將輸入 tokens 壓縮到 4096 個寄存器中，并通過與這些寄存器的交叉注意力解碼新視圖。

在場景級評估中，研究團(tuán)隊與 LongLRM 進(jìn)行對比，LongLRM 是一種結(jié)合了 Mamba 和全注意力機(jī)制的模型，可用于 3D 高斯濺射（3D Gaussian splat）預(yù)測。此外，他們還與純基于優(yōu)化的 3D 高斯濺射方法進(jìn)行了對比。

表丨對所有模型計算復(fù)雜性的總結(jié)（來源：arXiv）

在性能評估上，研究團(tuán)隊采用每 tokens 損失度量來評估模型有效使用完整上下文的能力。出現(xiàn)單調(diào)遞減的損失表示上下文利用成功，而處于平穩(wěn)狀態(tài)則表示上下文使用有限。

另據(jù)悉，他們從原始 LaCT 塊中移除了窗口注意力層，將滑動窗口注意力（SWA，sliding window-attention）層直接集成到大塊測試時訓(xùn)練層中，并將模型與全注意力模型、門控線性注意力（GLA，Gated Linear Attention）和 DeltaNet 進(jìn)行了比較。

為了確保公平性，他們?yōu)?GLA 和 DeltaNet 都增強(qiáng)了相同的滑動窗口注意力層，并采用 100 萬的 RoPE 庫進(jìn)行 32K tokens 上下文的培訓(xùn)。

表丨對所有方法機(jī)制和訓(xùn)練吞吐量的總結(jié)（來源：arXiv）

為了比較塊遞歸和逐 tokens 遞歸，在條件受控的實(shí)驗(yàn)中，研究團(tuán)隊的線性大塊遞歸策略在相同狀態(tài)大小下優(yōu)于線性逐 tokens 遞歸策略在。

由于語言本身并不天然存在塊狀結(jié)構(gòu)，研究團(tuán)隊提出的線性大塊遞歸變體在初始階段性能不如 GLA 和 DeltaNet 等逐 token 方法。然而，當(dāng)將其與大規(guī)模非線性狀態(tài)以及 Muon 優(yōu)化器相結(jié)合時，該變體的表現(xiàn)將超越這些逐 token 方法。

總的來說，本次成果凸顯了大塊測試時訓(xùn)練在計算效率和性能上的優(yōu)勢，為更高效且可擴(kuò)展的長上下文序列建模鋪平了道路。

通過消除對于低級硬件特定實(shí)現(xiàn)的依賴，LaCT 使人們能夠更廣泛地探索架構(gòu)設(shè)計空間。未來，研究團(tuán)隊希望這項工作能夠啟發(fā)并加速長上下文建模和測試時訓(xùn)練領(lǐng)域的新研究。

參考資料：

https://arxiv.org/abs/2505.23884

運(yùn)營/排版：何晨龍