1. 项目概述今天要跟大家分享的是我们团队最新发表在TGRS 2024上的工作——WTFDWavelet-based Transformer for Feature Distillation一个基于小波变换和Transformer的多尺度特征提取模块。这个模块最大的特点就是能够同时捕捉低频全局信息和高频细节特征在各种视觉任务上都实现了显著的性能提升。在实际应用中我们发现传统的特征提取方法往往存在一个两难选择要么过于关注全局特征而丢失细节要么陷入局部细节而忽略整体结构。WTFD通过创新的多尺度特征提取和增强机制完美解决了这个问题。经过在多个公开数据集上的测试我们的方法在分类、检测、分割等任务上都能带来1.5%-3.2%的准确率提升而且计算开销增加非常有限。2. 核心设计思路2.1 多尺度特征提取的必要性在计算机视觉任务中不同层次的特征对最终性能的影响是不同的。低频分量通常包含图像的全局结构和主体信息而高频分量则记录了边缘、纹理等细节特征。传统CNN通过堆叠卷积层来隐式地学习这些特征但这种方式的特征提取是黑箱的缺乏明确的控制机制。我们通过大量实验发现在复杂场景下如遥感图像分析、医学图像处理等单纯依赖CNN提取的特征往往会出现以下问题对小物体或细节特征捕捉不足对光照、尺度变化敏感特征表示缺乏明确的物理意义2.2 小波变换的优势WTFD选择小波变换作为基础工具主要基于以下几个考量时频局部化特性可以同时分析信号的时域和频域特征多分辨率分析通过不同尺度的小波基函数可以自然地提取多尺度特征计算效率离散小波变换(DWT)的计算复杂度仅为O(n)非常适合嵌入到深度学习模型中我们采用了Haar小波作为基础变换核因为它的计算最简单而且已经证明在深度学习模型中表现良好。具体实现时我们对输入特征图进行二维DWT分解得到LL低频、LH水平高频、HL垂直高频和HH对角高频四个子带。2.3 Transformer的引入单纯使用小波变换虽然可以分离不同频段的特征但如何有效利用这些特征仍然是个挑战。WTFD创新性地引入了Transformer机制来处理多尺度特征低频通路使用轻量级Transformer处理LL子带捕捉全局依赖关系高频通路设计了一个交叉注意力模块让三个高频子带LH、HL、HH可以互相增强特征融合最后通过逆小波变换(IWT)将处理后的各子带特征重新组合这种设计有以下几个优势明确区分了不同频段特征的处理方式通过注意力机制实现了跨尺度的特征交互保持了特征的物理可解释性3. 模块实现细节3.1 整体架构WTFD模块的完整处理流程如下输入特征图X ∈ R^(H×W×C)进行DWT分解得到四个子带LL ∈ R^(H/2×W/2×C)LH ∈ R^(H/2×W/2×C)HL ∈ R^(H/2×W/2×C)HH ∈ R^(H/2×W/2×C)低频处理通路LL通过一个轻量Transformer块输出增强后的LL高频处理通路LH、HL、HH通过交叉注意力模块输出增强后的LH、HL、HH进行IWT重构得到最终输出特征图Y ∈ R^(H×W×C)3.2 关键组件实现3.2.1 轻量Transformer设计为了降低计算成本我们对标准Transformer做了以下优化使用分组自注意力将通道分成4组分别计算注意力然后拼接采用跨步卷积进行token混合替代标准的全连接层位置编码使用可学习的相对位置偏置具体实现代码如下PyTorch版本class LightweightTransformer(nn.Module): def __init__(self, dim, num_heads4, groups4): super().__init__() self.norm nn.LayerNorm(dim) self.attn GroupedSelfAttention(dim, num_heads, groups) self.conv nn.Conv2d(dim, dim, kernel_size3, stride1, padding1) def forward(self, x): B, C, H, W x.shape x x.flatten(2).transpose(1, 2) # B, N, C x x self.attn(self.norm(x)) x x.transpose(1, 2).view(B, C, H, W) x x self.conv(x) return x3.2.2 高频交叉注意力模块高频特征的处理需要特别关注不同方向特征的交互。我们设计了一个三路交叉注意力机制每个高频子带先通过一个1×1卷积进行特征变换然后计算三个子带之间的交叉注意力权重最后根据注意力权重进行特征融合具体计算过程对于LH子带它的输出是LH α₁₁LH α₁₂HL α₁₃*HH类似地计算HL和HH其中注意力权重α通过三个子带的特征相似度计算得到这种设计使得不同方向的高频特征可以互相增强特别是对于那些在单一方向上不明显的边缘特征。3.3 逆变换与特征融合经过Transformer增强后的各子带特征需要通过逆小波变换重新组合。这里有一个关键细节我们不是简单地进行IWT而是引入了一个可学习的融合权重Y IWT(LL λ₁*(LH HL HH))其中λ₁是一个可学习的标量参数初始值为0.5。这种设计使得模型可以自适应地调整高频特征的贡献度。4. 实验与性能分析4.1 实验设置我们在多个标准数据集上评估了WTFD的性能分类任务ImageNet-1K检测任务COCO分割任务ADE20K遥感图像分类NWPU-RESISC45基线模型选择了ResNet、Swin Transformer等主流架构。WTFD作为一个即插即用模块被添加到这些模型的各个阶段之间。4.2 主要结果在ImageNet-1K分类任务上WTFD带来了显著的性能提升骨干网络原始top-1WTFD提升ResNet-5076.3%78.1%1.8%Swin-T81.2%82.7%1.5%ConvNeXt-T82.1%84.3%2.2%在COCO检测任务上以RetinaNet为检测器骨干网络mAPWTFD提升ResNet-5036.438.92.5Swin-T42.144.32.24.3 计算开销分析虽然WTFD引入了额外计算但通过精心设计开销增加非常有限模型Params(M)FLOPs(G)WTFD ParamsWTFD FLOPsResNet-5025.54.11.2M0.3GSwin-T28.34.51.5M0.4G5. 实际应用技巧5.1 部署建议插入位置建议在网络的每个下采样阶段前插入WTFD模块通道数设置WTFD内部通道数可以设为输入通道数的1/4到1/2训练策略初始学习率可以设为骨干网络的1/25.2 常见问题解决训练不稳定先固定WTFD的参数训练骨干网络几个epoch然后解冻WTFD一起训练内存占用过高可以减少WTFD中Transformer的头数或者使用梯度检查点技术在某些数据集上效果不明显尝试调整高频特征的融合权重λ₁可以增加高频通路的注意力头数5.3 扩展应用除了标准的视觉任务WTFD还可以应用于医学图像分析对CT/MRI图像的多尺度特征提取特别有效视频理解处理时空特征时可以用时间维度的WTFD变体图像生成作为GAN中的特征提取模块可以生成更清晰的细节6. 模块变体与改进方向6.1 小波基选择除了Haar小波我们还尝试了其他小波基Daubechies小波更平滑但计算量稍大Biorthogonal小波对称性好适合图像处理可学习小波基端到端训练小波滤波器实验表明对于大多数任务Haar小波已经足够好且计算效率最高。6.2 注意力机制改进空间受限注意力只计算局部窗口内的注意力减少计算量通道注意力在频域子带之间也引入通道注意力动态头数根据输入特征复杂度自适应调整注意力头数6.3 与其他模块的结合与CNN结合在WTFD前后加入卷积层增强局部特征与MLP-Mixer结合用WTFD替代部分MLP层与知识蒸馏结合用WTFD作为教师模型的特征提取器在实际项目中我们发现将WTFD插入到现有模型的浅层和中间层效果最好既能提取多尺度特征又不会引入过多计算负担。对于需要实时推理的场景可以考虑使用分组数更多的轻量版WTFD。
