神经形态扩散语言模型：通过稀疏性和块去噪解决计算与内存瓶颈

# 神经形态扩散语言模型：通过稀疏性和块去噪解决计算与内存瓶颈
来源：https://arxiv.org/html/2607.24841
邓宇 Wu¹，Clement Ruah²，Jiechen Chen¹，Bipin Rajendran²，以及 Osvaldo Simeone² ¹伦敦国王学院工程系，英国伦敦 ²东北大学伦敦分校智能网络系统研究所 \(INSI\)，英国伦敦 电子邮件：\{dengyu.wu, jiechen.chen\}@kcl.ac.uk，\{c.ruah, b.rajendran, o.simeone\}@nulondon.ac.uk 这项工作得到了欧洲研究理事会 \(ERC\) 根据欧盟地平线欧洲计划（资助协议编号 101198347）、EPSRC 开放奖学金（EP/W024101/1, EP/X011356/1）以及 EPSRC 项目（EP/X011852/1）的资助。

###### 摘要

自回归 \(AR\) 大型语言模型 \(LLMs\) 在推理时固有地低效，因为每个生成的 token 都需要访问全部模型参数，导致运算强度低且能耗高。掩码扩散语言模型 \(MDLMs\) 通过允许每次参数访问生成多个 token，部分解决了内存受限环境下的这一局限性。为了在现代具备大规模片上内存的平台上进一步提升推理效率，本文提出了神经形态 MDLMs \(N-MDLMs\)，它将块扩散与基于尖峰的神经形态计算相结合，共同提高吞吐量和能效。块扩散通过每次参数访问生成多个 token 来提升 token 吞吐量，而尖峰诱导的稀疏性则通过跳过非活跃通道来减少有效的参数流量和计算。为了分析稀疏性与扩散的协同效应，我们开发了一个 token 级的类屋顶线模型，该模型捕捉了块并行生成和尖峰稀疏性对解码效率的综合影响。在翻译任务上的实验结果表明，得益于尖峰诱导的稀疏性，N-MDLMs 在计算受限平台（对于此类平台，MDLMs 无法比 AR-LLMs 有所改进）上也能实现能效和吞吐量的显著提升。

## I引言

大型语言模型 \(LLMs\) 越来越多地部署在延迟敏感和成本受限的环境中，在这些环境中，内存利用率、吞吐量和能耗等方面的推理效率决定了系统性能。然而，传统的*自回归* \(AR\) LLMs 需要重复访问全部模型参数，导致运算强度有限且吞吐量低\[bi2026rooflinebench, verhelst2025keep\]。

为了缓解这些限制，*掩码扩散语言模型* \(MDLMs\) 在每次参数访问时生成多个 token\[arriola2025block\]。然而，每个去噪步骤中使用 transformer 模型仍然限制了能效。此外，吞吐量的提升依赖于系统在内存受限模式下运行，这可能无法转化为具有大量片上内存的现代推理平台\[appuswamy2024breakthrough, abts2022software\]。

参考图注 图 1：在所提出的神经形态掩码扩散语言模型 \(N-MDLM\) 中，神经形态编码器处理干净的前缀 token，并产生神经形态上下文表示。神经形态解码器通过去噪步骤迭代地精炼一个掩码 token 块，同时重用编码器的神经形态表示。每个神经形态 transformer 层都使用 IF 神经元动力学实现（MLP = 多层感知机）。*神经形态计算* 原理最近被证明在降低标准 AR transformers 的能耗方面具有潜在优势\[zhou2022spikformer, simeone2026modern, pan2025spikingbrain\]。神经形态模型以事件驱动的方式运行，计算由事件（也称为尖峰）触发，这些事件代表离散的激活。在没有尖峰的情况下，定制的硬件部署可以避免计算和内存访问，从而受益于动态稀疏性\[simeone2026modern\]。

在这项工作中，如图 1 (https://arxiv.org/html/2607.24841#S1.F1) 所示，我们引入了*神经形态掩码扩散语言模型* \(N-MDLMs\)，这是一种通过并行去噪步骤提升吞吐量，同时通过神经形态计算提供的动态稀疏性进一步提高推理效率的模型类别。如图 2 (https://arxiv.org/html/2607.24841#S1.F2) 所示，尽管传统的 AR-LLMs 每次参数访问生成一个 token，但 MDLMs 将此成本摊销到块内的多个 token 上\[arriola2025block\]。N-MDLMs 通过事件驱动的稀疏性进一步增强运算强度，从而减少了参数流量和计算成本。即使在现代神经形态硬件\[davies2018loihi, appuswamy2024breakthrough\]和 LLM 推理加速器\[abts2022software\]（这些硬件将内存和计算在芯片内共置）上，这也可以将操作转向内存受限模式，使得通过扩散进行的并行生成既高效又节能。

参考图注 图 2：从 AR-LLM 到所提出的 N-MDLM：(a) AR-LLM 需要为每个生成的 token 访问全部模型参数。(b) MDLM 每次参数访问生成 B 个 token，每次访问的操作数按比例 B 增加。(c) 所提出的 N-MDLM 进一步利用尖峰稀疏性，使得只有活跃通道触发突触操作，从而减少了参数流量和计算成本。这项工作的主要贡献如下：

- •*神经形态 MDLM：* 我们引入了 N-MDLMs，这是一类新颖的神经形态 LLMs，通过基于量化和 integrate-and-fire \(IF\) 尖峰模型\[pan2025spikingbrain\] 对传统 MDLM\[arriola2025block\] 进行转换而获得。
- •*分析：* 我们开发了一个 token 级的类屋顶线分析框架，该框架表征了 N-MDLMs 作为块大小和激活稀疏性函数可实现的解码效率，捕捉了块并行生成和事件驱动计算在计算和内存访问上的联合协同效应。
- •*实验验证：* 我们在翻译任务上验证了所提出的框架，表明与 AR-LLMs 和传统 MDLMs 相比，N-MDLMs 实现了更高的 token 吞吐量和能效，同时保持了有竞争力的任务性能。这些增益被认为源于稀疏性（用于缓解计算限制带来的瓶颈）和扩散（用于缓解内存传输带来的吞吐量限制）的协同使用。

## II 预备知识

在本节中，我们首先回顾 AR-LLMs 和 MDLMs，以及 IF 神经元模型。

### II-A 自回归与掩码扩散语言模型对比

AR-LLMs 将 token 序列 \(x=(x_1,\ldots,x_L)\) 上的联合分布 \(p_\theta(x)\) 分解为

\[
p_\theta(x)=\prod_{l=1}^L p_\theta(x_l \mid x_{<l}), \tag{1}
\]

其中 \(\theta\) 表示模型参数，每个 token \(x_l\) 取值在词汇表 \(\mathcal{V}\) 中。在 (1) 中，条件分布 \(p_\theta(x_l|x_{<l})\) 在给定前缀 \(x_{<l}=(x_1,\ldots,x_{l-1})\)（由之前的 token 组成）的情况下为 token \(x_l\) 分配一个概率。如 (1) 中的分解导致了严格的顺序推理过程，使得在每个解码步骤中模型只产生一个 token \(x_l\)。相应地，如图 2 (https://arxiv.org/html/2607.24841#S1.F2)(a) 所示，每次参数访问产生一个输出 token，限制了运算强度和吞吐量。

MDLMs\[arriola2025block\] 通过将 token 分组为大小为 \(B\) 的块，并通过多个去噪步骤并行生成每个块（参见第 III-A 节）来放宽顺序操作。具体来说，序列 \(x=(x_1,\ldots,x_L)\) 被划分为块 \(x^b=(x_{b,1},\ldots,x_{b,B})\)，其中 \(b=1,\ldots,\lceil L/B \rceil\)，联合分布在块级别分解为

\[
p_\theta(x)=\prod_{b=1}^{\lceil L/B \rceil} p_\theta(x^b \mid x^{<b}), \tag{2}
\]

其中 \(x^{<b}=(x^1,\ldots,x^{b-1})\) 表示先前生成的块。

### II-B 神经形态计算原理

神经形态计算系统使用由神经元动力学产生的离散尖峰事件来表示信息。一种常用的神经元模型是带有软重置机制的 integrate-and-fire \(IF\) 神经元\[wu2025optimizing\]。对于任何神经元，当内部状态变量 \(V_t\) 超过阈值 \(V_\mathrm{th}\) 时，会产生一个对应于输出 \(s_t=1\) 的尖峰事件，如下所示

\[
s_t=
\begin{cases}
1, & V_t \ge V_\mathrm{th}, \\
0, & V_t < V_\mathrm{th}.
\end{cases} \tag{3}
\]

给定输入突触电流 \(I_t\)，状态变量演变为

\[
V_{t+1}=V_t+I_t-s_t V_\mathrm{th}. \tag{4}
\]

突触电流 \(I_t\) 是由来自突触前神经元的传入尖峰与其突触权重相互作用产生的。

在事件驱动执行下，突触操作仅在尖峰发生时触发。因此，非活跃神经元既不产生计算也不产生内存访问，导致计算和参数流量中的*动态稀疏性*\[davies2018loihi\]。

## III 神经形态掩码扩散模型

在本节中，我们介绍 N-MDLMs，这是一种通过转换现有基于常规编码器-解码器的 MDLMs\[pan2025spikingbrain\] 构建的神经形态块扩散 LLM。

### III-A 架构

如图 1 (https://arxiv.org/html/2607.24841#S1.F1) 所示，所提出的 N-MDLM 遵循块扩散生成范式，在该范式中，通过迭代去噪过程在块内预测多个 token。我们特别采用了来自\[arriola2025encoderdecoder\]的编码器-解码器架构。其中，编码器处理当前上下文以产生表示，该表示在去噪过程中作为解码器的条件信息。

在基于分解 (2) 的编码阶段 \(b\)，模型在给定上下文 \(x^{<b}=(x^1,\ldots,x^{b-1})\) 的情况下生成块 \(x^b\)。编码器 \(f_{\theta_\mathrm{enc}}(\cdot)\) 首先处理上下文 \(x^{<b}\) 以产生潜在表示

\[
h^b = f_{\theta_\mathrm{enc}}(x^{<b}). \tag{5}
\]

使用潜在表示 \(h^b\)，解码器的生成从初始掩码 token 表示 \(z_0^b \in \mathbb{R}^{B \times D}\) 开始，其中 \(D\) 表示 transformer 的隐藏维度。矩阵 \(z_0^b\) 按列包含块中所有 \(B\) 个 token 的特殊掩码 token 的嵌入。这个初始的非信息性信号通过 \(S\) 个去噪步骤（索引为 \(s=1,\ldots,S\)）进行迭代精炼。在每个去噪步骤 \(s\) 中，选择 \(B/S\) 个 token（假设 \(B/S\) 为整数）用于剩余掩码位置以进行解掩码\[arriola2025block\]。为此，transformer 解码器 \(f_{\theta_\mathrm{dec}}(\cdot)\) 映射上下文表示 \(h^b\) 和当前噪声块 \(z_s^b=(z_{s,1}^b,\ldots,z_{s,B}^b)\) 以产生更新的 token 嵌入

\[
z_{s+1}^b = f_{\theta_\mathrm{dec}}(z_s^b, h^b) \in \mathbb{R}^{B \times D}, \tag{6}
\]

与 \(z_s^b\) 相比，有 \(B/S\) 个额外的 token 被解掩码。

### III-B 转换

从预训练的编码器-解码器 MDLM\[pan2025spikingbrain\] 开始，我们通过基于量化的转换方法获得 N-MDLM，该方法利用 IF 神经元模型 (4)\[pan2025spikingbrain\]。这种方法的主要原理是，IF 神经元或其变体的尖峰速率可以与量化函数的输出相匹配，其中量化步长等于 (3) 中的阈值 \(V_\mathrm{th}\)\[pan2025spikingbrain\]。

给定由编码器或解码器层产生的激活向量 \(\boldsymbol{a}=(a_1,\ldots,a_D)\)，具有自适应步长的量化函数获得量化整数值向量 \(\hat{\boldsymbol{a}} \in \mathbb{Z}^D\) 如下\[pan2025spikingbrain\]

\[
\hat{\boldsymbol{a}} = \mathrm{round}\left(\frac{\boldsymbol{a}}{V_\mathrm{th}(\boldsymbol{a})}\right), \qquad V_\mathrm{th}(\boldsymbol{a}) = \frac{1}{K} \mathrm{mean}(|\boldsymbol{a}|), \tag{7}
\]

其中取整函数 \(\mathrm{round}(\cdot)\) 和幅度函数 \(|\cdot|\) 是逐元素应用的。(7) 中的参数 \(K\) 控制量化级别的数量，\(K\) 值越大，分辨率越高。

每个量化值 \(\hat{a}_i\) 通过比特位编码表示，这是一种在 \(T\) 个时间步上使用三元尖峰 \(s_{i,t} \in \{-1,0,1\}\) 的时间编码方案。尖峰与整数 \(\hat{a}_i\) 的关系为

\[
\hat{a}_i = \sum_{t=1}^T 2^{t-1} s_{i,t}, \tag{8}
\]

其中 \(T = \left\lfloor \log_2\left(\max_i |\hat{a}_i|\right) \right\rfloor\) 表示时间步数。

设 \(\boldsymbol{w}=[w_1,\ldots,w_D]^\top \in \mathbb{R}^D\) 为一个权重向量。这可以对应线性层的行或 MLP 中的线性投影。量化激活向量 \(\hat{\boldsymbol{a}}\) 的内积 \(\boldsymbol{w}^\top \hat{\boldsymbol{a}}\) 可以使用一个改进的 IF 神经元在作用于 (8) 中的尖峰信号时进行评估。为此，设 \(I_t = \sum_{i=1}^D w_i s_{i,t}\) 为时间步 \(t\) 时 IF 神经元的输入电流，并考虑具有以下动力学的改进 IF 神经元

\[
V_{t+1} = V_t + 2^t I_t - 2^t \hat{s}_t V_\mathrm{th}(\boldsymbol{a}), \tag{9}
\]

其中 \(\hat{s}_t\) 表示发射的尖峰，计算为

\[
\hat{s}_t =
\begin{cases}
1, & V_t \ge 2^t V_\mathrm{th}(\boldsymbol{a}), \\
-1, & V_t \le -2^t V_\mathrm{th}(\boldsymbol{a}), \\
0, & \text{otherwise}.
\end{cases} \tag{10}
\]

因此，线性变换可以近似重构为\[pan2025spikingbrain\]

\[
\boldsymbol{w}^\top \hat{\boldsymbol{a}} \approx V_\mathrm{th}(\boldsymbol{a}) \sum_{t=1}^T 2^{t-1} \hat{s}_{i,t}. \tag{11}
\]

## IV 性能分析

与传统的 AR-LLMs 相比，N-MDLMs 可以通过选择块大小 \(B\)（摊销每次块的参数访问）和量化超参数 \(K\)（控制生成尖峰的稀疏性）来调节数据访问和计算。在本节中，我们在一个受屋顶线启发的框架\[verhelst2025keep\]中，研究块大小 \(B\) 和稀疏性超参数 \(K\) 对每个 token 的操作数和内存传输需求的影响，从而得出估计的 token 吞吐量。

### IV-A 计算与内存传输