1. 从“看局部”到“看全局”:为什么图像超分辨率需要混合注意力?

大家好,我是老张,在AI图像处理这个领域摸爬滚打了十来年,从最早的插值算法玩到现在的Transformer,感觉技术迭代的速度真是越来越快了。今天想和大家深入聊聊一个在CVPR 2023上大放异彩的新模型——HAT(Hybrid Attention Transformer)。这玩意儿在图像超分辨率(SR)任务上表现相当抢眼,尤其是处理那些纹理复杂、细节丰富的图片时,效果提升非常明显。

咱们先来打个比方,理解一下图像超分辨率的核心难点。想象一下,你手里有一张模糊的老照片,你的任务就是把它修复得清晰、有细节。传统的卷积神经网络(CNN)就像是一个拿着放大镜、一丝不苟的工匠,它非常擅长捕捉照片里局部的、重复的纹理,比如砖墙的纹路、树叶的形状。但是,这个工匠有个习惯,他只看自己放大镜框住的那一小块区域,对于整张照片的“全局布局”和“远距离关联”不太敏感。这就导致一个问题:当修复一些需要结合照片多个区域信息才能判断的细节时(比如远处窗户的轮廓是否和近处的线条连贯),CNN可能会力不从心,产生不自然的伪影或者模糊。

后来,Transformer架构横空出世,尤其是SwinIR这类模型,它引入了一种“自注意力”机制。这就像给修复师配了一个可以随时纵观全局的“上帝视角”。修复师在观察照片的某一块时,可以同时参考照片上任何其他位置的信息,来判断“这里应该是什么纹理”。这个机制对于理解图像的全局结构非常有力。但问题又来了,有研究者通过一种叫“局部归因图(LAM)”的分析方法发现,像SwinIR这样的纯Transformer模型,在实际计算时,并没有我们想象中那么充分地利用整张输入图片的所有像素信息。它可能还是过于关注某些局部窗口内的交互,导致“潜力”没有被完全激发,一些有用的像素信息被浪费了。

所以,HAT要解决的核心矛盾就来了:如何既保留CNN那种对局部特征精细提取的能力,又能像Transformer一样拥有强大的全局信息整合能力? 它的答案就是“混合注意力”。简单说,它不再非此即彼,而是把擅长捕捉全局统计信息的“通道注意力”和擅长建立长距离依赖的“基于窗口的自注意力”巧妙地融合在了一起。这种“双管齐下”的思路,目的就是为了“激活更多像素”,让模型在重建高清图像时,能调动起低清图中每一分有用的信息,从而在减少伪影、提升纹理真实感上实现突破。

2. HAT的核心武器:混合注意力模块(HAB)拆解

HAT的网络骨架是基于改进的Swin Transformer构建的,而它的灵魂就在于其中嵌入的混合注意力块(Hybrid Attention Block, HAB)。这个模块的设计非常精妙,咱们来仔细看看它到底是怎么工作的。

2.1 通道注意力:给每个特征通道“加权”

首先,我们来理解一下什么是通道注意力。你可以把一张图像经过卷积后得到的特征图,想象成一组叠在一起的胶片,每一层胶片(一个通道)都负责提取某种特定的特征,比如有的层专门响应边缘,有的层专门响应颜色,有的层专门响应纹理。

通道注意力机制的作用,就是给每一层胶片(每个通道)计算一个权重。这个权重是通过分析所有空间位置(即整张特征图) 的全局信息得来的。如果某一层胶片包含的信息对于当前的重建任务特别重要(比如在修复人脸时,包含眼睛细节的通道),那么它的权重就会被提高;反之,不那么重要的通道权重就会被降低。这个过程是全局的,它让模型能从整体上把握哪些特征类型更需要被关注。

在HAT的HAB模块里,通道注意力被集成在了标准Transformer块的前馈网络(FFN)部分之前。具体来说,模型会先对特征进行全局平均池化,得到一个代表每个通道全局信息的向量,然后通过一个小型神经网络(通常是两个全连接层)来生成权重,最后将这些权重乘回对应的特征通道上。这一步操作计算量不大,但效果显著,它确保了模型在像素级重建时,能够依据全局统计信息来“调配资源”。

2.2 基于窗口的自注意力:建立局部区域的深度关联

另一方面,HAT保留了Transformer核心的基于窗口的自注意力(Window-based Self-Attention)。为了控制计算复杂度,模型不会一次性计算整张大图所有像素之间的关系(那计算量是天文数字),而是将特征图划分成一个个不重叠的局部窗口(比如16x16像素),然后在每个窗口内部计算自注意力。

在这个窗口里,每一个像素(称为一个“token”)都会与窗口内的所有其他像素进行计算,评估它们之间的相关性。这就像在一个小会议室里,每个人都在倾听并回应其他人的发言,从而形成对这个会议室话题的深度共识。这种机制让模型能够非常精细地建模窗口内像素间复杂的依赖关系,对于恢复精细的局部纹理至关重要。

HAT在这里做了一个关键的改进:它增大了窗口尺寸。原始的SwinIR通常使用较小的窗口(如8x8),虽然节省计算成本,但每个像素的感受野受限。HAT将窗口大小扩大到16x16,让每个像素在第一步就能看到更大范围的邻居,这直接增强了模型在局部区域内的表征能力。当然,窗口变大计算量也会增加,但通过后续巧妙的模块设计,HAT在性能和效率之间找到了更好的平衡点。

2.3 混合的威力:1+1>2

那么,把通道注意力和窗口自注意力混合在一起,到底带来了什么好处呢?我自己的理解是,这形成了一种“宏观统筹”与“微观搞活”的协同。

  • 通道注意力(宏观统筹):它站在全局高度,告诉模型:“当前阶段,纹理通道比颜色通道更重要”或者“边缘信息需要被加强”。它激活的是对特征类型的全局选择。
  • 窗口自注意力(微观搞活):它在每个局部窗口内,让像素们充分交流,决定“这个角落的阴影应该如何与旁边的线条衔接”。它激活的是空间位置间的精细关联。

这两种注意力关注的维度不同,是互补的。通道注意力弥补了纯窗口自注意力可能缺乏的全局信息感知,帮助模型从整体上判断该激活哪些特征;而窗口自注意力则在通道注意力筛选出的重要特征内部,进行深度的、空间上的细节重建。实测下来,这种混合机制确实能引导模型去利用更多被传统Transformer忽略的输入像素信息,从而生成更自然、细节更丰富的高分辨率图像。

3. 打破窗口壁垒:重叠交叉注意力模块(OCA)

如果说HAB模块解决了“用什么方式看”的问题,那么HAT引入的另一个创新模块——重叠交叉注意力(Overlapping Cross-Attention, OCA)——就是为了解决“在哪里看”的边界问题。

3.1 移位窗口的局限与伪影的产生

在SwinIR等模型中,为了能让信息在不同窗口间传递,采用了“移位窗口”的策略。即在下一层,窗口的划分位置会移动一半窗口大小,这样新的窗口就会包含上一层不同窗口的部分像素,从而实现间接的信息交互。

但HAT的论文作者指出,这种间接交互并不完美。在实际生成的中间特征图上,有时能看到类似“马赛克”或“块状”的不连续痕迹,这就是阻塞伪影。它产生的根本原因在于,窗口与窗口之间的边界区域,像素缺乏直接、充分的交流。当一个窗口在独立计算时,它对其边缘像素的处理,可能无法与相邻窗口对同一边缘的处理完美对齐,在重建高频细节时就会产生接缝或不一致。

3.2 OCA的设计哲学:让窗口“手拉手”

OCA模块就是为了建立相邻窗口之间的直接对话通道。它的设计思想非常直观:既然窗口边界信息交互不足,那我就专门设计一个模块,让当前窗口的查询(Query)不仅关注自己内部的信息,也主动去“查询”和“关注”与之重叠的相邻窗口的特征。

具体是怎么实现的呢?我们避开复杂的数学公式,用流程来理解:

  1. 生成重叠区域特征:对于当前的输入特征图,OCA模块会使用一个展开(unfold)操作,为每个标准窗口(如16x16)生成一个更大的、与周边窗口有重叠的区域(例如20x20)的特征。这个更大的区域就包含了当前窗口及其周边上下左右的信息。
  2. 交叉注意力计算:然后,OCA执行一个交叉注意力计算。以标准窗口内的像素作为“查询(Q)”,而以这个更大的重叠区域内的所有像素作为“键(K)”和“值(V)”。
  3. 信息聚合:通过注意力计算,当前窗口的每个像素都会根据相似度,从重叠区域的所有像素(包括自身和邻居)那里聚合信息。这样,边界像素就能直接获取到邻居窗口的上下文,从而生成更连贯的特征。

这个过程就像让相邻的窗口不仅“背靠背”,还“手拉手”,在边界处进行了充分的磋商,确保大家对于边界该如何重建达成一致,从而有效消除了阻塞伪影。

3.3 为什么不全用OCA?效率的权衡

你可能会问,既然OCA这么好,为什么不把所有的注意力模块都换成OCA呢?原因很简单:计算成本

标准的窗口自注意力,其计算量大致与窗口内像素数的平方成正比。而OCA的交叉注意力,其查询来自标准窗口,但键值来自更大的重叠区域,计算量会显著增加。如果全部使用OCA,模型会变得非常笨重,训练和推理速度会让人难以接受。

因此,HAT采用了一种务实的设计:只在网络的部分深层阶段,稀疏地插入少数几个OCA模块。论文中发现,仅仅加入少量的OCA(比如在最后几个阶段),就能带来显著的性能提升,同时计算开销可控。这体现了很好的工程思维——将好钢用在刀刃上,在信息融合最需要、伪影最容易产生的地方,施加最直接的干预。

4. 从论文到实践:HAT的训练技巧与代码实战

看懂了原理,咱们手痒的工程师最关心的还是:这东西怎么用起来?HAT的成功不仅得益于新颖的结构,其训练策略也功不可没。这里结合官方的开源代码,分享一些实战要点。

4.1 同任务预训练策略

HAT论文中强调了一个训练技巧:相同任务预训练(Same-Task Pre-training)。这和我们常见的在ImageNet上做通用预训练不同。具体做法是,先用相对较小的模型(比如HAT-Light)在大型图像超分辨率数据集(如DIV2K)上训练到收敛。然后,将这个训练好的模型的权重,作为更大、更复杂模型(如HAT)的初始化起点。

我试过这个策略,效果确实比随机初始化或者用ImageNet权重初始化要好。我的理解是,超分辨率任务本身有很强的特异性,让模型先在“同宗同源”但更简单的任务上学习到一些基础的低分辨率到高分辨率的映射规律和特征提取能力,再让它去学习更复杂的映射,相当于提供了一个更优的起点,收敛更快,最终性能也更高。这有点像让运动员先进行基础体能训练,再专攻高难度技术动作。

4.2 关键代码片段解读

HAT的代码是开源的(GitHub上搜XPixelGroup/HAT),基于PyTorch实现,结构清晰。这里挑两个核心部分的代码,带大家感受一下:

首先是混合注意力块(HAB)中通道注意力的集成部分(简化示意):

class HAB(nn.Module):
    def __init__(self, dim, num_heads, window_size=16, ...):
        super().__init__()
        # 窗口自注意力层
        self.attn = WindowAttention(dim, window_size, num_heads, ...)
        # 前馈网络(FFN)中融入通道注意力
        self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
        # 注意:在FeedForward内部,通常会在两个全连接层之间加入通道注意力层

    def forward(self, x):
        # 残差连接 + 窗口自注意力
        x = x + self.attn(x)
        # 残差连接 + 前馈网络(内含通道注意力)
        x = x + self.ffn(x)
        return x

可以看到,通道注意力并不是一个独立的层,而是嵌入在标准Transformer块的前馈网络中的。你需要查看FeedForward类的具体实现,会发现它在扩展再压缩特征的过程中,加入了全局平均池化和全连接层来计算通道权重。

其次是重叠交叉注意力(OCA)模块中,如何生成重叠的键值对,这是其精髓:

# 假设 q 是标准窗口划分后的查询特征
# kv 是原始特征图经过一个卷积投影后得到的特征
# unfold 操作是关键,它用滑动窗口的方式提取重叠区域的特征块
kv_windows = self.unfold(kv) # 形状: (B, C*Kw*Kw, 窗口数量)
# 对kv_windows进行重组,使其便于与q_windows进行注意力计算
kv_windows = rearrange(kv_windows, 'b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch', 
                        nc=2, ch=C, owh=overlap_win_size, oww=overlap_win_size)
k_windows, v_windows = kv_windows[0], kv_windows[1] # 分别得到键和值

这里的unfold操作(在PyTorch中对应torch.nn.functional.unfold)就像用一个比标准窗口大的“刷子”,以一定的步长(小于刷子大小)滑动,从而提取出有重叠的区域特征。overlap_win_size就是重叠窗口的尺寸(大于标准窗口尺寸)。这样得到的k_windowsv_windows就包含了丰富的跨窗口上下文信息。

4.3 训练与调参经验

根据我个人和社区的经验,训练HAT这类大型SR模型有几个注意事项:

  1. 数据准备:高质量的数据集是关键。除了常用的DIV2K,结合使用Flickr2K、OST等更大规模的数据集进行预训练会有帮助。数据增强如随机旋转、翻转是标配。
  2. 损失函数:通常结合L1损失(保证像素级准确性)和感知损失(如VGG特征损失,保证视觉逼真度)。有的工作还会加入对抗性损失(GAN)来进一步提升纹理的真实感,但这会加大训练难度。
  3. 学习率与优化器:使用Adam或AdamW优化器,采用余弦退火或带热重启的学习率调度策略。对于HAT这样的模型,初始学习率不宜过大,一般在1e-4到5e-4之间试探。
  4. 硬件与批次大小:模型参数量较大,对显存要求高。在消费级显卡(如24GB显存)上,可能只能以较小的批次大小(如2-4)进行训练。可以考虑使用梯度累积来模拟更大的批次。

5. 效果对比与场景应用思考

说了这么多,HAT的实际效果到底怎么样?它在多个标准测试集(如Set5, Set14, Urban100, Manga109)上都取得了当时(CVPR 2023)最好的性能指标(PSNR/SSIM)。但指标只是一方面,从视觉上看,它的优势在于:

  • 纹理恢复更自然:对于动物毛发、建筑砖瓦、文字笔画等重复性纹理,HAT重建的结果更清晰、连贯,锯齿感更少。
  • 伪影抑制明显:特别是在处理有规则图案或高频细节的图片时,由窗口效应引起的块状伪影基本被消除。
  • 边缘更锐利:图像中物体的边缘轮廓保持得更好,过度平滑的情况有所改善。

那么,它适合用在什么场景呢?根据我的项目经验:

  • 老照片/经典影视修复:这是最直接的应用。HAT能够从低清、有损的源中恢复出令人惊喜的细节。
  • 医学影像增强:在允许的范围内,提升CT、MRI等影像的分辨率,可能帮助医生看清更细微的结构(需结合领域知识谨慎使用)。
  • 卫星与航拍图像:提升遥感图像的分辨率,用于城市规划、农业监测等领域。
  • 移动端图像后处理:虽然目前模型较大,但通过模型压缩、蒸馏等技术,其思想可以迁移到更轻量的模型中,用于手机相册的“超清放大”功能。

当然,HAT也不是没有缺点。它的模型复杂度高,推理速度相比一些轻量级CNN模型(如Real-ESRGAN的某些版本)要慢,对计算资源要求更高。在实际部署时,需要在效果、速度和资源消耗之间做权衡。

踩过几次坑之后,我最大的体会是,HAT的成功标志着图像超分辨率领域一个非常明确的趋势:从单纯的架构创新(CNN vs. Transformer),走向了更精细的模块化设计与融合创新。它没有彻底抛弃谁,而是思考如何让不同的优势机制协同工作。这种“混合”与“集成”的思路,或许比追求一个全新的、颠覆性的基础架构,在现阶段能带来更稳定、更有效的性能提升。对于从事相关研发的朋友来说,深入理解HAT的设计,不仅能帮你用好这个模型,更能启发你在其他视觉任务上,设计出属于自己的“混合注意力”方案。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐