企业官网建设流程全解析

脉冲神经网络和 Transformer 摆在一起，很多人第一反应是“不搭”：一个靠稀疏的 0/1 事件流省电，一个靠稠密浮点矩阵乘堆表达力，底层算子的习惯完全相反。2023 年那篇 Spike-driven Transformer 就是把这两套系统结合起来的成果，切入点很巧妙——把自注意力里的乘法换成加法。说白了，脉冲驱动 Transformer 的核心不是发明新算子，而是发现了一个等价变形：当 Q、K、V 三路都被压成 0/1 脉冲之后，(QK^T)V这个乘法主体可以改写成Q(K^T V)，而后者在两个二值矩阵相乘时，乘法本身就退化成“掩码累加”，也就是纯加法。这件事的价值在于，SNN 在专用硬件上真正省电的地方从来不是“网络小”，而是“没有乘法器被点亮”。它解决的问题也很明确：让 Transformer 这种高表达力结构能以事件驱动的方式跑起来，而不是只能跑在通用 GPU 上做稠密计算。适合谁来读？做过 SNN 训练、看过替代梯度、被脉冲不收敛折磨过的同学，以及想把注意力机制的能耗账算清楚、想复现一份能跑通的 SDSA（Spike-driven Self-Attention）代码的同学。下面我按“为什么这么设计—算子怎么算—配套模块怎么改—怎么跑起来—哪里会出问题”的顺序，把这份代码里我认为最值得抠的部分拆开讲。

1. 从乘法到加法：这次改造到底在省什么

1.1 脉冲神经网络的三条硬约束，决定了它天生排斥乘法

先把前提说清楚，不然下面所有设计看起来都像“为了炫技而故意绕路”。SNN 在工程上真正被看重的特性只有三条，而这三条每一条都在跟乘法作对。第一条是信息编码是二值的，神经元在某一时刻要么发放要么不发放，中间没有 0.37 这种值；第二条是计算是事件驱动的，没有脉冲到达就不做累加，所以空闲的神经元能耗接近零；第三条是硬件上乘法器的代价远高于加法器，业内常被引用的 45nm CMOS 数据里，一次 32 位浮点乘加大约是 3.7pJ 量级，而一次加法只有 0.9pJ 量级，差出接近四倍。

这三条约束叠在一起，结论就很清楚了：如果我只是把 Transformer 的权重换成低比特，然后继续在 GPU 上跑稠密矩阵乘，那 SNN 的省电优势一点都拿不到——GPU 里根本没有“跳过 0”这回事。真正想吃到红利，必须让计算图里乘法出现的次数本身变少，而不是让乘法变得“看起来像加法”。所以原始工作选的路子不是量化、不是剪枝，而是重新排列计算顺序，让乘法从结构上消失。这个选择背后的判断是：在 SNN 场景里，结构级的算子替换比数值级的近似更划算，因为它不引入额外的精度损失源。

我见过不少同学一上手就想“用低比特注意力加速”，结果模型精度掉了两个点，速度一点没变。原因就在这里——你以为的省算力，在通用硬件上不成立；而在专用硬件上，二值矩阵乘本来就不需要乘法器。所以后面所有的代码细节，你都要带着“这是为了配合事件驱动硬件”的视角去读，而不是“这是为了在 A100 上快 20%”。这个预期一旦摆正，很多看起来奇怪的写法就顺了。

1.2 乘法能被折叠掉的前提：结合律 + 二值性 + 去掉 softmax

这一步是整个工作的数学核心，我拆成三层讲。第一层是结合律。标准自注意力算的是softmax(QK^T)V，中间的QK^T是一个 N×N 的矩阵，计算量是 O(N²D)。但矩阵乘法满足结合律，(QK^T)V和Q(K^T V)在数值上是同一件事，而后者的计算量是 O(ND²)。这个技巧在线性注意力里早就有人用了，本身不新鲜。

第二层是二值性带来的算子退化。这一点才是关键。当 Q、K、V 都是 0/1 矩阵时，K^T V里的每一次乘加，本质上是“如果 k 的某一位是 1，就把 v 对应的那一行加进来；如果是 0，就跳过”。也就是说，乘法变成了条件累加。同理Q(K^T V)这一步，Q 是 0/1，就是在一堆行向量里做选择再加起来。整个注意力块里，没有一处需要真正撞击乘法器。这就是“加法 Transformer”这个名字的由来——它不是说把乘换成加号，而是说乘法在这套数据结构下压根不成立为运算。

第三层，也是最容易被忽略的一层：softmax 必须去掉。因为一旦有 softmax，softmax(QK^T)V就不等于Q(K^T V)了，归一化是非线性的，你没法把 V 提前乘进去。原始工作的处理方式是直接用脉冲计数做注意力权重，用一个固定的缩放系数压住幅值，替代 softmax 的归一化作用。这个取舍在视觉分类任务上是成立的，因为图像上没有因果性要求；但如果换成自回归的序列任务，K^T V会把未来时间步的信息混进来，这个变形直接失效。

提示：这套改法只适用于非因果的全局注意力。任何带 causal mask 的自回归场景，都不要直接套用这个顺序变换，否则你会得到“训练 loss 异常低、推理结果全乱”的诡异现象。

1.3 整体结构长什么样：四个模块的职责划分

从代码组织上看，一个 Spike-driven Transformer block 和我熟悉的 ViT block 骨架几乎一模一样，都是“注意力分支 + MLP 分支 + 两条残差”。差别全在细节里，我列个清单方便对照。

这里我要特别强调残差那条。很多人第一次写会想“既然全网络都是脉冲，残差也应该在脉冲域做加法吧”，结果就是精度崩掉。原因很直白：两个 0/1 向量相加得到的是 0/1/2，你再发射一次就是“或”运算，信息被严重压缩，深层的梯度几乎传不回去。原始工作的做法是让残差留在浮点域，也就是膜电位这一层做累加，只在进入下一个算子前才发射脉冲。这个设计不是可选项，是必须项。

2. 核心算子逐行拆解：Spike-driven Self-Attention

2.1 先看计算量与能耗对比，才知道该在哪里优化

在动手写代码之前，我习惯先把计算量算清楚，不然很容易在错误的地方做优化。假设输入是 N 个 token、通道维度 D、头数 h、每头维度 Dh = D/h，那么三种方案的复杂度对比如下。

表面上看，脉冲驱动版本和线性注意力的复杂度完全一样，那“省”在哪里？省在常数因子上：线性注意力里那两个矩阵乘仍然是实打实的浮点乘加，而脉冲版本因为操作数是 0/1，在事件驱动硬件上可以被编译成“按索引收集再累加”，乘法器全程不亮。这个差别在 GPU 上测不出来，但在能耗模型里是量级差异。

不过这里有个反直觉的点，我踩过：O(ND²) 不一定比 O(N²D) 快。两者的比值是 D/N。浅层特征图大、通道窄，比如 56×56 的特征图 N=3136、D=32，此时 D/N ≈ 0.01，收益是千倍级的；但到了深层，特征图缩到 7×7、N=49，通道涨到 384，D/N ≈ 7.8，也就是说这个“优化”反而比原始写法贵了近八倍。原始工作的模型在每一层都用了同样的注意力形式，但真正吃到红利的只有前两个 stage。这个事实在论文正文里不会被强调，但你复现的时候必须知道，否则会困惑“为什么我按论文写了却更慢”。

2.2 三路脉冲化：Q、K、V 必须在进入注意力前变成 0/1

这是整个实现里第一个必须做对的细节。很多人写完线性层就直接送进注意力，忘了发射脉冲，结果得到的是一个“带 BN 的普通线性注意力”，跑出来的精度甚至还不错，于是误以为复现成功了。判断方法很简单：在注意力入口打印一下q.unique()，如果里面出现了 0 和 1 之外的值，就是漏了。

import torch import torch.nn as nn class SurrogateSpike(torch.autograd.Function): """前向二值化，反向用矩形窗替代梯度。""" @staticmethod def forward(ctx, mem): ctx.save_for_backward(mem) return (mem >= 0).to(mem.dtype) @staticmethod def backward(ctx, grad_out): (mem,) = ctx.saved_tensors # 窗宽 0.5：只有膜电位落在阈值附近，梯度才允许通过 return grad_out * (mem.abs() <= 0.5).to(mem.dtype) class SpikeAct(nn.Module): def forward(self, x): return SurrogateSpike.apply(x)

为什么阈值取 0 而不是某个正数？因为前面接了 BatchNorm，BN 输出的分布已经被拉到零均值附近，阈值 0 就等价于“以 0 为判决边界”，这样平均发放率天然接近 50%，再通过 BN 的可学习偏移去调节。如果你跳过 BN 直接用阈值 1.0，那就要自己保证膜电位的量纲，很容易出现整层全 0 或全 1。

替代梯度的窗宽也值得说一句。矩形窗取 0.5 是常见起点，太窄（比如 0.1）会导致梯度只在极小区间内非零，深层几乎学不动；太宽（比如 1.0 以上）会让反向传播的梯度近似于恒等映射，模型确实能训起来，但你已经不是在训 SNN 了，精度虚高、部署时掉点。我的经验是先按 0.5 起步，如果发现 loss 前 10 个 epoch 几乎不动，再往 0.8 调；如果发现训练精度很高但验证集抖得厉害，往 0.3 收。

2.3 计算顺序换一换，乘法就从图里消失了

下面这段是整个仓库里最值得逐行看的代码。我把它写成可以直接跑的形式，注释里标了每一步的意图。

class SDSA(nn.Module): """Spike-driven Self-Attention""" def __init__(self, dim, num_heads=8): super().__init__() assert dim % num_heads == 0 self.h = num_heads self.dh = dim // num_heads self.scale = self.dh ** -0.5 # 压住累加幅值 self.q = nn.Linear(dim, dim, bias=False) self.k = nn.Linear(dim, dim, bias=False) self.v = nn.Linear(dim, dim, bias=False) self.proj = nn.Linear(dim, dim) self.bn_q = nn.BatchNorm1d(dim) self.bn_k = nn.BatchNorm1d(dim) self.bn_v = nn.BatchNorm1d(dim) self.act = SpikeAct() def _to_spike(self, x, bn): # x: [TB, N, C] -> BN 在通道维 -> 再发射 x = bn(x.transpose(1, 2)).transpose(1, 2) return self.act(x) def forward(self, x): # x: [T, B, N, C]，本身已经是脉冲 T, B, N, C = x.shape x = x.flatten(0, 1) # [T*B, N, C] q = self._to_spike(self.q(x), self.bn_q) k = self._to_spike(self.k(x), self.bn_k) v = self._to_spike(self.v(x), self.bn_v) def heads(t): return t.view(T * B, N, self.h, self.dh).permute(0, 2, 1, 3) q, k, v = heads(q), heads(k), heads(v) # [TB, h, N, dh] # 关键：先 K^T V，再让 Q 去"取" kv = k.transpose(-2, -1) @ v # [TB, h, dh, dh] out = q @ kv # [TB, h, N, dh] out = out * self.scale out = out.permute(0, 2, 1, 3).reshape(T * B, N, C) out = self.proj(out) return out.view(T, B, N, C)

注意kv = k.transpose(-2, -1) @ v这一行，形状是[TB, h, dh, dh]，它和 N 完全无关。这意味着注意力的中间结果不再随序列长度平方增长，显存占用这一块就已经省下来了，尤其是在浅层大分辨率特征图上，这个差别非常明显。

那“乘法消失”体现在哪？在 PyTorch 里你写@，底层当然还是调 GEMM。但如果你按 0/1 的语义去看，这行代码等价于下面这种索引写法——而索引写法就是事件驱动硬件上实际编译出来的东西。

# 语义等价示意：乘法退化为按脉冲索引的选择性累加 kv = torch.zeros(T * B, self.h, self.dh, self.dh, device=x.device) for i in range(N): idx = k[:, :, i, :].bool() # 该 token 哪些通道发了脉冲 kv[idx] += v[:, :, i, :][idx] # 只做加法，没有乘法

实际部署时当然不会真的写 Python 循环，硬件上是把k的 1 位置翻译成地址，把v的对应元素送进累加器。你只要理解这一层对应关系，就能看懂为什么论文敢说“注意力里乘法数量为零”。

还有一个绕不过去的细节：缩放系数。K^T V的结果是计数，最大值上界是 N，随着 N 和 head_dim 增大，这个计数会线性膨胀，不做缩放的话，后面接 BN 时方差会随分辨率漂移。原始工作里用的是一个固定常数（1.5 量级）来压幅值，工程上换成head_dim ** -0.5也能跑通，但这两个选择在大分辨率输入上的表现不一样，我建议按你的实际输入尺寸在验证集上二选一，不要两个混着用。

2.4 一个必须提前想清楚的问题：因果性被破坏了

前面提到 softmax 被去掉，但还有一层更隐蔽的东西——Q(K^T V)这个顺序变换把时间因果性彻底抹平了。在(QK^T)V里，第 i 个 token 的输出只和 0..i 的 token 有关（如果加了 mask）；但换成Q(K^T V)之后，K^T V是一个全局聚合量，第 i 个 token 拿到的信息里必然混入了 i 之后的内容。

对图像分类任务这无所谓，因为图像上没有“未来”这个概念，全局注意力本来就是全连通的，这也是原始工作选择视觉任务作为主战场的原因之一。但如果你把这套代码搬到语音流、事件相机时序流、或者任何需要在线推理的场景里，就会出问题：离线训练时一切正常，上线做逐帧推理时结果对不上。我的处理办法是在 SDSA 外面加一层时间掩码，把 T 个时间步拆成独立的注意力计算，只在通道维度共享权重，不跨时间做聚合。代价是模型丢失了跨时间步的长程建模能力，需要靠堆层数或者加一个轻量的时间聚合模块来补。

注意：不要试图给K^T V加一个下三角 mask 来恢复因果性，因为K^T V已经把 N 维压掉了，你没法在 dh×dh 的矩阵上表达“第 i 个 token 只能看前 i 个”这种约束。因果性只能靠拆时间步来保，不能靠 mask 来救。

3. 与注意力配套的其它模块：这些地方比注意力更容易翻车

3.1 MLP 分支：降采样为什么只能用池化

先说一个我在复现时踩得最深的坑——降采样不能用 stride=2 的卷积。标准的 ViT 或者 ConvNeXt 里，stage 之间用 stride-2 卷积下采样，卷积核在特征图上滑动，每个输出位置都会覆盖 2×2 的输入区域。在浮点网络里这没问题，因为输出几乎不可能是精确的 0。但在脉冲网络里，输入是 0/1，如果某个 2×2 窗口里恰好全 0（在浅层高分辨率、低发放率的情况下这非常常见），卷积输出就是 0，再经过 BN 和脉冲发射，这个神经元在后续所有时间步里都可能再也不发火——这就是业内说的“神经元死亡”。

原始工作的解法是用 max-pool 做降采样。max-pool 的语义是“窗口内只要有一个 1，输出就是 1”，它保证了信息不会被完全抹掉。这个改动看起来很小，但它直接决定了深层能不能训起来。我在一个 4 stage 的模型上做过对比：用 stride-2 卷积的版本，跑到第 3 个 stage 时整层发放率掉到 0.001 以下，loss 卡住不动；换成 max-pool 之后发放率稳定在 0.2 到 0.4 之间，训练正常收敛。

MLP 分支本身的结构是Linear → BN → Spike → Linear → BN → Spike，和注意力分支保持一致的范式。这里有个细节：两个 Linear 之间不能加 Dropout。Dropout 会把一部分激活直接置零，在浮点网络里这是正则化，在脉冲网络里这是人为制造“不发火”，BN 的统计量会跟着漂移，训练和推理的发放率对不上。要用正则化的话，用 weight decay 或者 stochastic depth 更安全。

3.2 归一化：BN 必须在脉冲发射之前，而且不能在脉冲之后补

这条看起来是常识，但我见过至少三个复现仓库把它写反了。正确顺序是Linear → BN → Spike，错误顺序是Linear → Spike → BN。为什么错了？因为脉冲之后张量只有 0 和 1，你在这个分布上做 BN，等于把“发放率”强行拉到零均值，而 BN 的输出是浮点，接着又需要一次脉冲发射才能维持二值性——你就多了一次发射，信息被砍了两刀。

还有一个更隐蔽的问题：推理阶段的 BN 用的是 running stats，如果训练时 T 维和 B 维混在一起统计，而推理时 batch=1，统计量会对不上。原始代码里 BN 的输入是把 T 和 B 展平成一个维度（flatten(0, 1)），所以BatchNorm1d实际是在 T×B 个样本上统计的。推理时如果你只喂一帧（T=1, B=1），running stats 虽然还在，但和训练时的分布差异会放大。我的做法是推理时也保持 T=4 一起喂，再做时间维平均，这样和训练分布一致，实测分类精度能稳住。

至于为什么不用 LayerNorm——LN 是在每个 token 内部做归一化，对 SNN 来说，膜电位的幅值信息被 LN 抹掉之后，脉冲发射就接近于随机的伯努利采样，训练会非常不稳定。BN 因为跨样本统计，保留了样本间的相对幅值关系，在 SNN 社区里是更主流的选择。如果你的 batch size 被迫很小（比如显存不够），可以考虑 GroupNorm，但要注意它的统计维度和 BN 不同，阈值需要重新调。

3.3 时间维怎么处理：直接复制还是加时间注意力

静态图像做 SNN 训练时，最常见的做法是把一张图复制 T 次（T 一般取 4）当成 T 个时间步喂进去。这个做法简单，但它带来一个必须想清楚的问题：这 T 帧之间的注意力是独立算的，还是联合算的？

原始代码的做法是把 T 和 B 展平，也就是每一帧独立做空间注意力，帧与帧之间只通过膜电位状态（LIF 的 V）传递信息。这样做的好处是显存友好、实现简单；坏处是模型没有显式的跨时间建模能力，T 增大带来的收益很快饱和。我实测下来，T 从 2 加到 4 有稳定提升（大概 0.8 到 1.2 个点），从 4 加到 8 基本持平，再往上甚至会掉——因为梯度要穿过更多的时间步，替代梯度的误差累积起来了。

# 典型的静态图训练写法 T = 4 for img, label in loader: x = img.unsqueeze(0).repeat(T, 1, 1, 1, 1) # [T, B, C, H, W] out = model(x) # [T, B, num_classes] logits = out.mean(0) # 时间维平均 loss = criterion(logits, label)

如果你的任务本身就是事件流（比如事件相机采集的数据），那 T 帧之间天然有真实的时间演化，这时候更值得做的是让不同时间步共享注意力权重但各自独立计算，再用一个轻量的时间聚合（比如可学习的加权和）来融合。我试过直接在时间维上做注意力，参数量涨了不少，收益不明显，性价比不高。

4. 完整复现流程与参数选择

4.1 环境与数据管线的最小可用配置

环境上没什么特殊要求，PyTorch 2.x 加一个常规的视觉训练栈就够了，不需要额外的 SNN 框架——这套代码里的脉冲逻辑完全可以手写，用自定义 autograd Function 就行。我建议不要一上来就上大规模数据集，先在 32×32 的小数据集上把流程跑通，确认发放率、梯度、loss 曲线都正常，再往大分辨率迁移。

这里有一条我强烈建议：先把 T 设成 1 跑一遍。T=1 时模型退化成一个“带二值激活的 Transformer”，如果这个版本都训不收敛，那问题一定在 BN 位置、替代梯度窗宽或者学习率上，和时序机制无关。排除掉这些之后再把 T 加到 4，问题范围就小得多。这个“降维排查”的习惯帮我省了大量时间。

4.2 关键超参怎么定：给几个能直接用的起点

超参我按重要程度排了个序，附上选取理由，你可以直接抄，但要按自己的数据规模微调。

关于头数，我要多说一句。head_dim太小（比如小于 16）的时候，K^T V就是一个很小的计数矩阵，表达能力受限，而且缩放系数会变得很大，数值上不稳。我一般保证head_dim不低于 32。如果模型总维度不够，宁可减少头数，也不要硬切。

关于学习率还有一个经验：SNN 的训练前期，如果学习率给大了，会出现“整层脉冲全 1”的塌缩现象——所有神经元都饱和发放，BN 的方差趋近于零，梯度消失。判断方法是看第一层 BN 的 running_var，如果它掉到 1e-4 以下，基本就是塌了。遇到这种情况，把学习率砍一半，warmup 拉长到 30 个 epoch，通常能救回来。

4.3 训练循环与监控指标：光看 loss 是不够的

训练脚本本身不复杂，但监控指标必须加，否则你没法定位问题。

def train_one_epoch(model, loader, opt, crit, T=4): model.train() for img, label in loader: img = img.unsqueeze(0).repeat(T, 1, 1, 1, 1) out = model(img) # [T, B, C] loss = crit(out.mean(0), label) opt.zero_grad() loss.backward() # 梯度裁剪对 SNN 几乎是必须的 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step()

必须盯的四个指标，我按优先级列一下：每层的平均发放率（健康区间 0.1 到 0.5，低于 0.05 说明快死了，高于 0.8 说明饱和）、替代梯度的非零比例（太低说明窗宽不合适）、各层 BN 的 running_var（判断是否塌缩）、梯度范数（SNN 的梯度范数波动比 ANN 大得多，突然飙升通常是某层发放率跳变）。

我自己写了个小 hook，在每个 SpikeAct 后面统计x.mean()，训练时每 100 个 iteration 打印一次。这个习惯的价值在于，当 loss 卡住时，你能立刻判断是“整体学不动”还是“某几层死了”，而不是盲目调学习率。

4.4 能耗怎么估：别拿 GPU 的时钟当指标

如果你要对外说“我这套方案省了多少能耗”，千万别用 GPU 上的推理时间去换算，那只能反映稠密计算的效率，和 SNN 的能耗模型是两回事。合理的估算方式是算子计数 + 单次算子能耗：统计一次前向里乘法和加法的次数，乘法按 3.7pJ、加法按 0.9pJ 折算，再乘上脉冲稀疏度带来的跳过比例。

举个具体的算法。假设某一层 head_dim=32、token 数 N=196、头数 8，那么K^T V的加法是h × dh × dh × N = 8 × 32 × 32 × 196 ≈ 1.6M次，Q(K^T V)同量级，合计约 3.2M 次加法。如果换成标准注意力，QK^T是h × N × N × dh = 8 × 196 × 196 × 32 ≈ 9.8M次乘加。只算这一层，算子数量差了三倍，再乘上“乘法 vs 加法的能耗差”，总差距在十倍量级。这个账算出来才有说服力，而且它和你在 GPU 上测到的耗时完全可能相反——我就遇到过“算子数少了三倍但 GPU 上慢了 40%”的情况，原因就是形状不规整导致 GEMM 效率低。

5. 常见问题与排查实录

5.1 训练不收敛：先按这个顺序查

我把踩过的问题整理成了一张速查表，按“出现频率 × 排查成本”排序，从上往下查效率最高。

这里我想单独说一下“loss 完全不动”这个症状。它最常见的原因不是学习率，而是脉冲发射的位置写错了，导致梯度根本传不到前面的层。验证方法很直接：在模型最前面加一个 hook，打印第一层权重的梯度范数，如果恒等于 0，那就是替代梯度没生效——大概率是自定义 Function 的backward里忘了乘grad_out，或者条件判断写反了。这个 bug 我用半天时间才定位到，因为它不报错，只是安静地不学习。

5.2 精度掉点：按这个归因顺序排查

如果你的模型能训起来，但比同规模的浮点基线低了五六个点，别急着改结构，按下面的顺序归因，通常能找回大部分差距。

第一顺位是时间步 T。T=1 的时候，模型其实是个二值激活网络，表达力天然弱于 ReLU 网络，掉点 1 到 2 个点是正常的；T 加到 4 之后，靠时间维的冗余能把这块补回来。第二顺位是训练轮数。SNN 的收敛曲线比 ANN 平缓得多，同样的数据集，ANN 训 100 轮收敛，SNN 可能要 300 轮，早期停训是掉点大户。第三顺位是数据增强的强度。RandAugment 在脉冲网络上的效果不太稳定，我一般把强度调低一档。

第四顺位才是结构上的取舍，比如去掉 softmax 带来的表达力损失。这一块确实存在，原始工作也是靠更深的网络和更大的 T 来补的，不是靠算子层面能解决的。所以如果你的目标是“精度追平浮点基线”，那大概率要失望；如果你的目标是“在给定能耗预算下拿到最高的精度”，这才是这套方案的正确评价方式。

5.3 速度没变快：先确认你的瓶颈在哪

这个我在 4.4 节提过，但值得单独展开，因为它是复现者最容易产生“被论文骗了”的情绪的地方。

第一，GPU 上不会变快。0/1 矩阵在 GPU 里还是被当成 fp32 稠密矩阵做 GEMM，乘法器照样点亮，你只是换了个写法。真正变快需要事件驱动的加速器，或者至少用稀疏算子库把 0 跳过。第二，形状不匹配会让 GEMM 效率暴跌。dh×dh的小矩阵乘在 GPU 上是非常不友好的形状，尤其是 dh 只有 32 的时候，Tensor Core 的利用率很低，实测耗时可能比标准注意力还高。第三，BN 的 transpose 会带来非连续内存。代码里bn(x.transpose(1,2)).transpose(1,2)这个写法，两次 transpose 都会改变 stride，后续算子要么触发隐式拷贝，要么走慢路径。我一般会在 BN 之后加一次.contiguous()，虽然多一次拷贝，但整体反而更快。

如果你的目标就是想在自己的设备上测出速度收益，那正确的做法是关掉所有 BN 的 transpose，把通道维放到最后（NHWC 布局），然后看看能不能用上稀疏矩阵库。这条路我试过，收益不稳定，取决于你的硬件和稀疏度，不建议作为主要优化方向。

5.4 一个容易被忽略的坑：推理时的时间步对齐

最后分享一个部署阶段才会暴露的问题。训练时 T=4 一起喂，推理时如果为了降低延迟只喂 T=1，会出现精度断崖式下跌——不是缓慢下降，是直接掉十来个点。原因有两层：一是 BN 的 running stats 是在 T×B 上统计的，单帧输入的分布对不上；二是膜电位在 T 个时间步之间有累积过程，只跑一步相当于只用了 1/4 的信息量。

解决办法有两个方向。一是推理时也保持完整的 T 步，代价是延迟变成 T 倍，适合对延迟不敏感的场景。二是做时间维的权重融合，把 T 步的膜电位累积近似折算成一个等效的阈值偏移，这样单步推理也能接近多步的精度。第二种做法需要重新校准阈值，工作量不小，我建议除非有硬性延迟要求，否则优先选第一种。

我在实际复现这套代码时最深的体会是，这套方案的价值不在“更快”，而在“更省”，而且这个“省”必须放到正确的硬件模型和正确的评价口径下才成立。如果你拿 GPU 的 wall-clock 时间去衡量它，会得出完全相反的结论；但如果你按算子能耗去算，十倍级的差距是实打实的。另外一个经验是，SNN 的调试和 ANN 完全是两种手感——ANN 出问题你调学习率、调正则，SNN 出问题你第一件事应该是去看发放率和梯度范数，这两个指标正常了，剩下的才是超参问题。最后再提一句，这套注意力形式在浅层大分辨率上的收益最大，深层其实可以考虑换回标准注意力或者干脆用更轻的结构，没必要为了“全篇统一”而在 D/N 很大的层上硬套，那样只会白白增加算力。