把“不定期读一篇Paper”这个系列捡起来的时候我第一个想重读的就是Non local Neural Network。这篇论文在我看来是视觉注意力机制的承上启下之作往上它能接住传统手工特征里非局部均值滤波的思路往下它几乎直接孕生了后来席卷各个领域的Self-Attention结构。如果你正打算理解全局建模、或者想在自己的网络里加入长距离依赖这篇绝对绕不开。本篇文章不谈鸡汤只拆公式、看实验、聊工程里那些论文没写明白的细节。我习惯每篇Paper都带着三个问题去读它解决了什么它的方案比之前强在哪把它搬到我自己场景里会踩什么坑下面围绕这三个问题展开。1. 为什么这篇Paper值得专门拆开读1.1 局部卷积的天生缺陷感受野往上堆全局依赖却仍旧“够不着”做视觉任务的朋友对“感受野”这个概念都不陌生。卷积一层层叠下去理论上高层特征能覆盖全图但实际操作中一个3x3的卷积核在深层能看到的区域并不“均匀”它更多是局部的堆叠而且随着网络加深梯度穿过这么多层之后远距离信息早被稀释了。如果你做过视频任务体会会更明显两个在空间上离得很远的物体比如远处的人和近处的水花在时间上又有先后关系靠堆卷积去建模这种依赖简直像让一个人通过一条条短信描述来拼出全貌中间漏一个字都会偏。Non local Neural Network想做的事情就是给网络一个直接“看全部”的通道计算某个位置输出时不只参考周围邻居而是参考整个输入特征图上所有位置并且根据每个位置与当前点的相关程度自动分配权重。这样一来长距离依赖的传播路径从“逐层接力”变成了“一步直达”。论文把这个操作叫做non-local是相对于标准卷积“只考虑局部邻域”而言的。更关键的是这种全局交互不是靠增加卷积核尺寸硬堆出来的。增加卷积核尺寸确实能扩大局部窗口但计算量随核大小平方增长而且对视频里跨时间帧的依赖依然无能为力。Non-local把目光从“核大小”转移到了“位置间关系”这其实是一次建模思路的转变。1.2 非局部思想的渊源从图像去噪里的Non-local Means说起“non-local”这个词其实不是这篇论文原创的它来自经典图像处理里的非局部均值滤波Non-local Means。去噪场景中一个像素的干净值可以用图像中所有相似patch的像素值加权平均得到纹理区域的重复结构会被自动匹配。当年的核心思想是不要只用一个局部窗口去估计当前像素而是让整幅图像里所有相似的patch都来投票。这篇论文把这个思想搬到了神经网络里不再用固定大小的卷积窗口而是让特征图上的任意两个位置直接建立关系。这种“所有位置两两交互”的结构后来看就是注意力机制的雏形。我每次给学生讲这篇论文都会先讲非局部均值滤波因为它能让non-local这个名字立刻变得很直观——你确实是在全局范围内寻找和自己“长得像”的特征。连着读几遍论文你会发现它的方法论非常清晰。作者没有发明复杂的模块只把一个公式放进网络然后做了大量消融实验。这种风格其实很适合作为论文精读样本既看思想源头又能看到如何把想法变成可训练的模块。2. 核心公式拆解Non-local操作都做了什么2.1 公式本身并不复杂把“看整个画面”写成一个加权和论文的核心公式非常简洁用文字可以写成对于输出位置i遍历输入特征图上的所有位置j先计算位置i和位置j之间的相似度再把这个相似度作为权重对位置j经过线性变换后的特征做加权求和最后除以一个归一化因子。如果用符号表示就是y_i 1/C(x) * Σ_j f(x_i, x_j) * g(x_j)其中x是输入特征图i是当前要计算的位置j遍历所有位置f是相似度函数g是单点特征变换C是归一化因子。由于j要遍历所有位置所以不管两个点在空间上隔多远、时间上隔几帧只要它们特征相似就会直接产生一条信息传递路径。这种设计和卷积的本质区别在于卷积的权重是固定的、位置无关的而且只在局部邻域计算。non-local的权重是f实时计算出来的完全依赖于输入内容本身。你可以把卷积理解成“拿着固定模板去抠图”而non-local是“先看清画面内容再决定谁跟谁配对”。后者天然具有数据自适应性这也是它后来被各种注意力机制吸收的核心原因。2.2 g、f、C(x)三个零件各管什么先看g函数。它通常被实现为1x1卷积在视频任务里是1x1x1卷积作用是把原始特征投影到一个相对低维的空间。这个降维不是为了效果好而是为了控制后续矩阵乘法的计算量。你可以把g理解成一份“精简版个人信息”原始特征太冗余先提炼出关键信息再去做全局比较。再看f函数。它负责回答“位置i和位置j有多相关”。论文给了四个候选形式后面会专门比较。这里先记住f的输出是一个相似度标量它直接决定了加权方式是non-local模块里最灵活的部分。最后是归一化因子C(x)。它存在的意义是让输出尺度不随位置数量变化而剧烈抖动。不同f形式对应不同CGaussian系列会用所有相似度的和Dot product这类就直接用位置总数N。归一化方式不同会影响训练的稳定性和最终效果。2.3 一个可运行的PyTorch参考实现理解公式最好的方式是自己写一遍。下面这个是Embedded Gaussian版本的参考实现也是我平时最常用的版本import math import torch import torch.nn as nn class NonLocalBlock(nn.Module): def __init__(self, in_channels, inter_channelsNone): super().__init__() self.in_channels in_channels self.inter_channels inter_channels if inter_channels else in_channels // 2 # 三个投影分支g是value变换theta是query变换phi是key变换 self.g nn.Conv3d(in_channels, self.inter_channels, kernel_size1) self.theta nn.Conv3d(in_channels, self.inter_channels, kernel_size1) self.phi nn.Conv3d(in_channels, self.inter_channels, kernel_size1) # 输出投影升回原始通道数 self.W nn.Conv3d(self.inter_channels, in_channels, kernel_size1) nn.init.zeros_(self.W.weight) nn.init.zeros_(self.W.bias) def forward(self, x): batch, c, t, h, w x.size() n t * h * w # [B, C, N] g_x self.g(x).view(batch, self.inter_channels, n) # [B, C, N] - [B, N, C] theta_x self.theta(x).view(batch, self.inter_channels, n).permute(0, 2, 1) # [B, C, N] phi_x self.phi(x).view(batch, self.inter_channels, n) # 两两相似度B*N*N f torch.matmul(theta_x, phi_x) # 论文原版这一步没有除 sqrt(C)复现时建议加上避免大通道数下 softmax 输入幅度过大 f f / math.sqrt(self.inter_channels) f torch.softmax(f, dim-1) # 加权聚合B*N*C y torch.matmul(f, g_x.permute(0, 2, 1)) # 恢复空间结构 y y.permute(0, 2, 1).contiguous().view(batch, self.inter_channels, t, h, w) # 输出投影 残差连接 y self.W(y) return y x有几点需要说明。第一theta和phi就是公式里的两个embedding对应了query和key第二softmax放在最后一维等于把归一化因子C(x)设为所有相似度的和第三输出层的W被初始化为0这样网络刚开始训练时这个block等价于恒等映射不会破坏backbone的稳定状态。我在实际写代码时习惯保留缩放因子。原论文为了对标Transformer前身的思路没有用1/sqrt(d)但inter_channels比较大的时候比如512softmax之前的输入分布会影响反向传播的尺度。加一个缩放损耗在数值上更稳效果也基本不会变差。这是我复现时的经验论文里没写。2.4 为什么Embedded Gaussian会走向Self-Attention注意上面代码中fsoftmax(θ(x_i)^T φ(x_j))这正是点积注意力的核心形式。Transformer里面的Scaled Dot-Product Attention也是先做内积、再缩放、再softmax区别只在Transformer会把通道拆成多个头。换句话说Non local Neural Network其实已经迈进了self-attention的门只是没有把“多头”和“位置编码”作为显式设计。很多后来的视觉Transformer论文都会把non-local block当作自己的一个等价或前身来引用。这样你再回头读公式会更容易理解后面那些改进工作它们都是在回答同一个问题——如何用更小代价获得更好的全局交互。3. 四种相似度函数怎么选实验结果比预想温和3.1 四个候选形式一览论文里f函数有四种选择我把它们的公式和归一化方式整理成一张表函数名称f(x_i, x_j) 的具体计算归一化因子 C(x)特点Gaussianexp(x_i^T x_j)所有j的f之和最简单直接内积加指数Embedded Gaussianexp(θ(x_i)^T φ(x_j))所有j的f之和各自投影后再内积softmax等价于注意力Dot productθ(x_i)^T φ(x_j)位置总数N无指数计算量略低线性关系ConcatenationReLU(w_f^T [θ(x_i), φ(x_j)])位置总数N拼接后回归一个相似度标量其中Embedded Gaussian是我最常用的一档。它比Gaussian多学了两个投影矩阵能对特征先做“问题”和“答案”的分离表达能力更强。Dot product省了指数运算适合对训练速度有要求但不在意那零点几个点差的场景。Concatenation的实现稍微麻烦需要处理拼接的维度效果也不比其他档显著好。3.2 论文消融的结论形式引起的差距并不大作者在视频分类任务上把这四种形式都跑了一遍结果比我预想的温和得多。不同f形式之间的精度差距基本上在个位数的小数级别里打转Embedded Gaussian是表现最稳的一个但Gaussian也没有差到不能用的程度。这个结论透露出的信息是non-local核心的收益来自“两两交互”这个结构本身而不是相似度函数的具体花样。对做工程的人来说这其实是个好消息。你不用为了那零点几个点去折腾复杂的相似度设计把框架搭起来、位置放对收益就拿到了。反过来如果你在某次实验里发现换一个f之后涨了很多那大概率不是f的功劳而是换f的同时改了初始化、归一化方式或者训练配置中的其他东西。3.3 我的实验室体会优先选简单形式把预算留给位置和尺度我自己的项目里默认配置是Embedded Gaussian但这不是因为它在所有任务上都最强而是它的解释性最好softmax出来的权重可以直接可视化能很直观地告诉合作方“模型在看哪里”。可视化对排障非常有用。如果显存紧张我才会改成Dot product能省一小块指数运算主计算还是矩阵乘法那部分。另一个容易被忽略的点是不同的f形式对归一化因子C(x)的依赖不同。Gaussian系列的C(x)是动态求和的体现为softmax而Dot product、Concatenation用固定N意味着这些相似度没有被归一化到0到1输出尺度可能随特征幅度变化而波动。如果你选择的不是softmax形式后面接的归一化层一定要仔细调。我第一次用Dot product时直接搬了Embedded Gaussian的norm参数结果训练前几步loss直接飞了。4. 这个Block放在哪里效果最好视频分类与检测里的实证细节4.1 视频模型里堆在深处收益更明显原论文在视频行为识别任务上的实验用了C2D作为baseline。基线本身是2D卷积处理每一帧再靠时间维池化做融合对长时程的运动关系建模非常弱。加入non-local block后模型可以直接在时空特征上建立任意两帧之间的依赖效果提升很明显。更值得关注的是加在哪一层。论文中的消融显示把non-local block加在比较靠后的stage收益高于加在浅层。我理解原因有两个浅层特征图分辨率高位置数量N很大相似度矩阵容易爆炸而且浅层特征语义不充分算出来的相似度也许只是低层纹理的匹配深层特征已经经过了多层抽象位置之间做内积能捕捉到更有意义的关系。这个规律在后来的很多注意力网络里也被反复验证。4.2 在检测和分割任务中插入位置影响明显在目标检测和实例分割任务里作者把non-local block插在主干网络的不同阶段之间。这类任务的特点是特征图本身下采样倍数高N不会太大所以算得动。加在主干网络后半段的非局部交互对检测大目标和需要全局上下文的目标都有帮助。关键点检测上也受益因为它需要基于整个人体姿态的全局关系来预测关节点而不是只看局部patch。对于我自己的实践经验检测任务里最常见的做法是在backbone最后一个stage的输出上加一个block。那块特征图分辨率一般在14x14或28x28级别N只有几百到一千左右BNN的矩阵完全可接受。如果你想更进一步可以在特征金字塔的最高两层各加一个但多数任务上一个就够用。4.3 复现时最容易被坑的两个位置问题第一个是batch normalization的行为。non-local block内部通常也会接归一化层当batch size比较小时BN的统计量会不稳定导致训练初期震荡。小batch训练时我建议把block内部换成GroupNorm或者直接把BN的momentum调大。第二个是加block的位置与上游特征的感受野是否匹配。如果你把它加在太靠近输出的预测层可能只对最后一层特征做了全局加权但底层细节信息没有参与交互效果就会打折。我见过有人直接在检测器的检测头里堆两个non-local结果显存爆了精度还没涨这就是位置选错了。正确顺序一般是先在backbone的深层做全局建模然后再让下游头去用。5. Non-local与Self-Attention、全连接层的关系辨析5.1 与全连接层的关键差异权重是动态还是静态很多人会把non-local理解为在全图上做全连接。只从“所有位置都参与”这一点看两者确实都覆盖全部位置但本质差别在权重。全连接层的参数是训练完成后固定的任给一张输入同样的位置拿到的权重一模一样non-local的权重由当前输入特征实时计算位置i和位置j的交互强度取决于它们各自的内容。你甚至可以把它看成一种“输入自适应全连接层”。这个差异带来的实际影响是non-local可以处理不同尺寸的输入特征图而全连接层一旦确定N就无法改变输入分辨率。所以non-local天然适合插入CNN的不同阶段这是它比FC灵活的根本原因。5.2 与Self-Attention几乎同源的Embedded GaussianEmbedded Gaussian版本的non-local整个计算流程就是query、key、value三件套theta是queryphi是keyg是valuesoftmax之后加权求和。这和Transformer的attention计算几乎一模一样。区别在于Transformer在通道维度上切了多个头并设计了显式的位置编码non-local论文的原始版本没有这些。因此很多地方会把这篇论文看成视觉领域最早的self-attention实践之一。理解这条脉络对你的价值在于现在铺天盖地的各类视觉Transformer模型里的注意力模块它们改进的很多点相对位置编码、窗口注意力、分头机制其实都可以拿non-local当基准回测。你在自己任务里试一个新模块如果效果不如预期先回头看看non-local baseline往往能帮你快速定位是新机制的问题还是全局建模本身就不适合这个任务。5.3 与RNN式时序建模相比交互路径长度完全不同RNN类模型建模时序依赖要一步一步往后传递两个相隔较远的帧需要经过很多时间步才能相互作用信息一路损失。non-local在空间和时间维度上同时建立两两连接任意两帧之间的路径长度都为1信息可以直接匹配。这也是为什么它在视频任务里比传统循环模型更有优势。这一点在实践中的直接体现是输入视频时间维较长时non-local可以捕捉到开头动作与结尾结果的关联。当然路径短不代表效果好代价是计算与N的平方成正比。当你需要建模的视频分辨率高、时间长时还是得在局部稀疏连接和全局连接之间做折中。5.4 理解误区Non-local不是全局池化的花哨版还有一个常见误区是把non-local跟全局平均池化混淆。全局池化确实把整张特征图压成单个向量但它没有保留空间位置差异也没有不同位置的选择性。non-local保留了输出特征图的所有位置并且每个位置的输出都带着不同的全局上下文权重。换句话说池化是“一刀切”地汇总non-local是“戴着有色眼镜”地汇总重要位置看更多无关位置看更少。这个区别对某些任务很关键。比如一个人蹲在角落的小目标检测全局池化会把它平均掉non-local会因为该位置特征与查询点高度匹配而加大它的权重从而保留住信息。6. 计算开销与显存优化把理论变成能跑的网络6.1 复杂度藏在矩阵乘法里一个反直觉的数字示例很多人以为non-local只是加了一个“注意力”不会有多贵。真正算一下就知道问题在哪。假设输入特征是[B, C, T, H, W] [1, 512, 8, 56, 56]位置总数N 8 x 56 x 56 25088。光一个相似度矩阵尺寸就是B x N x N 1 x 25088 x 25088约6.3亿个元素。用float32存储接近2.5GB显存这还没有算特征投影和后续聚合。所以论文以及所有工程实现都不会直接在原始分辨率上硬算一定会先做通道压缩。这也是为什么non-local block里inter_channels往往设置为输入通道的一半甚至四分之一。把512压缩到128矩阵乘法本身的两个输入从512降到128但相似度矩阵大小仍然由N决定大头没变。真正决定显存的是N不是C。6.2 论文实践中怎么控制中间张量论文在主干网络上选择插入位置时实际上已经天然避开了超大N。视频任务的特征图经过多次下采样后空间和时间维都缩小到14x14或7x7级别再算全局交互才可行。目标检测任务里backbone最后一层特征图通常是32倍下采样也远小于输入分辨率。配套的还有归一化和dropout。复现的时候要注意论文中的block内部会在输出投影前加正则防止全局建模把信息都集中到少数热点位置。丢掉这部分你会在训练后期看到输出特征方差变大测试时的精度也会有波动。6.3 工程上的三个优化技巧第一个优化是降采样N。在计算相似度矩阵之前对作为key和value的特征图做一次池化或下采样把N从25088降到6272或更小相当于只让查询点去匹配一小部分代表性位置。这种方式会损失一些精细匹配能力但在视频任务里通常能保留绝大部分效果。后来很多简化工作都是这个思路的变体。第二个优化是张量内存的连续布局。写代码时注意permute之后加contiguous否则后面view会触发隐式拷贝显存瞬间翻倍。我第一次实现时就在这边吃过亏看起来只多了半行代码实际占用多了好几个G。第三个优化是混合精度。全局交互主要计算是矩阵乘法用float16在支持张量加速的设备上有接近翻倍的吞吐。如果发现训练后期精度波动可以把block内部的中间计算保留在fp32只在矩阵乘法里用fp16。这类细节对长训练任务影响很大。6.4 后来者都在做减法简化和变体思路论文出来之后围绕non-local的改进主流方向不是做加法而是做减法。有的工作发现softmax归一化不是必须的直接用全局平均池化加1x1卷积加广播加法就能模拟大部分上下文建模收益但计算量小了一个数量级。还有工作把全局稠密连接改成稀疏交叉连接每次只让位置与同行同列的位置交互再叠加两轮达到近似全局覆盖。这些变体对我们实际选型很有参考价值不用一上来就上原版non-local。先在显存预算内跑一个简化版再看效果缺口有多明显。如果简化版能追平原版说明任务本身不要求特别复杂的全局交互如果拉开差距再上更强的版本也不迟。最后说一点个人体会。这篇论文我前前后后读过几次每一次关注点都不一样。第一次读的时候满脑子都是怎么把block塞进自己的项目里跑得快第二次读开始琢磨它的消融实验为什么这么设计最近再读反倒最感兴趣的是它和后续工作之间的谱系关系。我的建议是如果你刚开始接触注意力机制不要急着去追最新的Transformer变体先把这个non-local的公式和代码吃透后续很多模块你都会觉得眼熟。如果你已经有项目在使用类似模块也可以找个时间回来翻翻原文对照一下公式和实现里有没有被你忽略的细节。这个系列我会继续更新下一篇大概率会挑一个计算更轻的全局建模工作来拆到那时候再一起聊。
