超長(zhǎng)序列推理時(shí)的巨大開銷如何降低？

中國(guó)科學(xué)院自動(dòng)化所李國(guó)齊、徐波團(tuán)隊(duì)發(fā)布的類腦脈沖大模型SpikingBrain (瞬悉)-1.0提出了新思路。

SpikingBrain借鑒大腦信息處理機(jī)制，具有線性/近線性復(fù)雜度，在超長(zhǎng)序列上具有顯著速度優(yōu)勢(shì)。

在GPU上1M長(zhǎng)度下TTFT 速度相比主流大模型提升26.5x，4M長(zhǎng)度下保守估計(jì)速度提升超過100x；在手機(jī)CPU端64k-128k-256k長(zhǎng)度下較Llama3.2的同規(guī)模模型Decoding速度提升4.04x-7.52x-15.39x。

SpikingBrain適配了面向沐曦MetaX國(guó)產(chǎn)GPU集群的高效訓(xùn)練和推理框架、Triton算子庫、模型并行策略以及集群通信原語，表明了構(gòu)建國(guó)產(chǎn)自主可控的新型非Transformer大模型架構(gòu)生態(tài)的可行性。

為什么類腦計(jì)算有望解決這一難題？

現(xiàn)有的主流大模型基于簡(jiǎn)單的神經(jīng)元和復(fù)雜的網(wǎng)絡(luò)架構(gòu)，在Scaling law驅(qū)動(dòng)下取得了巨大成功。

中國(guó)科學(xué)院團(tuán)隊(duì)在前期的工作中證明，具有復(fù)雜內(nèi)生動(dòng)態(tài)的脈沖神經(jīng)元可以在數(shù)學(xué)上等價(jià)為若干簡(jiǎn)單脈沖神經(jīng)元的組合。

這說明了存在使用由復(fù)雜神經(jīng)元組成的小規(guī)模網(wǎng)絡(luò)代替由簡(jiǎn)單神經(jīng)元組成的大規(guī)模網(wǎng)絡(luò)的可能性。

受此啟發(fā)，一條“基于內(nèi)生復(fù)雜性”的通用智能實(shí)現(xiàn)思路被提出，即找到一條融合神經(jīng)元內(nèi)部豐富動(dòng)力學(xué)特性的類腦模型發(fā)展路徑。

SpikingBrain-1.0就是這一思路下的初步嘗試。

大模型時(shí)代的新視角

人腦是目前唯一已知的通用智能系統(tǒng)，包含約1000億神經(jīng)元和約1000萬億突觸數(shù)量、具有豐富的神經(jīng)元種類、不同神經(jīng)元又具有豐富的內(nèi)部結(jié)構(gòu)，但功耗僅20W左右。

現(xiàn)有大模型通常基于Transformer架構(gòu)，增加網(wǎng)絡(luò)規(guī)模、算力資源和數(shù)據(jù)量提升智能水平，但二次方復(fù)雜度使其訓(xùn)練和推理開銷巨大，超長(zhǎng)序列處理能力受限。例如，當(dāng)前國(guó)內(nèi)外主流大模型僅支持64k或128k長(zhǎng)度的序列訓(xùn)練，支持1M以下長(zhǎng)度推理。

其基本計(jì)算單元為點(diǎn)神經(jīng)元模型：簡(jiǎn)單乘加單元后接非線性函數(shù)，這條簡(jiǎn)單神經(jīng)元加網(wǎng)絡(luò)規(guī)模拓展的技術(shù)路徑可以被稱為“基于外生復(fù)雜性”的通用智能實(shí)現(xiàn)方法。

相比之下， “基于內(nèi)生復(fù)雜性”的通用智能實(shí)現(xiàn)方法的目標(biāo)是，充分利用生物神經(jīng)網(wǎng)絡(luò)在神經(jīng)元和神經(jīng)環(huán)路上的結(jié)構(gòu)和功能特性，找到構(gòu)建具有生物合理性和計(jì)算高效性的神經(jīng)網(wǎng)絡(luò)新路徑。

因此，探索腦科學(xué)與人工智能基礎(chǔ)模型架構(gòu)之間的橋梁、構(gòu)建新一代非Transformer的類腦基礎(chǔ)模型架構(gòu)，或?qū)⒁I(lǐng)下一代人工智能的發(fā)展方向、為實(shí)現(xiàn)國(guó)產(chǎn)自主可控類腦大模型生態(tài)提供基礎(chǔ)積累。

核心技術(shù)

SpikingBrain-1.0基于脈沖神經(jīng)元構(gòu)建了線性（混合）模型架構(gòu)，具有線性（SpikingBrain-7B）及近線性復(fù)雜度（SpikingBrain-76B，激活參數(shù)量12B）的類腦基礎(chǔ)模型（圖1）。

△圖1. SpikingBrain框架概覽

為解決脈沖編碼時(shí)的性能退化問題，構(gòu)建了自適應(yīng)閾值神經(jīng)元模型，模擬生物神經(jīng)元脈沖發(fā)放的核心過程，隨后通過虛擬時(shí)間步策略實(shí)現(xiàn)“電位-脈沖”的轉(zhuǎn)換，將整數(shù)脈沖計(jì)數(shù)重新展開為稀疏脈沖序列。

借助動(dòng)態(tài)閾值脈沖化信息編碼方案，可以將模型中計(jì)算量占比90%以上的稠密連續(xù)值矩陣乘法，替換為支持事件驅(qū)動(dòng)的脈沖化算子，以實(shí)現(xiàn)高性能與低能耗二者兼顧：脈沖神經(jīng)元僅在膜電勢(shì)累積達(dá)到閾值時(shí)發(fā)放脈沖事件，脈沖到達(dá)時(shí)觸發(fā)下游神經(jīng)元活動(dòng)，無脈沖時(shí)則可處于低能耗靜息狀態(tài)。

進(jìn)一步，網(wǎng)絡(luò)層面的MoE架構(gòu)結(jié)合神經(jīng)元層面的稀疏事件驅(qū)動(dòng)計(jì)算，可提供微觀-宏觀層面的稀疏化方案，體現(xiàn)按需計(jì)算的高效算力分配。

該團(tuán)隊(duì)在理論上建立了脈沖神經(jīng)元內(nèi)生動(dòng)力學(xué)與線性注意力模型之間的聯(lián)系，揭示了現(xiàn)有線性注意力機(jī)制是樹突計(jì)算的特殊簡(jiǎn)化形式，從而清晰地展示了一條不斷提升模型復(fù)雜度和性能的新型可行路徑。

基于這一理解以及團(tuán)隊(duì)前期工作,團(tuán)隊(duì)構(gòu)建了與現(xiàn)有大模型兼容的通用模型轉(zhuǎn)換技術(shù)和高效訓(xùn)練范式，可以將標(biāo)準(zhǔn)的自注意力機(jī)制轉(zhuǎn)換為低秩的線性注意力模型，并適配了所提出的脈沖化編碼框架。

此外，為實(shí)現(xiàn)國(guó)產(chǎn)算力集群對(duì)類腦脈沖大模型的全流程訓(xùn)練和推理支持，團(tuán)隊(duì)開發(fā)了面向沐曦MetaX國(guó)產(chǎn)GPU集群的高效訓(xùn)練和推理框架、Triton算子庫、模型并行策略以及集群通信原語。

SpikingBrain-7B 和SpikingBrain-76B分別為層間混合純線性模型和層內(nèi)混合的混合線性 MoE 模型（圖2）。

△圖2. SpikingBrain網(wǎng)絡(luò)架構(gòu)

其中SpikingBrain-7B由線性注意力和滑窗注意力1:1層間堆疊而成。

而SpikingBrain-76B則包含128個(gè)sink token、16個(gè)路由專家以及1個(gè)共享專家；對(duì)于線性層，在第[1, 2, 3, 5, 7, 9, 11] 層布置了7個(gè)稠密 FFN，其余層均實(shí)現(xiàn)為MoE層；對(duì)于注意力模塊在第[7, 14, 21, 28]層采用線性注意力+Softmax注意力(LA+FA)組合，在其他層均采用線性注意力+滑窗注意力（LA+SWA）組合。

在推理階段，SpikingBrain利用脈沖編碼將激活值轉(zhuǎn)換為整數(shù)計(jì)數(shù)用于GPU執(zhí)行，或轉(zhuǎn)換為脈沖序列用于事件驅(qū)動(dòng)的神經(jīng)形態(tài)硬件。

性能亮點(diǎn)

SpikingBrain1.0的長(zhǎng)序列訓(xùn)練效率顯著提升。SpikingBrain-1.0-7B模型能以極低的數(shù)據(jù)量（約為主流大模型的2%）,實(shí)現(xiàn)與眾多開源Transformer模型相媲美的通用語言建模性能（表1）。

SpikingBrain-1.0-76B混合線形模型通過擴(kuò)展更多的參數(shù)量和更精細(xì)的注意力設(shè)計(jì)，基本保持了基座模型的性能，能使用更少的激活參數(shù)接近甚至優(yōu)于Llama2-70B、Mixtral-8*7B、Gemma2-27B等先進(jìn)的Transformer模型（表2）。

SpikingBrain-1.0-7B模型在Huggingface框架下適配了多卡序列并行推理（使用ZeCO加上P2P通信），并支持4M長(zhǎng)度的Prefill。結(jié)果顯示，相比于使用標(biāo)準(zhǔn)注意力和A2A通信的Qwen baseline，SpikingTime-1.0-7B在512K和1M長(zhǎng)度下TTFT（提交提示到生成第一個(gè)Token所需的時(shí)間）加速分別達(dá)到13.88倍和26.5倍，且隨序列長(zhǎng)度和卡數(shù)擴(kuò)展具有幾乎恒定的時(shí)間開銷，在4M長(zhǎng)度下Qwen baseline已經(jīng)難以評(píng)測(cè)，根據(jù)擬合scaling曲線，保守估計(jì)速度提升超過100倍（表4）。

團(tuán)隊(duì)將壓縮到1B的SpikingBrain-1.0部署到CPU手機(jī)端推理框架上，在64k-128k-256k長(zhǎng)度下較Llama3.2的1B模型Decoding速度分別提升4.04x-7.52x-15.39x。

圖3 基于CPU移動(dòng)推理框架下，不同輸出長(zhǎng)度的解碼速度比較

為了在國(guó)產(chǎn)曦云C550算力集群上進(jìn)行訓(xùn)練/推理適配，團(tuán)隊(duì)對(duì)訓(xùn)練框架（Megatron、Colossal-AI）進(jìn)行Triton算子加速和通信適配等優(yōu)化，能在集群上保持百卡規(guī)模訓(xùn)練的數(shù)周穩(wěn)定運(yùn)行；在訓(xùn)練效率上，7B模型進(jìn)行8k長(zhǎng)度訓(xùn)練的MFU達(dá)到23.4%，TGS per GPU達(dá)到1558 tokens/s（Megatron框架、DP rank為8，PP rank為4，PP-micro batch size為2，global batch size為512）。

對(duì)話Demo和網(wǎng)絡(luò)試用端口：團(tuán)隊(duì)提供了SpikingBrain-1.0-76B模型的網(wǎng)絡(luò)端的試用端口供大家體驗(yàn)，該模型基于vLLM推理框架部署在沐曦MetaX GPU集群上，可以支持?jǐn)?shù)百人的并發(fā)請(qǐng)求。為支持類腦研究生態(tài)的構(gòu)建，團(tuán)隊(duì)開源了SpikingBrain-1.0-7B模型（詳見技術(shù)報(bào)告）。

總結(jié)

本次發(fā)布的國(guó)產(chǎn)自主可控類腦脈沖大模型探索了脈沖神經(jīng)元內(nèi)生復(fù)雜神經(jīng)動(dòng)力學(xué)與線性注意力模型之間的機(jī)制聯(lián)系，設(shè)計(jì)了線性模型架構(gòu)和基于轉(zhuǎn)換的異構(gòu)模型架構(gòu)，通過動(dòng)態(tài)閾值脈沖化解決了脈沖驅(qū)動(dòng)限制下的大規(guī)模類腦模型性能退化問題，實(shí)現(xiàn)了國(guó)產(chǎn)GPU算力集群對(duì)類腦脈沖大模型訓(xùn)練和推理的全流程支持。

超長(zhǎng)序列的建模在復(fù)雜多智能體模擬、DNA序列分析、分子動(dòng)力學(xué)軌跡等超長(zhǎng)序列科學(xué)任務(wù)建模場(chǎng)景中將具有顯著的潛在效率優(yōu)勢(shì)。未來該團(tuán)隊(duì)將進(jìn)一步探索神經(jīng)元內(nèi)生復(fù)雜動(dòng)態(tài)與人工智能基礎(chǔ)算子之間的機(jī)制聯(lián)系，構(gòu)建神經(jīng)科學(xué)和人工智能之間的橋梁，期望通過整合生物學(xué)見解來突破現(xiàn)有人工智能瓶頸，進(jìn)而實(shí)現(xiàn)低功耗、高性能、支持超長(zhǎng)上下文窗口的類腦通用智能計(jì)算模型，啟迪更低功耗的下一代神經(jīng)形態(tài)計(jì)算理論和芯片設(shè)計(jì)。

網(wǎng)絡(luò)端試用端口網(wǎng)址：https://controller-fold-injuries-thick.trycloudflare.com

中文技術(shù)報(bào)告網(wǎng)址：https://github.com/BICLab/SpikingBrain-7B/blob/main/SpikingBrain_Report_Chi.pdf

英文技術(shù)報(bào)告網(wǎng)址：https://arxiv.org/abs/2509.05276 

模型代碼網(wǎng)址：https://github.com/BICLab/SpikingBrain-7B