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

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

出品丨AI 科技大本營（ID：rgznai100）

超對稱技術(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ì)世界建模的上下文長度

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

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

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

同樣，對于涉及格點(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。這兩種場景都要求上下文長度比當(dāng)前大語言模型的能力（101?-102? token vs. 典型的 10?-10? token 限制）高出多個(gè)數(shù)量級(jí)，這需要像 Monte Carlo Attention 這樣的新型注意力機(jī)制來實(shí)現(xiàn)宇宙尺度科學(xué)建模和模擬的上下文長度要求。

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

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

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

最后，這種方法促進(jìn)了專門為長序列處理設(shè)計(jì)的存內(nèi)計(jì)算架構(gòu)的發(fā)展。在不考慮 GPU 內(nèi)存約束的情況下，該方法可以實(shí)現(xiàn)與完整預(yù)訓(xùn)練語料庫序列長度相匹配的有效上下文長度。這一范式轉(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 通過 one-hot 編碼轉(zhuǎn)換為 259 維稀疏向量，其中對應(yīng) token 索引位置為 1。該 one-hot 向量通過無偏置的線性層投影到維度 D 的稠密嵌入空間，形成最終的 token 嵌入表示。

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

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

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

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

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

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

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

每個(gè)代表字節(jié)包含來自其原始?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 來自紅色塊，字節(jié) 5 來自藍(lán)色塊，字節(jié) 9 來自綠色塊。通過 Delegate 和重組，信息在全局范圍內(nèi)流動(dòng)，而計(jì)算復(fù)雜度僅取決于塊大小的平方。

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

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

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

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

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

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

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

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

第三，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 是通過線性投影從增強(qiáng)表示 c 導(dǎo)出的查詢、鍵和值矩陣， 是鍵的維度。

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

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

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

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

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

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

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

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

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

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

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

然而，這種間接的注意力計(jì)算可能導(dǎo)致在迭代表征傳播過程中信息退化。為了解決這一局限性，超對稱團(tuán)隊(duì)引入了循環(huán) Monte Carlo Attention，通過重復(fù)單層操作多次來實(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ò)，超對稱團(tuán)隊(duì)提出了一個(gè)具有增強(qiáng)模式捕捉能力的時(shí)序卷積塊（TCN）。TCN 通過堆疊核大小為 K 的一維卷積來實(shí)現(xiàn)多尺度處理，從而提取分層的時(shí)間特征。

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

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

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

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

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

最后，TCN 塊的輸出通過殘差連接與來自注意力模塊的已關(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ī)制，通過過濾鍵值對來降低計(jì)算復(fù)雜度，即選擇 token 子集進(jìn)行注意力計(jì)算。

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

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

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

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

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

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

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

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

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

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

雖然線性注意力實(shí)現(xiàn)了 的復(fù)雜度，但這種效率是以有限的全局上下文集成為代價(jià)的，因?yàn)楫?dāng)序列長度 L 變得極大時(shí)，低秩近似無法充分捕捉長程依賴。然而，Monte Carlo Attention 通過其塊間代表機(jī)制實(shí)現(xiàn)了近乎無限的全局信息流，允許有效上下文長度隨層數(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ù)文章中，我們將帶來更多關(guān)于其核心技術(shù)、前沿應(yīng)用與未來規(guī)劃的深度解讀，敬請關(guān)注 CSDN AI 科技大本營和文章合集。