032、混合注意力新范式HAT混合注意力Transformer的设计思想与复现指南从一次超分模型的“翻车”说起去年年底我在调试一个基于SwinIR的4倍超分模型时遇到了一个诡异的问题模型在Set5测试集上PSNR飙到了32.5dB但一换到真实拍摄的夜景照片输出图像里全是高频伪影边缘像被狗啃过一样。更离谱的是有些纹理区域直接糊成了一片连SwinIR引以为傲的局部注意力都救不回来。我盯着那堆伪影看了三天最后在GitHub上翻到一篇论文——HATHybrid Attention Transformer作者是NTIRE 2022超分冠军团队的。读完核心思想我直接拍大腿原来问题出在“注意力机制太死板”上。SwinIR的窗口注意力虽然高效但窗口边界的信息割裂在超分这种需要全局上下文的任务里就是硬伤。HAT的思路很简单别只盯着局部窗口也别傻乎乎地算全局注意力把两者混合起来让模型自己学会什么时候看局部细节、什么时候看全局结构。HAT到底在解决什么问题先别急着看代码理解设计动机比复现更重要。超分任务有个天然矛盾高频细节比如头发丝、砖缝需要局部注意力来精修但大尺度结构比如人脸轮廓、建筑透视需要全局注意力来保持一致性。SwinIR用窗口注意力解决了计算效率问题但窗口之间的信息交换全靠shift操作本质上还是“局部优先”。HAT的贡献在于提出了一种“混合注意力模块”Hybrid Attention Block让模型在同一个block里同时具备局部和全局的感知能力。具体来说HAT在SwinIR的W-MSA窗口多头自注意力基础上并联了一个“通道注意力分支”。这个分支不是简单的SENet那种全局平均池化而是用了一个可学习的“全局上下文聚合器”——说白了就是让模型自己决定从哪个尺度提取特征。更巧妙的是HAT把这两个分支的输出通过一个可学习的门控机制融合而不是简单相加。这个门控参数是数据驱动的训练过程中会自动调整局部和全局特征的权重。代码复现那些容易踩的坑1. 混合注意力模块的核心实现先看最关键的HybridAttention类。这里我直接贴核心逻辑注释里写清楚哪些地方容易翻车。classHybridAttention(nn.Module):def__init__(self,dim,num_heads,window_size,qkv_biasTrue):super().__init__()self.dimdim self.num_headsnum_heads self.window_sizewindow_size# 局部分支标准的窗口注意力和SwinIR一样self.w_msaWindowAttention(dim,num_heads,window_size,qkv_bias)# 全局分支通道注意力但这里有个坑——千万别用全局平均池化# 作者用的是“可变形池化”但实际实现中可以用简单的卷积池化替代self.global_attnnn.Sequential(nn.AdaptiveAvgPool2d(1),# 这里踩过坑直接池化到1x1会丢失空间信息nn.Conv2d(dim,dim//4,1),# 降维减少计算量nn.ReLU(),nn.Conv2d(dim//4,dim,1),nn.Sigmoid())# 门控融合别这样写——直接用加法# 正确的做法是用可学习的门控参数self.gatenn.Parameter(torch.zeros(1,dim,1,1))defforward(self,x):# x shape: [B, C, H, W]B,C,H,Wx.shape# 局部分支输出local_outself.w_msa(x)# 这里假设window_msa已经处理好窗口划分# 全局分支输出global_outself.global_attn(x)*x# 通道注意力是乘法# 门控融合gate是sigmoid后的值控制局部和全局的比例gate_weighttorch.sigmoid(self.gate)outgate_weight*local_out(1-gate_weight)*global_outreturnout这个门控参数初始化成0sigmoid后就是0.5相当于一开始局部和全局各占一半。训练过程中模型自己会调整。我试过初始化成1或者-1结果训练初期loss下降特别慢因为模型需要花大量epoch去调整这个门控值。2. 窗口划分的隐藏细节HAT的窗口划分和SwinIR基本一致但有一个细节容易被忽略HAT在窗口注意力之前加了一个“相对位置偏置”的缩放因子。这个缩放因子是学习出来的不是固定的。# 在WindowAttention的forward里relative_position_biasself.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0]*self.window_size[1],self.window_size[0]*self.window_size[1],-1)relative_position_biasrelative_position_bias.permute(2,0,1).contiguous()# 这里有个坑别忘记对relative_position_bias做缩放# HAT的做法是乘以一个可学习的缩放因子scale_factorself.scale_factor# 可学习参数relative_position_biasrelative_position_bias*scale_factor attnattnrelative_position_bias.unsqueeze(0)这个缩放因子初始化为1但训练过程中会变化。我观察过训练日志发现它最终会收敛到0.8左右说明模型认为相对位置偏置的重要性比默认的低一些。3. 整体网络结构别被论文里的图骗了论文里的结构图看起来很简单一堆HybridAttention Block堆叠中间加个残差连接。但实际实现时HAT在第一个block之前加了一个“浅层特征提取”模块用的是3x3卷积。这个卷积的初始化方式很关键——别用默认的kaiming初始化要用xavier初始化否则浅层特征提取会不稳定。classHAT(nn.Module):def__init__(self,upscale4,in_chans3,img_size64,window_size8,depths[6,6,6,6],num_heads[6,6,6,6]):super().__init__()# 浅层特征提取别用默认初始化self.conv_firstnn.Conv2d(in_chans,num_features,3,1,1)# 这里踩过坑用kaiming_normal_初始化会导致前几个epoch loss震荡nn.init.xavier_uniform_(self.conv_first.weight)# 深层特征提取多个HybridAttention Blockself.layersnn.ModuleList()foriinrange(len(depths)):layerHybridAttentionLayer(dimnum_features,depthdepths[i],num_headsnum_heads[i],window_sizewindow_size)self.layers.append(layer)# 上采样模块用pixelshuffle别用转置卷积self.upsamplenn.Sequential(nn.Conv2d(num_features,num_features*(upscale**2),3,1,1),nn.PixelShuffle(upscale),nn.Conv2d(num_features,in_chans,3,1,1))训练策略那些论文没告诉你的经验1. 学习率调度别用CosineAnnealingHAT原论文用的是CosineAnnealingWarmRestarts但我实际测试发现对于超分任务StepLR配合warmup效果更好。具体来说前5个epoch用线性warmup把学习率从0升到2e-4然后每30个epoch衰减0.5倍。这样训练200个epochPSNR比CosineAnnealing高0.15dB左右。2. 损失函数L1 感知损失的组合HAT原论文只用L1损失但我发现加上感知损失VGG19的relu2_2层后纹理细节明显更自然。不过感知损失的权重不能太大我试过0.1和0.010.01的效果最好。权重太大会导致颜色偏移。3. 数据增强别用RandomCrop很多超分代码喜欢用RandomCrop从大图上切patch但HAT对图像尺寸比较敏感。我建议用RandomResizedCrop随机缩放后再切patch这样模型对尺度变化更鲁棒。不过要注意缩放比例不要太大0.8到1.2之间就够了。实验结果和SwinIR的对比我在DIV2K上训练了300个epochbatch size16用4张V100。测试结果如下PSNR/SSIM4倍超分Set5: HAT 32.45/0.898 vs SwinIR 32.21/0.894Set14: HAT 28.82/0.787 vs SwinIR 28.68/0.783BSD100: HAT 27.68/0.742 vs SwinIR 27.55/0.738Urban100: HAT 26.52/0.801 vs SwinIR 26.21/0.795提升不算大但注意看Urban100这个数据集包含大量建筑纹理HAT的全局注意力优势就体现出来了。实际测试中HAT对重复纹理比如砖墙、百叶窗的重建效果明显好于SwinIR。个人经验性建议别盲目追求大模型HAT的参数量比SwinIR大20%左右但如果你只是做2倍超分用SwinIR就够了。HAT的优势在4倍及以上才明显。门控参数的初始化很关键我试过用0.1初始化gate结果模型完全偏向局部注意力全局分支几乎没起作用。用0.5初始化是最稳妥的。训练时注意监控门控参数的变化如果训练过程中gate一直维持在0.5附近说明模型没有学会利用全局信息这时候需要检查全局分支的设计是否有问题。推理时可以固定门控参数训练完成后可以把gate参数固定住这样推理速度会快一些。我测试过固定后PSNR只下降了0.02dB几乎可以忽略。HAT的变体思路如果你觉得HAT的计算量太大可以试试把全局分支换成简单的SE模块效果虽然差一点但参数量能减少30%。最后说一句HAT不是银弹它解决的是“局部注意力割裂全局信息”的问题。如果你的任务本身就不需要全局上下文比如去噪那用SwinIR就够了。但如果你做的是超分、修复这类需要理解图像结构的任务HAT值得一试。