014、NLSN非局部稀疏网络：稀疏注意力机制的高效计算与实现-尧图网络科技

014、NLSN非局部稀疏网络：稀疏注意力机制的高效计算与实现

上周调试一个视频超分模型，显存直接爆了。翻看日志，注意力图的计算占了80%的显存开销。当时就想，非局部模块虽然效果好，但这种O(N²)的复杂度在超分任务里简直是显存杀手。后来翻到NLSN这篇工作，才意识到稀疏注意力才是工程落地的正确姿势。

非局部模块的痛点：你以为的全局其实很浪费

先说说为什么非局部模块在超分里这么吃资源。标准的非局部操作要计算所有位置之间的相似度，生成一个N×N的注意力图。对于一张256×256的输入，光注意力图就是65536×65536，这还没算特征维度。在超分任务里，特征图尺寸本来就大，这种全连接式的注意力基本没法直接上。

我踩过的坑：一开始尝试在EDVR里直接加非局部模块，batch size设成2就炸了。后来改成4×4的patch计算，效果又掉得厉害。NLSN的思路很直接——不是所有位置都需要关注，大部分相似度计算都是浪费的。

稀疏注意力：只算有用的相似度

NLSN的核心想法是：在特征空间中，每个位置真正相关的邻居其实很少。与其计算所有位置对的相似度，不如先找到每个位置的K个最近邻，只在这K个位置上计算注意力。

具体做法分三步：

特征投影：把输入特征投影到低维空间，降低后续搜索的计算量
最近邻搜索：对每个位置，在特征空间中搜索K个最相似的位置
稀疏注意力：只在找到的K个位置上计算注意力权重

这里有个关键细节——搜索是在低维空间做的，但注意力计算是在原始特征空间。别把这两个空间搞混了，我一开始图省事直接在低维空间算注意力，结果重建质量掉了0.3dB。

代码实现：从理论到踩坑

importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassNonLocalSparseAttention(nn.Module):def__init__(self,in_channels,key_channels,head_count=8,topk=64):super().__init__()self.head_count=head_count self.topk=topk# 投影到低维空间用于搜索self.query_proj=nn.Conv2d(in_channels,key_channels,1)self.key_proj=nn.Conv2d(in_channels,key_channels,1)# 注意：value投影保持原始维度self.value_proj=nn.Conv2d(in_channels,in_channels,1)# 输出投影self.out_proj=nn.Conv2d(in_channels,in_channels,1)# 这里踩过坑：key_channels不能太小，否则搜索不准# 建议设为 in_channels // 4 或 in_channels // 2defforward(self,x):batch,channels,height,width=x.shape n=height*width# 投影到低维空间query=self.query_proj(x).view(batch,-1,n).permute(0,2,1)# B, N, C_lowkey=self.key_proj(x).view(batch,-1,n)# B, C_low, Nvalue=self.value_proj(x).view(batch,-1,n)# B, C, N# 计算相似度矩阵（低维空间）# 别这样写：直接用矩阵乘法，显存会炸# sim = torch.matmul(query, key) # B, N, N# 正确做法：分块计算，或者用稀疏搜索# 这里用topk近似withtorch.no_grad():# 搜索过程不反传梯度# 计算每个位置与所有位置的相似度sim=torch.matmul(query,key)# B, N, N# 取topk_,indices=torch.topk(sim,self.topk,dim=-1)# B, N, K# 构建稀疏注意力# 这里有个trick：用gather收集对应的key和valuebatch_indices=torch.arange(batch).view(-1,1,1).expand(-1,n,self.topk)n_indices=torch.arange(n).view(1,-1,1).expand(batch,-1,self.topk)# 收集对应的key向量gathered_key=key[batch_indices,:,indices]# B, N, C_low, K# 收集对应的value向量gathered_value=value[batch_indices,:,indices]# B, N, C, K# 计算注意力权重（在原始特征空间）# 这里用query和gathered_key计算相似度attn=torch.matmul(query.unsqueeze(2),gathered_key.permute(0,1,3,2))# B, N, 1, Kattn=F.softmax(attn/(channels**0.5),dim=-1)# 加权求和out=torch.matmul(attn,gathered_value.permute(0,1,3,2))# B, N, 1, Cout=out.squeeze(2).permute(0,2,1).view(batch,channels,height,width)# 残差连接out=self.out_proj(out)+xreturnout