用于神经解码的Von-Neumann State-Space Transformer

# 用于神经解码的冯·诺依曼状态空间Transformer
来源：https://arxiv.org/html/2608.25088
Morteza Sarafyazd††thanks: 职位：BrainCo，美国萨默维尔\.

###### 摘要

皮层计算呈现出显著的低维特性：由神经群体活动承载的一小部分潜变量，引导着单个神经元的高维响应。我们的目标是提升样本效率——构建能在有限数据和较小参数预算下实现良好解码的模型。在标准Transformer层中，前馈模块对每个token应用相同的算子。我们提出一种受冯·诺依曼启发的高效计算假设，作为神经解码的替代方案：一个控制器解码指令，然后执行token特定的算子；而常见的实现方式——软专家混合（soft mixture of experts）——仅混合它们的输出，而非算子本身。我们引入了冯·诺依曼状态空间Transformer（VN–SST），这是一种记忆增强型Transformer，其前馈模块是一个低秩指令库：一个共享基础算子加上一组小型学习的低秩指令，每个token的代码从该库中合成该token实际使用的权重矩阵。代码通过携带的状态空间记忆的低维投影读取，因此一个缓慢的潜在轨迹充当指令指针——反映了低维动态如何可能引导皮层计算。在三个运动皮层神经解码基准测试中，VN–SST的数据效率远高于现代Transformer，两者都联合预测神经脉冲并解码行为。该模型在最稀缺的基准测试上以巨大优势获胜，在另外两个基准测试上处于领先地位，并将更长的上下文窗口转化为逐渐上升而非下降的准确率。我们评估发现，该网络将大型指令库压缩为每个token仅需几比特，因此程序容量作为控制通道，而非精度杠杆。相同的模型在用于语言建模（LLMs）的两个小型文本基准测试中也更具参数效率，这表明其具有通用性机制。

## 1 引言

皮层计算呈现出显著的低维特性：跨越数千个神经元的群体活动由一小部分潜变量组织，这些变量的缓慢轨迹引导着单个细胞的更快、更高维响应。我们关注的是样本高效的神经解码——从有限数据和较小模型中实现准确的行为读出。为此，我们采用受冯·诺依曼启发的高效计算假设，并将其构建到Transformer主干中。冯·诺依曼（存储程序）机具有相同的结构：一个小型、缓慢的控制器获取并解码一条指令，然后一个执行单元运行该指令指定的算子，因此一个紧凑的程序引导着大规模、输入特定的计算。相比之下，在标准Transformer中，前馈网络（FFN）对每个token应用相同的固定权重；我们则为FFN配备一个控制器——一个缓慢变化的状态空间记忆，其读出逐token选择FFN执行的算子。

一个自然的初步尝试是软专家混合（soft–MoE），其中控制器混合多个固定FFN专家的输出\[1 (https://arxiv.org/html/2608.25088#bib.bib1)\]。这是一个有用的容量调节手段，但从冯·诺依曼的角度看，它并未执行程序：它混合的是固定算子的结果，而不是构建一个token特定的算子。我们选择让控制器从共享的小型低秩“指令库”中逐token合成FFN的权重本身，该指令库添加到共享基础算子中（公式7 (https://arxiv.org/html/2608.25088#S2.E7)，第2.2节 (https://arxiv.org/html/2608.25088#S2.SS2)）。每个token的代码在算子的低维流形上选择一个点，因此执行的映射是真正token特定的，而每条添加的指令成本很低——这暴露了指令库大小作为一个独立的扩展轴，与参数和数据并列，并将程序转化为我们可以测量的具体控制信号。

程序是生成而非查找的。我们不从原始token读取指令；而是根据跨序列携带的缓慢变化选择性状态空间（SSM）记忆的低维读出来生成它。然后，缓慢的潜在轨迹充当指令指针——一个获取要运行的哪个算子的低维动态——将神经科学直觉转化为机制：一个小型、缓慢的潜在变量驱动着更高复杂度、token特定的计算。

我们在VN–SST（冯·诺依曼状态空间Transformer）中实现了这一点：一个可直接替换的、记忆增强型Transformer——包含局部注意力、一个选择性SSM和一个快速权重记忆，由一个协调的控制器驱动——其前馈算子由携带的低维状态驱动的低秩指令库逐token合成（第2.2节 (https://arxiv.org/html/2608.25088#S2.SS2)）。与扩展律研究类似\[2 (https://arxiv.org/html/2608.25088#bib.bib2)\]，我们在三个运动皮层解码编解码器上，从参数、数据和上下文三个维度将VN–SST与现代Transformer进行对比（第4.1节 (https://arxiv.org/html/2608.25088#S4.SS1)–4.3节 (https://arxiv.org/html/2608.25088#S4.SS3)）。这个可编程算子在样本效率上显著更高：在有限数据预算下，它在数据最稀缺的编解码器上以巨大优势击败Transformer（解码R² 0.35 vs. 0.21），在所有三个编解码器的每个数据预算下均领先，并且——独特之处在于——将更长的上下文窗口转化为逐渐上升而非下降的解码准确率。然后，一个控制比特诊断工具（第4.4节 (https://arxiv.org/html/2608.25088#S4.SS4)）使得冯·诺依曼的主张变得可测量：网络将32条指令库压缩为每个token仅需几比特（≈3.4–6.5个有效算子），因此程序容量表现为控制通道，而非精度杠杆——这是任何输出混合机制无法报告的信号。

## 2 模型

两个模型共享一个任务无关的骨干网络，将隐藏序列h∈RB×T×dh\in\mathbb{R}^{B\times T\times d}（批量大小B，序列长度T，隐藏宽度d）映射到相同形状的输出；只有输入/输出头不同。对于神经编解码器，输入头是从N个分箱发放率到隐藏维度的线性映射，有两个输出头：一个线性映射回N个发放率（脉冲延续）和一个线性行为读出头。

### 2.1 Transformer基线

基线是一个现代仅解码器Transformer\[3 (https://arxiv.org/html/2608.25088#bib.bib3),4 (https://arxiv.org/html/2608.25088#bib.bib4)\]：由相同层堆叠而成，每层结合了预层归一化RMSNorm\[5 (https://arxiv.org/html/2608.25088#bib.bib5)\]、旋转位置嵌入（RoPE）\[6 (https://arxiv.org/html/2608.25088#bib.bib6)\]、因果多头自注意力和SwiGLU\[7 (https://arxiv.org/html/2608.25088#bib.bib7)\]前馈模块。每一层通过两个残差更新应用注意力和前馈模块：

h←h+Attn(RMSNorm(h)),h←h+SwiGLU(RMSNorm(h)),\qquad h\leftarrow h+\mathrm{Attn}\!\big(\mathrm{RMSNorm}(h)\big),\qquad h\leftarrow h+\mathrm{SwiGLU}\!\big(\mathrm{RMSNorm}(h)\big),\tag{1}

其中前馈模块用另一个线性投影门控其输入的一个线性投影，

SwiGLU(x)=W2(SiLU(W1x)⊙W3x).\mathrm{SwiGLU}(x)=W_{2}\big(\mathrm{SiLU}(W_{1}x)\odot W_{3}x\big).\tag{2}

三个权重矩阵W1,W3,W2W_{1},W_{3},W_{2}在位置间共享，因此相同的算子作用于每个token。

### 2.2 VN–SST：记忆、状态空间与可编程计算

VN–SST保留了基线的层骨架，但替换了其两个模块：用三个并行记忆通路取代普通注意力，用一个可编程前馈算子取代固定的SwiGLU。这些选择是对此工作所依据的存储程序假设的直接解读。冯·诺依曼机分离出一个获取和解码指令的控制器、一个运行解码后算子的执行单元，以及一个在步骤间保持状态的记忆；计算是读取下一条指令并应用它的循环。VN–SST将这些角色映射到单个序列层：持久的选择性状态空间和快速权重通路是记忆，跨token携带状态；缓慢状态的低维读出充当指令指针，选择接下来运行哪个算子；控制器将该指针解码为每个token的指令代码；可编程的SwiGLU是执行单元，其权重由该代码实时合成。下文描述的三个通路和低秩指令库正是这些角色的具体实现。

每一层首先将其输入归一化，u=RMSNorm(h)u=\mathrm{RMSNorm}(h)，一个小控制器读取每个token的门控值——读取混合门gt∈Δ2g_{t}\in\Delta^{2}（对三个通路的单纯形权重）和记忆写入门wtΔ,wtM=σ(⋅)w^{\Delta}_{t},w^{M}_{t}=\sigma(\cdot)。三个通路随后并行运行。

#### 通路 1 —— 局部感觉缓冲区。

第一条通路是局限于宽度为w的因果带内的普通多头自注意力，

yloc=LocalAttnw(u),y^{\text{loc}}=\mathrm{LocalAttn}_{w}(u),\tag{3}

其计算成本仅为O(Tw)\mathcal{O}(Tw)，捕捉当前窗口内的短程结构，但不跨越窗口携带信息——那是后面两条持久通路的任务。

#### 通路 2 —— 选择性状态空间（慢动态）。

第二条通路是一个对角、输入依赖的SSM\[8 (https://arxiv.org/html/2608.25088#bib.bib8),9 (https://arxiv.org/html/2608.25088#bib.bib9)\]，它在片段间携带每个通道的状态s∈Rd×ns\in\mathbb{R}^{d\times n}（每个通道n个潜在状态）。输入投影x=Wxux=W_{x}u，每token步长Δt=wtΔsoftplus(WΔut+b)\Delta_{t}=w^{\Delta}_{t}\,\mathrm{softplus}(W_{\Delta}u_{t}+b)（写入门wtΔw^{\Delta}_{t}；学习的WΔW_{\Delta}和偏置bb），输入依赖的选择向量Bt,Ct∈RnB_{t},C_{t}\in\mathbb{R}^{n}，以及对角衰减A=−exp(Alog)A=-\exp(A_{\log})（学习的AlogA_{\log}），

Āt=exp(Δt⊙A),st=Āt⊙st−1+(Δtxt)⊗Bt,ytssm=⟨st,Ct⟩+D⊙xt.\displaystyle\bar{A}_{t}=\exp(\Delta_{t}\odot A),\quad s_{t}=\bar{A}_{t}\odot s_{t-1}+(\Delta_{t}x_{t})\otimes B_{t},\quad y^{\text{ssm}}_{t}=\langle s_{t},C_{t}\rangle+D\odot x_{t}.\tag{4}

这里⊙\odot是逐元素（Hadamard）积，⊗\otimes是外积，⟨·,·⟩\langle\cdot,\cdot\rangle在n个状态维度上缩并，DD是学习的逐通道跳跃连接。读出ytssmy^{\text{ssm}}_{t}是缓慢状态sts_{t}的低维投影；它正是我们在下面用作指令指针的信号。

#### 通路 3 —— 快速权重联想记忆（情景性）。

第三条通路是一个delta规则矩阵记忆M∈Rdk×dvM\in\mathbb{R}^{d_{k}\times d_{v}}（键和值维度dk,dvd_{k},d_{v}），同样在片段间携带，使用l2\ell_{2}归一化的键和查询kt,qt∈Rdkk_{t},q_{t}\in\mathbb{R}^{d_{k}}以及值vt∈Rdvv_{t}\in\mathbb{R}^{d_{v}}（写入门wtMw^{M}_{t}）。在每个步骤中，它将当前预测误差写入MM，并通过输出投影WoW_{o}按查询读回内容，

Mt=Mt−1+wtMkt(vt−kt⊤Mt−1)⊤,ytmem=Wo(qt⊤Mt).\displaystyle M_{t}=M_{t-1}+w^{M}_{t}\,k_{t}\big(v_{t}-k_{t}^{\top}M_{t-1}\big)^{\top},\quad y^{\text{mem}}_{t}=W_{o}\big(q_{t}^{\top}M_{t}\big).\tag{5}

写入误差而非原始值使召回具有内容可寻址性：一个类似于已存储键的键会检索其关联的值。

#### 融合。

三个读出通过控制器的读取门混合，并通过一个线性映射WfW_{f}加回残差流，

rt=gtlocytloc+gtssmytssm+gtmemytmem,h←h+Wfr.r_{t}=g_{t}^{\text{loc}}y^{\text{loc}}_{t}+g_{t}^{\text{ssm}}y^{\text{ssm}}_{t}+g_{t}^{\text{mem}}y^{\text{mem}}_{t},\qquad h\leftarrow h+W_{f}r.\tag{6}

因为状态(s,M)(s,M)跨片段持久，一个仅包含ww个token的窗口可以携带远比ww长的依赖关系。

#### 可编程计算：低秩指令库。

代替SwiGLU专家的软混合\[1 (https://arxiv.org/html/2608.25088#bib.bib1)\]，SwiGLU的两个投影都从共享的低秩库中逐token合成。对于通用投影W，写作，

W(t)=W0+∑k=1Kct,kUkVk⊤,ct∈RK,W(t)\;=\;W_{0}\;\+\\;\sum_{k=1}^{K}c_{t,k}\,U_{k}V_{k}^{\top},\qquad c_{t}\in\mathbb{R}^{K},\tag{7}

其中W0W_{0}是共享的基础算子，\{UkVk⊤\}\{U_{k}V_{k}^{\top}\}是KK个学习的秩–r“指令”，由每个token的码ctc_{t}选择。对于输入映射Win:Rd→R2dffW_{\text{in}}:\mathbb{R}^{d}\!\to\!\mathbb{R}^{2d_{\text{ff}}}（连接的门控/上投影），库包含Vk∈Rd×rV_{k}\in\mathbb{R}^{d\times r}和Uk∈R2dff×rU_{k}\in\mathbb{R}^{2d_{\text{ff}}\times r}；对于下映射Wdown:Rdff→RdW_{\text{down}}:\mathbb{R}^{d_{\text{ff}}}\!\to\!\mathbb{R}^{d}，第二个库包含Vkd∈Rdff×rV^{d}_{k}\in\mathbb{R}^{d_{\text{ff}}\times r}和Ukd∈Rd×rU^{d}_{k}\in\mathbb{R}^{d\times r}，两者都由相同的码ctc_{t}门控。关键的是，永远不会实例化token特定的权重：每个编程的投影都是一个基础映射加上一个低秩修正，

Win(t)ut=Win,0ut+∑k=1Kct,kUk(Vk⊤ut),W_{\text{in}}(t)\,u_{t}\;=\;W_{\text{in},0}\,u_{t}\;\+\\;\sum_{k=1}^{K}c_{t,k}\,U_{k}\big(V_{k}^{\top}u_{t}\big),\tag{8}

即，将utu_{t}投影到K×rK\times r库上，用码ct,kc_{t,k}缩放块kk，并通过UU读出。将这个2dff2d_{\text{ff}}维输出分成两个dffd_{\text{ff}}维半部\[gate;up\]=Win(t)ut[\,\text{gate};\text{up}\,]=W_{\text{in}}(t)\,u_{t}，然后隐藏激活ht=SiLU(gate)⊙uph_{t}=\mathrm{SiLU}(\text{gate})\odot\text{up}通过同样编程的Wdown(t)W_{\text{down}}(t)映射出去。用Uk,UkdU_{k},U^{d}_{k}初始化为零，该层从共享的基础SwiGLU精确开始。程序的每token成本为O(Kr(d+dff))\mathcal{O}\!\big(Kr(d+d_{\text{ff}})\big)，因此程序容量KK独立于基础FFN的O(ddff)\mathcal{O}(d\,d_{\text{ff}})计算进行缩放——每添加一条指令的成本约为完整MoE专家的∼r/dff\sim\!r/d_{\text{ff}}。

#### 指令指针：一个流形条件化的代码。

代码从token和携带的SSM状态的低维读出（公式4 (https://arxiv.org/html/2608.25088#S2.E4)）解码，

ct=tanh(MLP([ut;Pytssm]))∈[−1,1]K,P∈Rm×d,m≪d,c_{t}\;=\;\tanh\!\Big(\mathrm{MLP}\big([\,u_{t}\,;\,P\,y^{\text{ssm}}_{t}\,]\big)\Big)\;\in\;[-1,1]^{K},\qquad P\in\mathbb{R}^{m\times d},\ m\ll d,\tag{9}

因此，一个缓慢、低维的潜在轨迹(PyssmP\,y^{\text{ssm}})选择在流形上哪个算子在每个token执行——“获取–解码–执行”循环，SSM状态充当程序计数器。tanh\tanh将合成的算子保持在仿射流形的有界区域内。

### 2.3 记忆携带与训练

序列被处理为长度L的连续片段，并使用截断时间反向传播：状态(s,M)(s,M)跨片段传递，并每隔几个片段分离一次；Transformer运行相同的循环但是状态。