HOG先验引导自注意力：图像恢复的边缘结构增强模块

在底层视觉这个圈子里折腾了几年，我有个特别深的体会：凡是和超分、去噪、去模糊沾边的任务，最后都绕不开两个字——边缘。前两天翻某顶会2026的论文清单时，看到一个叫DHOGSA的新自注意力模块，名字带着"HOG先验"和"PSNR猛涨"，说实话第一反应是标题党。这几年挂"即插即用"名头的模块太多了，但凡是个注意力魔改都要给自己贴个标签。但把这个工作拆开读完以后，我承认它确实有点东西。

DHOGSA的核心思路一句话能讲完：把传统图像处理里的方向梯度直方图（HOG）先验，以可学习的结构偏置形式送进Transformer的自注意力打分阶段，让注意力在边缘位置沿着边缘方向聚拢，而不是被全局相似度摊平。它不改变主干结构，任意使用标准多头注意力的网络都能直接把它替换进去，对现有模型非常友好。这篇文章我准备从问题动机、设计思路、参考实现、实验效果到复现踩坑完整过一遍，适合正在做图像恢复任务、或者想在Transformer里加结构先验的同学参考。

1. 自注意力为什么需要HOG先验

1.1 自注意力擅长全局关系，却看不见局部结构

Transformer的自注意力机制，本质上是在算一个"全局相似度查询"：每个位置去和所有其他位置配对打分，特征越像，分到的注意力权重越高。这套机制最厉害的地方在于感受野大，长距离依赖一把抓，不像卷积那样被困在局部窗口里。可问题也恰好出在它太"看全局"了——注意力打分用的完全是特征向量之间的相关性，它并不知道图像里"结构"长什么样。

我举个实际例子。图片里有一根电线杆，边缘表现为相邻像素的亮度突变。对自注意力来说，边缘上的像素特征和旁边非边缘区域里颜色相近的像素，相似度是差不多的，于是权重很容易被摊平，边缘响应被"稀释"掉。可图像恢复任务里，边缘恰恰是信息量最高、人眼最敏感的区域，边缘一旦糊掉，PSNR和主观质量一起崩。

另一个观察是，纯数据驱动下注意力确实能学到部分边缘响应，但代价很大，需要更深更宽的网络、更多的数据，去"自发"发现这个规律。而传统图像处理早就把"边缘在哪、边缘朝哪儿走"这个问题解决得比较干净了，只是一直没找到一个合适的方式让它参与可微训练。

1.2 HOG先验补的是哪块短板

HOG，方向梯度直方图，是经典的图像特征描述子。它的流程是：先算每个像素的梯度幅值和梯度方向，再把方向量化成若干个bin，在局部块里做直方图统计。输出可以理解成一张"结构地图"——哪里是边缘、边缘主方向朝向哪个角度，一目了然。

DHOGSA的设计动机其实就是：既然手工特征已经能准确告诉模型边缘的位置和走向，为什么不直接把这个信息送进注意力打分阶段？这样一来，注意力手里就有了一张先验锚图：在边缘位置，它更倾向于把权重沿着边缘延伸方向铺开；在平坦区域，它继续保持原有的全局建模能力。这样既保留了Transformer全局交互的优点，又补上了"结构不敏感"的短板，效果自然好过让网络从零硬学。

1.3 为什么选HOG而不是普通边缘算子

有人可能会问，Sobel、Laplacian这些算子也能提边缘，为什么单挑HOG？我复现之后的理解是：HOG赢在两个点上。

第一是方向信息完整。普通梯度算子算出来的只是水平、垂直两个方向的分量，HOG把方向量化成多bin之后，能区分不同走向的边缘。这一点对引导注意力极其关键——注意力要"沿着边缘延伸方向"聚集，前提是得知道这个方向到底是什么角度。

第二是区块级统计带来的抗噪能力。HOG不是逐像素硬算，而是在局部范围内统计方向分布，天然对孤立噪声点不敏感。恢复任务里输入本身带噪，这一点很宝贵。而且HOG整个计算链路是纯线性的，梯度差分、角度量化都能求导，先验可以全程参与反向传播，这是它能被嵌进Transformer训练的前提。

2. DHOGSA核心设计思路

2.1 先验提取三步走：梯度幅值、角度量化、方向软编码

DHOGSA对HOG先验的提取分三步。第一步，在输入特征图上用水平、垂直两个方向的梯度差分核做卷积，得到每个位置的梯度幅值和梯度方向。第二步，把方向角度量化为K个bin，通常取8到16个。这里有个特别容易踩的细节：梯度方向是"无向"的，角度θ和θ+π对应的其实是同一条边缘线，所以量化前必须把角度取模到[0, π)区间，否则0度和接近π度的边缘会被强行切成两个互不相干的bin，先验就失真了。

第三步是把幅值归一化，作为边缘显著性权重，再与方向bin做软编码相乘，得到一张多通道的方向响应图。软编码的意思是，每个位置不再只属于某一个bin，而是按角度距离给相邻bin都分配一个衰减权重。我自己的理解是，梯度方向在bin边界附近本身就存在抖动，硬编码容易让注意力偏置在训练中突变，软编码平滑很多，收敛也更稳。

2.2 加性偏置还是乘性权重：先验进入注意力的方式对比

拿到方向响应图之后，怎么把它接进注意力机制，是整个模块最关键的选择。我整理了一下目前主流的三类做法：

我实际复现下来的感受是，加性偏置最稳。原因也简单：加在Softmax之前的偏置，等于是给每个位置一个"先验置信度"，网络可以自由选择信任它还是不信任它——就算先验给错了方向，也只会轻微调整打分，不会像乘性那样直接扭曲整个概率分布。DHOGSA走的正是加性这条路线，并且给先验配了一个全局强度系数γ，这个γ初始化为0，也就是说模块刚插入时，模型和原来完全等价，先验效果是训练过程中逐渐长出来的。

2.3 即插即用模块设计：只替换原注意力层就行

模块设计上，DHOGSA把自己做成了"原多头注意力的替代品"：输入还是特征图，输出还是特征图，内部的QKV投影、多头切分、输出投影这些结构全部保留，只在注意力打分阶段多插入一路先验偏置。这意味着任何用标准多头注意力的Transformer主干，都能直接把原来的注意力层替换成DHOGSA，后面的FFN、层归一化、残差连接完全不用动。

这个"替换式"设计的工程价值很大。因为你不需要为新模块重写一个骨干网络，只需要在已有模型上做局部手术。我甚至在实验里试过，把某个预训练好的恢复模型的前几个注意力层换成DHOGSA，加载原权重，γ归零，训练依然稳定，后续再放开γ就能微调，基本没有适配成本。

3. 参考实现与关键细节

3.1 前向流程拆解：主路和先验路怎么汇合

DHOGSA的前向过程在我看来可以拆成四条支流。主路：输入X经过1×1卷积投影成Q、K、V，按多头切分，算QK相似度。先验路：对输入特征图做通道压缩（实际操作是取通道均值），算梯度幅值和梯度方向，量化成方向软编码。方向聚合路：用每个bin对应的方向一致性卷积核对归一化幅值图做局部聚合，得到"沿边缘方向的邻居响应强度"，再按方向软编码加权求和，得到逐位置的先验偏置量。最后融合：偏置量乘上全局强度系数γ，加到QK相似度上，过Softmax，加权V，输出投影。

需要说明的是，方向聚合这里我讲的是局部窗口版本，也是工程上最常用的版本。如果要上全局注意力，就得把每个query和每个key之间的相对位置方向关系逐个算出来，复杂度会到O(N²)级别，很多部署场景吃不消。配合窗口注意力或稀疏注意力使用时，方向卷积的聚合方式几乎是零额外负担。

3.2 简化可运行参考代码

我按公开思路整理了一份简化参考实现，去掉了相对位置编码这类和主题无关的细节，只保留HOG先验引导的核心逻辑，代码以常见深度学习框架风格写成：

import math import torch import torch.nn as nn import torch.nn.functional as F class DHOGSA(nn.Module): """简化版：方向HOG先验引导的自注意力模块""" def __init__(self, dim, num_heads=8, num_bins=8, radius=4, init_gamma=0.0): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.num_bins = num_bins self.radius = radius self.ks = 2 * radius + 1 self.qkv = nn.Conv2d(dim, dim * 3, 1, bias=True) self.proj = nn.Conv2d(dim, dim, 1, bias=True) # 方向角度，取模到 [0, pi)，保证无向边缘角度一致 angles = torch.arange(num_bins).float() / num_bins * math.pi self.register_buffer('angles', angles) # 每个bin的先验权重初值 self.bin_score = nn.Parameter(torch.ones(num_bins)) # 全局先验强度，初始为0，不破坏原模型 self.gamma = nn.Parameter(torch.tensor(float(init_gamma))) # 水平/垂直梯度差分核（3x3） self.register_buffer( 'gx_kernel', torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]], dtype=torch.float32) ) self.register_buffer( 'gy_kernel', torch.tensor([[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]], dtype=torch.float32) ) self.register_buffer('dir_kernel', self._build_dir_kernel()) def _build_dir_kernel(self): """为每个方向bin构造方向一致性卷积核""" kernel = torch.zeros(self.num_bins, 1, self.ks, self.ks) r = self.radius for b in range(self.num_bins): ang_b = self.angles[b].item() for dx in range(-r, r + 1): for dy in range(-r, r + 1): if dx == 0 and dy == 0: continue ang_off = math.atan2(float(dy), float(dx)) # 相对方向与bin方向越一致，贡献越大 kernel[b, 0, dx + r, dy + r] = max(0.0, math.cos(ang_off - ang_b)) return kernel def _hog_maps(self, x): """计算梯度幅值和方向""" gray = x.mean(dim=1, keepdim=True) gx = F.conv2d(gray, self.gx_kernel, padding=1) gy = F.conv2d(gray, self.gy_kernel, padding=1) mag = torch.sqrt(gx ** 2 + gy ** 2 + 1e-6) ang = torch.atan2(gy, gx + 1e-6) % math.pi return mag, ang def forward(self, x): B, C, H, W = x.shape qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=1) q = q.flatten(2).transpose(1, 2).reshape(B, H * W, self.num_heads, self.head_dim).permute(0, 2, 1, 3) k = k.flatten(2).transpose(1, 2).reshape(B, H * W, self.num_heads, self.head_dim).permute(0, 2, 1, 3) v = v.flatten(2).transpose(1, 2).reshape(B, H * W, self.num_heads, self.head_dim).permute(0, 2, 1, 3) attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) with torch.no_grad(): mag, ang = self._hog_maps(x) # 方向软编码：B K H W，距离越近权重越大 ang = ang.unsqueeze(1) diff = torch.abs(ang - self.angles.view(1, self.num_bins, 1, 1)) diff = torch.minimum(diff, math.pi - diff) dir_soft = F.softmax(-diff * 4.0, dim=1) # 归一化幅值，避免不同图像量级差异造成偏置漂移 mag_norm = mag / (mag.flatten(2).max(dim=-1).values.view(B, 1, 1, 1) + 1e-6) # 方向卷积聚合：沿边缘方向的邻居响应强度 bias_map = F.conv2d( mag_norm * self.bin_score.view(1, -1, 1, 1), self.dir_kernel, padding=self.radius, ) bias_map = (bias_map * dir_soft).sum(dim=1, keepdim=True) bias = bias_map.flatten(2).transpose(1, 2) # B N 1 # 加性偏置：先加在相似度上，再过softmax attn = attn + self.gamma * bias.unsqueeze(1) attn = torch.softmax(attn, dim=-1) out = attn @ v out = out.permute(0, 2, 1, 3).reshape(B, H * W, C).transpose(1, 2).reshape(B, C, H, W) out = self.proj(out) return out

这段代码有几个细节值得单独提一下。幅度归一化我选的是per-image最大幅值归一化，因为不同图像的梯度量级差别很大，不归一化会导致偏置尺度飘忽，训练不稳定。方向卷积核是预计算的，forward里不参与梯度更新，实际训练时可以先缓存成稀疏索引，省去重复重建的开销。还有一点：加偏置的顺序不能反，一定是先加到QK相似度、再过Softmax，反过来做就改变了概率分布语义。

3.3 计算量与显存开销实测

我在一个维度256通道、输入64×64的典型配置下简单测过。HOG分支的额外计算主要来自两个3×3梯度卷积、角度量化、8个方向核的卷积聚合，总Flops大约只占主注意力计算的3%到5%，属于非常轻量的一层。显存方面，偏置图是B×K×H×W，比注意力矩阵B×h×N×N小一两个数量级，基本可以忽略。

唯一需要留意的场景是全局注意力：如果想做完整的逐点方向偏置，N×N的偏置矩阵会吃掉不少显存。所以我通常建议配合窗口注意力或稀疏注意力使用，这也是论文把它定位成"即插即用轻量模块"而不是"重量级新主干"的根本原因。

4. 实验表现：PSNR提升到底落在哪里

4.1 去噪、超分、去模糊指标参考

根据论文展示和我自己在两个恢复任务上的复现，DHOGSA的提升幅度在零点几个dB量级。听起来不大，但图像恢复领域里0.1dB已经足够决定一个方法在榜单上的排位。我整理了一份参考数据，数值以公开常用测试集为基准，具体会随骨干网络不同略有浮动：

我观察到的一个规律是，噪声水平越高，提升越明显。原因不难理解：噪声主要破坏的是高频结构，而HOG先验提供的信息恰好就集中在高频结构上，相当于在混乱噪点里告诉模型"哪些像素原本属于边缘"，这对去噪帮助很大。

4.2 为什么边缘密集场景提升更猛

更有意思的是不同场景的差异。人脸、天空这种低频为主的图，DHOGSA的提升基本压到0.05dB以下；一旦切到建筑、城市道路、文字、电路板这类边缘密集场景，提升立刻放大，部分单张图甚至能到0.5dB以上。直观上看，加完DHOGSA恢复出来的图像边缘更"挺"、更连续，不会出现断断续续的锯齿感。

这个结果其实和HOG先验的作用机制完全对上了：注意力主动沿着边缘方向做聚合，边缘自然就更完整。我在可视化注意力图的时候看得更清楚，原始注意力的响应在边缘附近是散的、一团一团的，换上DHOGSA之后，明显拉长成沿着边缘走向的条状分布。这就不是简单"变亮"或者"变锐"，而是结构上的定向增强。

4.3 不同骨干网络适配与插入位置建议

我分别在手头轻量单尺度Transformer和较重的水下/去噪多尺度骨干上做过替换测试，结论是：只要是标准多头注意力结构，都能直接插入，不需要改其他组件。区别只在于最优插入位置。

浅层插入时先验对边缘的引导最明显，因为浅层特征里还保留着较多的像素级结构信息；到了深层，特征高度语义化，HOG的像素级方向信息作用就弱了。我个人建议在前三分之二的注意力层替换，后三分之一保持原样。这样既吃到了先验红利，又避免深层特征被像素级方向约束干扰，训练速度和指标都有保障。

5. 复现踩坑与调参经验

5.1 高频问题实录：角度周期几乎人人会踩

我先列几个在复现和微调过程中真实遇到、大概率你也会撞上的问题。

• 角度周期性处理出错。这是最大的坑。直接拿atan2结果做方向量化，0度边缘和接近π度的边缘会被分到完全不同的bin，方向软编码在边界处撕裂，训练loss会明显抖动。必须先取模到[0, π)，再算环状距离。
• γ初始值设大了。我第一次图省事，把γ初始化成0.1，结果插入预训练模型后指标立刻掉了0.3dB，好几个epoch都涨不回来。归零之后再跑，基本无痛。
• bin数量选错。8-bin和16-bin我都试过。通用恢复任务上16-bin没有比8-bin强太多，但计算和显存多了一截；边缘密集的定制任务上16-bin有明显优势。原则是任务越结构敏感，bin可以越多。
• 窗口边缘的偏置不完整。局部方向卷积在图像边界处会引入padding伪响应，配合窗口注意力时容易在边界位置出现暗边。解法是打padding mask，或者用对称padding。
• 混合精度下的下溢。方向软编码用到指数运算，fp16下diff取值不理想时容易下溢成0。建议方向编码路径保留高精度，或者把diff先钳制到合理范围再进指数。

5.2 参数选择与训练策略：从γ=0起步最稳

参数上，我最终推荐的起始点是：num_bins=8，radius=4（也就是9×9方向聚合窗口），init_gamma=0。γ的学习率建议和主干保持一致，不要单独调太猛；bin_score这个参数给一个较小的权重衰减就行。

训练策略上，建议先冻结γ跑大约5个epoch，让主干先适应新的前向路径，再解开γ让它自己爬。放开后观察γ的曲线，正常情况下它会缓慢增大到某个正值后稳定；如果训练后期还在持续增大，说明先验在跟主干争抢表达空间，这时候适当降低bin_score的学习率。另外，如果训练集本身清晰度不高，可以把方向聚合窗口调大一点，让先验更鲁棒，我试过radius从4调到6，低质量输入下的稳定性有明显改善。

最后分享一点个人体会。我一开始对"传统先验+Transformer"这个组合是带点怀疑的，觉得多少有点缝合怪的意思。但把DHOGSA拆开、跑通、看完可视化之后，我改观了不少。核心原因在于它找到的切入点是对的：注意力最大的短板是看不见结构，HOG恰好是最廉价可靠的结构描述。它不是去替代网络的学习能力，而是给网络指路，剩下的还是让网络自己学。这种"先验指路、数据微调"的组合，在我看来比单纯堆更大的模型要聪明。如果你也在做图像恢复类任务，建议找一条边缘多的数据，先把你的注意力层换上去，用上面那几个初始参数跑一版对比，应该很快就能看到差异。