AGI 新技術(shù)路線：下一代稀疏注意力機(jī)制 Monte Carlo Attention 開源

作者 | 超對(duì)稱技術(shù)

出品丨AI 科技大本營(yíng)（ID：rgznai100）

超對(duì)稱技術(shù)公司在新版基座模型 BigBang-Proton 使用的 Monte Carlo 注意力，在二進(jìn)制塊編碼（Binary Patch Encoding）技術(shù)上，用巧妙的塊間代表交流機(jī)制（Inter-Patch Delegation Mechanism），實(shí)現(xiàn)了線性復(fù)雜度，兼具了傳統(tǒng)基于 QKV 調(diào)整的稀疏注意力、狀態(tài)空間和線性注意力的優(yōu)點(diǎn)，且規(guī)避其缺點(diǎn)，為宇宙尺度的建模探索有效的注意力方案。

物質(zhì)世界建模的上下文長(zhǎng)度

Monte Carlo Attention 是為了解決 BigBang-Proton 框架的理論需求而開發(fā)的，盡管實(shí)際實(shí)現(xiàn)面臨硬件約束。推動(dòng)這一發(fā)展的基本假設(shè)包括幾個(gè)關(guān)鍵考慮因素。首先，對(duì)于自回歸預(yù)訓(xùn)練，二進(jìn)制塊編碼作為一種原生多模態(tài)方法，可以無(wú)縫地將所有數(shù)字?jǐn)?shù)據(jù)格式轉(zhuǎn)換為標(biāo)準(zhǔn)二進(jìn)制序列，從而對(duì)超長(zhǎng)上下文長(zhǎng)度提出了嚴(yán)格要求。

其次，理論-實(shí)驗(yàn)學(xué)習(xí)范式提供了在預(yù)訓(xùn)練期間跨尺度、結(jié)構(gòu)和學(xué)科整合來(lái)自歷史和正在進(jìn)行的科學(xué)實(shí)驗(yàn)的實(shí)驗(yàn)數(shù)據(jù)的潛力，這需要遠(yuǎn)超純自然語(yǔ)言預(yù)訓(xùn)練的上下文長(zhǎng)度。最后，在將宇宙視為單一實(shí)體的終極場(chǎng)景中，如果能將所有原子（10??）的信息轉(zhuǎn)換為用于預(yù)訓(xùn)練的單個(gè)序列，上下文長(zhǎng)度能否達(dá)到宇宙尺度？

將復(fù)雜物質(zhì)結(jié)構(gòu)轉(zhuǎn)換為序列的 token（本文用 token 同時(shí)指代傳統(tǒng)的 BPE token 和二進(jìn)制塊編碼中的 patch） 長(zhǎng)度估算提出了前所未有的上下文長(zhǎng)度需求。對(duì)于包含多組學(xué)數(shù)據(jù)和細(xì)胞結(jié)構(gòu)的綜合虛擬細(xì)胞整合，大約有 101? 個(gè)原子，每個(gè)原子需要 10-20 個(gè) token 來(lái)完整表示其位置、鍵合、相互作用和動(dòng)態(tài)狀態(tài)信息，總序列長(zhǎng)度達(dá)到約 101? 個(gè) token（1 千萬(wàn)億 token）。

同樣，對(duì)于涉及格點(diǎn) QCD 數(shù)據(jù)的 QCD 建模，包含 ~10? 個(gè)構(gòu)型，每個(gè)構(gòu)型有 10? 個(gè)格點(diǎn)和每個(gè)格點(diǎn) ~100 個(gè)浮點(diǎn)值，總計(jì) ~101? 字節(jié)的數(shù)據(jù)在考慮完整參數(shù)和理論描述表示后，轉(zhuǎn)換為約 102? 個(gè) token。這兩種場(chǎng)景都要求上下文長(zhǎng)度比當(dāng)前大語(yǔ)言模型的能力（101?-102? token vs. 典型的 10?-10? token 限制）高出多個(gè)數(shù)量級(jí)，這需要像 Monte Carlo Attention 這樣的新型注意力機(jī)制來(lái)實(shí)現(xiàn)宇宙尺度科學(xué)建模和模擬的上下文長(zhǎng)度要求。

為了在預(yù)訓(xùn)練中擴(kuò)展上下文窗口，主流大語(yǔ)言模型（如 DeepSeek V3、Qwen3、Llama3）通常采用兩階段訓(xùn)練，先在小部分?jǐn)?shù)據(jù)上訓(xùn)一個(gè)長(zhǎng)上下文預(yù)訓(xùn)練階段，在大部分?jǐn)?shù)據(jù)上訓(xùn)短的上下文，將上下文長(zhǎng)度從典型的 4096 token 擴(kuò)展到 128K token。

相比之下，Monte Carlo Attention 無(wú)須分階段訓(xùn)練，通過(guò)其塊間代表機(jī)制實(shí)現(xiàn)了理論上無(wú)限的上下文長(zhǎng)度，從根本上改變了預(yù)訓(xùn)練方法。這一突破對(duì)現(xiàn)有預(yù)訓(xùn)練技術(shù)和硬件設(shè)計(jì)具有深遠(yuǎn)影響。

首先，傳統(tǒng)的批次處理約束得到緩解，因?yàn)樽⒁饬τ?jì)算可以跨批次分布，從而能夠高效處理超長(zhǎng)序列。其次，計(jì)算復(fù)雜度從 O(L2) 降低到 O(L)，顯著減少了收斂所需的訓(xùn)練步數(shù)，可能改善損失收斂率和困惑度曲線。第三，通過(guò)將上下文長(zhǎng)度與 GPU 內(nèi)存限制解耦，Monte Carlo Attention 能夠訓(xùn)練比設(shè)備內(nèi)存容量長(zhǎng)多個(gè)數(shù)量級(jí)的序列。

最后，這種方法促進(jìn)了專門為長(zhǎng)序列處理設(shè)計(jì)的存內(nèi)計(jì)算架構(gòu)的發(fā)展。在不考慮 GPU 內(nèi)存約束的情況下，該方法可以實(shí)現(xiàn)與完整預(yù)訓(xùn)練語(yǔ)料庫(kù)序列長(zhǎng)度相匹配的有效上下文長(zhǎng)度。這一范式轉(zhuǎn)變需要開發(fā)能夠支持真正宇宙尺度序列處理的下一代硬件架構(gòu)。

Monte Carlo 架構(gòu)

BigBang-Proton 架構(gòu)由三個(gè)核心組件構(gòu)成：

1. Binary Patch Encoding 的 embedding 輸入 ：輸入詞匯表包含 0-255 的字節(jié)值和三個(gè)特殊 token，總大小為 259。每個(gè)輸入 token 通過(guò) one-hot 編碼轉(zhuǎn)換為 259 維稀疏向量，其中對(duì)應(yīng) token 索引位置為 1。該 one-hot 向量通過(guò)無(wú)偏置的線性層投影到維度 D 的稠密嵌入空間，形成最終的 token 嵌入表示。

2. Monte Carlo Attention ：利用塊間代表機(jī)制驅(qū)動(dòng)局部和全局信息交換，使上下文長(zhǎng)度隨層數(shù)呈指數(shù)級(jí)增長(zhǎng)，同時(shí)保持線性計(jì)算復(fù)雜度。

3. 前饋時(shí)序卷積網(wǎng)絡(luò) (TCN)：取代 Transformer 中傳統(tǒng)的前饋全連接網(wǎng)絡(luò)，以捕捉局部空間和時(shí)間模式。由于 TCN 能夠?qū)W習(xí)位置信息，因此消除了 Transformer 中使用的位置嵌入。

上下文長(zhǎng)度定義為 Transformer 單層在一次完整注意力計(jì)算中可讀取的極限，它與 Transformer 層數(shù)的深度無(wú)關(guān)。Transformer 中的直接信息流被限制在上下文長(zhǎng)度內(nèi)的 token 之間。在預(yù)訓(xùn)練中，批次間的信息流依賴于共享權(quán)重，而非注意力計(jì)算。

相比之下，卷積神經(jīng)網(wǎng)絡(luò)（CNN）的感受野隨網(wǎng)絡(luò)深度而擴(kuò)展。超對(duì)稱團(tuán)隊(duì)受此啟發(fā)，采用逐層操作來(lái)增強(qiáng)所有輸入嵌入之間的信息流，甚至跨越批次。Monte Carlo Attention 的關(guān)鍵創(chuàng)新在于每一層中的 Delegate 操作，從而實(shí)現(xiàn)動(dòng)態(tài)的 token 重組。

圖：嵌入向量在塊之間被重組。每個(gè)塊向其他塊發(fā)送代表，并從其他塊接收代表，通過(guò)注意力計(jì)算進(jìn)行信息交換。

塊間代表交流機(jī)制 Inter-Patch Delegation Mechanism

輸入嵌入被分組為特定塊大小（patch size）的塊，這與字節(jié)級(jí)語(yǔ)言模型（如 BLT、Megabyte、SpaceByte 和 BGPT）中在轉(zhuǎn)換為嵌入向量之前實(shí)施分塊的方法不同。受代議制政治系統(tǒng)的啟發(fā)，在每次逐層操作中，由 P 個(gè)字節(jié)輸入組成的每個(gè)塊會(huì)隨機(jī)（蒙特卡羅方式）或有選擇地指定一個(gè)字節(jié)作為代表 Delegate，與其他塊交換信息。在所有塊之間的 Delegate 過(guò)程之后，每個(gè)塊最多從其他塊接收 P-1 個(gè)代表，同時(shí)向其他塊發(fā)送相同數(shù)量的代表，而塊內(nèi)的字節(jié)數(shù)保持不變。注意力計(jì)算在每個(gè)塊內(nèi)執(zhí)行，其復(fù)雜度為 O(P2)。

每個(gè)代表字節(jié)包含來(lái)自其原始?jí)K和接收其他代表的塊的上下文信息，然后返回其原始?jí)K進(jìn)行后續(xù)的注意力計(jì)算。在圖所示的 toy model 中，為 3 個(gè)塊（紅色、藍(lán)色和綠色）設(shè)置 P=4。12 個(gè)字節(jié)的序列被重塑為一個(gè) 4×3 矩陣，然后轉(zhuǎn)置為一個(gè) 3×4 矩陣，再展平回一個(gè) 12 字節(jié)的序列。這種重組將序列 [1,2,3,4] 轉(zhuǎn)換為 [1,5,9,2]，其中字節(jié) 1 和 2 來(lái)自紅色塊，字節(jié) 5 來(lái)自藍(lán)色塊，字節(jié) 9 來(lái)自綠色塊。通過(guò) Delegate 和重組，信息在全局范圍內(nèi)流動(dòng)，而計(jì)算復(fù)雜度僅取決于塊大小的平方。

下圖展示了通過(guò)逐層塊間代表操作，信息流的上下文長(zhǎng)度如何增加。將塊大小 P 設(shè)為 32，序列長(zhǎng)度設(shè)為 40906，則每個(gè)塊可以容納 P-1 個(gè) Delegate 字節(jié)用于信息交換，每個(gè) Delegate 包含 P 長(zhǎng)度的上下文信息。

當(dāng)層數(shù)深度變大時(shí)，經(jīng)過(guò) N 層后的有效上下文長(zhǎng)度的遞歸關(guān)系為：

當(dāng) P=32 且 C(0)=0 時(shí)，可以計(jì)算不同層數(shù)的上下文長(zhǎng)度。結(jié)果如圖所示。

圖：逐層塊間代表操作驅(qū)動(dòng)信息流的上下文長(zhǎng)度以 的速度增長(zhǎng)，其中 P 是塊大小，N 是層數(shù)。對(duì)于塊大小=32，在第一層，信息可達(dá) 992，在第二層可達(dá) 32736。

Delegate 操作可以形式化地定義為一個(gè)包含四個(gè)關(guān)鍵步驟的分層過(guò)程，這些步驟反映了實(shí)際的計(jì)算流程。首先，輸入序列 被分解為 個(gè)塊，每個(gè)塊包含 P 個(gè) token：

其中 表示長(zhǎng)度為 的輸入序列， 是批次大小， 是隱藏維度， 表示分解為 個(gè)塊，每個(gè)塊包含 個(gè) token，且 。

其次，對(duì)每個(gè)塊應(yīng)用 1×1 卷積操作以生成將被發(fā)送到其他塊的 Delegate token，并且 Delegate 映射從每個(gè)塊中選擇和分發(fā)代表性 token：

其中 表示 Delegate token 組， 表示從第 個(gè)塊 生成的 Delegate token， 表示從塊 中為塊 選擇的代表性字節(jié)。這些 Delegate token 被置換以與原始局部塊對(duì)齊。

第三，Delegate token 被分發(fā)到其他塊，并與原始局部 token 連接，形成一個(gè)包含局部和全局信息的增強(qiáng)表示：

其中 表示第 個(gè)塊的增強(qiáng)上下文， 表示從鄰近塊 接收的 Delegate token 集。

最后，在增強(qiáng)表示上計(jì)算自注意力，以促進(jìn)局部和全局上下文之間的信息交換。標(biāo)準(zhǔn)的自注意力機(jī)制可以描述為

其中 Q,K,V 是通過(guò)線性投影從增強(qiáng)表示 c 導(dǎo)出的查詢、鍵和值矩陣， 是鍵的維度。

結(jié)合塊間代表操作，每層的注意力計(jì)算可以形式化地表示為：

其中 表示第 層第 個(gè)塊的注意力輸出， 是 Delegate 操作后第 個(gè)塊的查詢、鍵和值矩陣。

最終的輸出隱藏狀態(tài)通過(guò)局部塊表示和重組后的塊表示之和計(jì)算得出，并使用殘差連接，以確保穩(wěn)定的梯度傳播和信息持久性：

其中 表示注意力計(jì)算和殘差連接后的第 個(gè)塊的最終輸出，Linear 表示注意力計(jì)算后應(yīng)用的線性變換。

塊重組變換可以表示為一個(gè)雙射映射：

每個(gè)塊內(nèi)的注意力計(jì)算復(fù)雜度為：

全局信息流的復(fù)雜度為：

其中 是塊的數(shù)量。

表：當(dāng) P=32 時(shí)，不同層數(shù)對(duì)應(yīng)的上下文長(zhǎng)度

表征退化、循環(huán)注意力與稀疏性

標(biāo)準(zhǔn) Transformer 依賴于完整注意力計(jì)算，使輸入嵌入能夠在預(yù)定義的上下文長(zhǎng)度內(nèi)從其他 token 獲得表征。Monte Carlo Attention 采用塊間代表機(jī)制來(lái)實(shí)現(xiàn)全局表示交換，從而達(dá)到超大的有效上下文長(zhǎng)度。

然而，這種間接的注意力計(jì)算可能導(dǎo)致在迭代表征傳播過(guò)程中信息退化。為了解決這一局限性，超對(duì)稱團(tuán)隊(duì)引入了循環(huán) Monte Carlo Attention，通過(guò)重復(fù)單層操作多次來(lái)實(shí)現(xiàn)。

• 標(biāo)準(zhǔn) Transformer : Layer? → Layer? → Layer? → ...

• 循環(huán) Monte Carlo Attention : Layer? N→ Layer? N→ Layer? N→ ...

前饋 TCN

替換傳統(tǒng)全連接前饋網(wǎng)絡(luò)，超對(duì)稱團(tuán)隊(duì)提出了一個(gè)具有增強(qiáng)模式捕捉能力的時(shí)序卷積塊（TCN）。TCN 通過(guò)堆疊核大小為 K 的一維卷積來(lái)實(shí)現(xiàn)多尺度處理，從而提取分層的時(shí)間特征。

令 表示輸入張量。TCN 對(duì)該輸入應(yīng)用多層一維卷積。每一層可以表示為：

其中 ， 表示核大小為 K 的一維卷積操作，ReLU 是逐元素應(yīng)用的激活函數(shù)。 這些卷積層的堆疊使網(wǎng)絡(luò)能夠在不同尺度上捕捉分層的時(shí)間特征。具體來(lái)說(shuō)，由于連續(xù)卷積的累積效應(yīng)，網(wǎng)絡(luò)中更深的層可以捕捉更長(zhǎng)距離的依賴關(guān)系。

此外，所提出的 TCN 保持了局部-全局平衡，當(dāng)注意力機(jī)制處理長(zhǎng)距離依賴時(shí)，TCN 專注于細(xì)粒度的局部模式發(fā)現(xiàn)。這是通過(guò)在 TCN 架構(gòu)中應(yīng)用擴(kuò)張卷積來(lái)實(shí)現(xiàn)的，這使得模型能夠在不顯著增加參數(shù)數(shù)量的情況下覆蓋更寬的感受野。對(duì)于擴(kuò)張因子 d，卷積操作變?yōu)椋?/p>

其中 表示擴(kuò)張因子為 d 的擴(kuò)張卷積操作。

值得注意的是，由于 TCN 固有的卷積特性，BigBang-Proton 能夠直接從輸入序列中學(xué)習(xí)空間和位置信息，從而消除了傳統(tǒng) Transformer 架構(gòu)中通常需要的顯式位置嵌入。

最后，TCN 塊的輸出通過(guò)殘差連接與來(lái)自注意力模塊的已關(guān)注特征相結(jié)合：

其中 是 TCN 最后一層的輸出， 表示從注意力機(jī)制獲得的已關(guān)注特征。 這種設(shè)計(jì)確保了模型能夠同時(shí)受益于注意力機(jī)制提供的全局上下文和 TCN 捕捉的細(xì)粒度局部模式，從而提升整體性能。

與稀疏注意力、狀態(tài)空間、線性注意力的比較

現(xiàn)有稀疏注意力與 Monte Carlo Attention 之間的根本區(qū)別在于其核心計(jì)算機(jī)制。稀疏注意力方法（比如 NSA 和 MoBA）采用基于選擇的機(jī)制，通過(guò)過(guò)濾鍵值對(duì)來(lái)降低計(jì)算復(fù)雜度，即選擇 token 子集進(jìn)行注意力計(jì)算。

NSA 利用三種復(fù)雜的策略：通過(guò)塊級(jí)聚合進(jìn)行 token 壓縮，通過(guò)塊級(jí) top-n 識(shí)別進(jìn)行 token 選擇，以及滑動(dòng)窗口機(jī)制以保留局部上下文，從而從 N 個(gè) token 中動(dòng)態(tài)構(gòu)建緊湊表示，通過(guò)選擇 top-K 子集實(shí)現(xiàn)。

MoBA 采用受混合專家啟發(fā)的方法，將上下文劃分為塊，并應(yīng)用 top-k 門控機(jī)制進(jìn)行選擇性注意力。相比之下，Monte Carlo Attention 通過(guò)塊間代表采用基于重組的機(jī)制，通過(guò)將全局上下文壓縮為代表性 token 并在塊之間交換，從而在塊之間實(shí)現(xiàn)間接的信息傳播，而不是在選定的 token 之間進(jìn)行直接注意力。

這些核心機(jī)制的差異導(dǎo)致了稀疏注意力方法的關(guān)鍵劣勢(shì)。首先，稀疏注意力存在選擇偏差和信息丟失問題，未被選中的 token 被丟棄，導(dǎo)致忽略全局依賴和協(xié)調(diào)碎片化，而 Monte Carlo Attention 通過(guò)受控的 Delegate 機(jī)制保留了關(guān)鍵信息。

其次，稀疏注意力遇到計(jì)算瓶頸，需要在選定的 token 之間進(jìn)行 復(fù)雜度的計(jì)算，而 Monte Carlo Attention 通過(guò)局部 2P 注意力實(shí)現(xiàn)了高效的全局交換，具有更優(yōu)的 上 下 文 長(zhǎng) 度 復(fù)雜度。第三，由于選擇約束，稀疏注意力的上下文建模能力有限，而 Monte Carlo Attention 能夠?qū)崿F(xiàn)指數(shù)級(jí)的上下文長(zhǎng)度擴(kuò)展。

結(jié)構(gòu)化狀態(tài)空間序列模型（S4）及其后繼者（包括 RetNet、RWKV 和 Mamba）是一類將隱藏張量視為狀態(tài)空間的序列模型，S4 通過(guò)結(jié)構(gòu)化線性動(dòng)力系統(tǒng)（使用對(duì)角 A 矩陣）來(lái)高效地建模長(zhǎng)程依賴，而 Mamba 引入了輸入依賴的參數(shù)選擇，RetNet 則結(jié)合了門控機(jī)制以增強(qiáng)選擇性信息傳播。Monte Carlo Attention 與狀態(tài)空間模型之間的根本區(qū)別在于其核心信息流機(jī)制。

S4 采用順序狀態(tài)傳播，信息通過(guò)線性遞推關(guān)系 流動(dòng)，創(chuàng)建了一個(gè)馬爾可夫依賴鏈，限制了每個(gè)狀態(tài)只能直接訪問前一個(gè)狀態(tài)。相比之下，Monte Carlo Attention 通過(guò)塊間代表實(shí)現(xiàn)直接的全局信息交換，允許任何塊通過(guò) delegate token 重組訪問來(lái)自任何其他塊的信息。

S4 存在固有的局限性，例如由于線性時(shí)不變性導(dǎo)致的建模靈活性有限，由有限維狀態(tài)向量引起的信息瓶頸，通過(guò)多步傳播捕捉長(zhǎng)程依賴的困難，以及對(duì)復(fù)雜非線性關(guān)系建模的表達(dá)能力不足。雖然 Mamba 通過(guò)輸入依賴的參數(shù)選擇解決了 S4 的一些局限性，引入了選擇性狀態(tài)空間機(jī)制 （其中 是動(dòng)態(tài)計(jì)算的），但它仍然繼承了低秩表示的基本近似誤差，并且在狀態(tài)轉(zhuǎn)換計(jì)算期間容易受到數(shù)值不穩(wěn)定的影響。

Monte Carlo Attention 通過(guò)在增強(qiáng)的局部-全局上下文中進(jìn)行精確計(jì)算來(lái)保持完整注意力的表達(dá)能力，從而超越了這些局限性，避免了 S4 的線性時(shí)不變性約束和低秩近似相關(guān)的精度損失。這種方法能夠在保持局部精度的同時(shí)實(shí)現(xiàn)真正的全局上下文傳播，繞過(guò)了線性注意力模型和狀態(tài)空間近似中固有的信息瓶頸，后者為了計(jì)算效率而犧牲了表達(dá)能力。

作為 S4 模型的一種變體，傳統(tǒng)的線性注意力模型通過(guò)從根本上進(jìn)行近似來(lái)實(shí)現(xiàn)計(jì)算效率，但這犧牲了信息保真度。這些方法消除了 softmax 操作，并通過(guò)核函數(shù)將查詢和鍵映射到隱藏表示，然后通過(guò)鍵和值的右乘積計(jì)算注意力，將復(fù)雜度從 降低到 (其中 N 是序列長(zhǎng)度，d 是矩陣維度)。

TransNormer 進(jìn)一步通過(guò)用歸一化操作替代縮放來(lái)解決無(wú)界梯度問題。然而，與 Monte Carlo Attention 相比，這些方法存在關(guān)鍵局限性。線性注意力方法通過(guò)核函數(shù)將高維 Q、K、V 矩陣映射到低維特征空間，從根本上損害了表示質(zhì)量，不可避免地丟失了高階統(tǒng)計(jì)信息和復(fù)雜的 token 交互。相比之下，Monte Carlo Attention 在增強(qiáng)的局部-全局上下文中通過(guò)精確的自注意力計(jì)算保持了完整的注意力表達(dá)能力，確保了完整的信息保留。

雖然線性注意力實(shí)現(xiàn)了 的復(fù)雜度，但這種效率是以有限的全局上下文集成為代價(jià)的，因?yàn)楫?dāng)序列長(zhǎng)度 L 變得極大時(shí)，低秩近似無(wú)法充分捕捉長(zhǎng)程依賴。然而，Monte Carlo Attention 通過(guò)其塊間代表機(jī)制實(shí)現(xiàn)了近乎無(wú)限的全局信息流，允許有效上下文長(zhǎng)度隨層數(shù)深度呈指數(shù)級(jí)擴(kuò)展，同時(shí)保持線性計(jì)算復(fù)雜度。

開源鏈接

• 論文 : https://arxiv.org/abs/2410.00129

• GitHub : https://github.com/supersymmetry-technologies/BigBang-Proton

• Hugging Face : https://huggingface.co/SuperSymmetryTechnologies/BigBang-Proton

* 本文為 BigBang-Proton 系列報(bào)道第二篇。在后續(xù)文章中，我們將帶來(lái)更多關(guān)于其核心技術(shù)、前沿應(yīng)用與未來(lái)規(guī)劃的深度解讀，敬請(qǐng)關(guān)注 CSDN AI 科技大本營(yíng)和文章合集。