ATD模型：基于令牌字典的高效图像修复Transformer架构

1. 项目概述：ATD模型的核心创新

在图像修复领域，Transformer架构近年来展现出超越传统CNN的性能优势，但其自注意力机制存在两个根本性瓶颈：一是计算复杂度随图像尺寸呈二次方增长，二是局部窗口注意力限制了感受野范围。ATD（Adaptive Token Dictionary）模型通过引入可学习的令牌字典系统，实现了三个关键突破：

1. 线性复杂度的全局建模 ：传统Transformer的自注意力计算需要处理所有像素点之间的两两关系（复杂度O(N²)）。ATD通过将图像特征与紧凑的字典特征交互，将复杂度降至O(N×M)，其中M是固定大小的字典容量（通常M<<N）

2. 内外先验融合机制 ：模型维护一个在训练过程中动态更新的令牌字典（典型设置M=256-1024），这些字典条目逐渐学习到各类典型图像结构的特征表达。通过创新的Token Dictionary Cross-Attention（TDCA）机制，输入图像可以与这些"视觉单词本"进行知识交互

3. 内容感知的特征分组 ：基于TDCA产生的注意力图，模型自动将图像区域按语义相似性划分为不同类别组，在组内实施自注意力计算。这种动态分组相比固定窗口划分更符合自然图像的稀疏相似性特性

实际测试表明，在4倍超分辨率任务中，ATD-light模型在Set5数据集上仅用15.6G MACs就达到32.18dB PSNR，比SwinIR节省47%计算量同时提升0.23dB。这种效率优势在处理2K/4K高分辨率图像时更为显著。

2. 核心架构设计解析

2.1 令牌字典的构建与学习

令牌字典D∈ℝ^(M×C)是ATD的核心组件，其学习过程体现以下设计考量：

1. 初始化策略 ：字典条目采用正交初始化，确保初始阶段各条目表征不同的视觉模式。实验显示，使用Modified Kaiming初始化比随机初始化最终PSNR提升约0.15dB

2. 动态更新机制 ：字典与模型参数一起通过梯度下降更新，但采用较低的学习率（通常为其他参数的1/5）。这种设计防止字典条目过快收敛到局部最优

3. 容量选择原则 ：字典大小M需要平衡表达能力和计算开销。对于256×256输入图像，M=512可在模型大小和性能间取得较好平衡（如图1所示）

表1：不同字典大小在Urban100数据集上的表现

2.2 三支路注意力架构

ATD的每个Transformer层包含三个并行的注意力分支（如图2所示）：

1. TDCA分支 ：处理图像→字典的交叉注意力

   • 使用降维后的查询（通常r=4）提升效率
   • 采用余弦相似度而非点积计算注意力
   • 引入对数缩放因子应对大字典稀释问题

2. AC-MSA分支 ：基于TDCA结果的类别内自注意力

   • 先按最大响应字典条目对图像区域分类
   • 每类再分固定大小的子组（典型ns=64）
   • 组内计算标准多头自注意力

3. SW-MSA分支 ：保留的局部窗口注意力

   • 窗口大小通常设为8×8
   • 补偿TDCA可能丢失的局部细节
   • 与Swin Transformer的窗口机制兼容

class ATDLayer(nn.Module):
    def __init__(self, dim, dict_size=512, num_heads=8, window_size=8):
        super().__init__()
        self.tdca = TokenDictCrossAttn(dim, dict_size)
        self.ac_msa = AdaptiveCategoryMSA(dim, num_heads)
        self.sw_msa = WindowMSA(dim, window_size, num_heads)
        self.ffn = CategoryAwareFFN(dim)
        
    def forward(self, x, token_dict):
        x_norm = self.norm1(x)
        tdca_out = self.tdca(x_norm, token_dict)
        ac_msa_out = self.ac_msa(x_norm, tdca_out.attn_map)
        sw_msa_out = self.sw_msa(x_norm)
        x = x + tdca_out + ac_msa_out + sw_msa_out
        x = x + self.ffn(self.norm2(x), tdca_out.category_idx)
        return x

3. 关键技术创新细节

3.1 令牌字典交叉注意力（TDCA）

TDCA机制的核心改进在于其稀疏化设计：

1. 对数缩放注意力 ：原始注意力计算修改为：

   A_{ij} = \frac{\exp(\tau'\cdot \text{sim}(q_i,d_j))}{\sum_k \exp(\tau'\cdot \text{sim}(q_i,d_k))}, \quad \tau'=1+\tau\log(M)
   

   其中τ是可学习参数，这种设计使得字典越大时注意力越稀疏。实验显示当M=512时，top-1注意力值比标准softmax高3-5倍

2. 跨层级字典共享 ：所有Transformer层共享同一字典，但通过独立的WK/WV投影矩阵实现层级特定表达。这种设计既节省参数又保持灵活性

3. 可视化分析 ：图3展示字典条目激活模式，可见不同条目确实捕捉到不同方向的边缘、纹理等基础视觉模式

3.2 自适应类别划分策略

基于TDCA结果的动态分组包含以下关键步骤：

1. 响应最大化分类 ：对每个空间位置(i,j)，选择使其注意力响应最大的字典条目：

   c_{ij} = \arg\max_k A_{ij}^k
   

2. 均衡子组划分 ：将同类别的token按固定大小ns分组，不足时循环填充。实验发现ns=64在并行效率和特征一致性间取得较好平衡

3. 位置编码保留 ：在组内计算自注意力时，保留原始相对位置编码，防止空间信息丢失。这比纯内容注意力提升约0.4dB

3.3 类别感知前馈网络（CFFN）

标准FFN扩展为：

\text{CFFN}(x_i) = \text{MLP}([x_i \| e_{c_i}])

其中e_ci是第ci个字典条目对应的嵌入向量。这种设计带来两方面优势：

1. 增强特征变换的类别特异性
2. 促进字典知识与局部特征的深度融合

消融实验显示CFFN对纹理恢复效果显著，在包含丰富高频细节的Urban100数据集上带来0.6dB提升。

4. 实现与优化技巧

4.1 训练策略优化

1. 渐进式字典学习 ：分三阶段训练：

   • 第一阶段（0-50k迭代）：冻结字典，仅训练其他参数
   • 第二阶段（50k-100k）：字典学习率设为其他参数的1/3
   • 第三阶段（100k后）：正常训练

2. 混合精度训练 ：对字典使用FP32精度，其他部分用FP16，平衡数值稳定性与训练速度

3. 注意力掩码正则化 ：对TDCA注意力图施加熵正则：

   L_{reg} = \lambda \sum_{i,j} A_{ij}\log A_{ij}
   

   防止注意力过度稀疏（λ=0.01）

4.2 推理加速技术

1. 字典条目剪枝 ：统计验证集上的注意力激活频率，移除长期未被激活的字典条目（约20%）

2. 类别缓存机制 ：对视频等连续输入，复用前一帧的类别划分结果，减少50%以上TDCA计算

3. 动态分辨率处理 ：对大尺寸输入，先在低分辨率进行粗略分类，再上采样细化，提升3-5倍速度

5. 实验结果分析

5.1 超分辨率性能对比

在DIV2K验证集上的4×超分辨率结果：

表2：与主流方法的性能对比

特别在纹理恢复方面，ATD展现出明显优势。如图4所示，对于建筑物外墙的规则网格结构，ATD能更好地保持线条的连续性和锐度。

5.2 多任务适应能力

ATD-U在去噪和JPEG伪影去除任务上也表现优异：

表3：在DND和JPEG-AI基准上的结果

6. 实际应用建议

1. 字典定制化 ：针对特定领域（如医学影像），可在预训练字典基础上进行领域自适应微调，通常只需1-2k迭代即可显著提升效果

2. 混合精度部署 ：字典部分保持FP16精度时，建议采用动态缩放因子防止数值下溢

3. 硬件适配 ：在部署到不同硬件时：

   • GPU：增大批次提升并行度
   • NPU：适当减少字典大小换取更高吞吐
   • 移动端：使用ATD-light并量化到8bit

在实际业务场景中，我们发现两个典型应用模式：

• 质量优先模式 ：使用完整ATD，开启所有注意力分支
• 效率优先模式 ：仅保留TDCA+SW-MSA，牺牲少量质量换取30%速度提升

对于需要处理4K视频的场合，建议采用空间分块处理，每块重叠32像素以避免边界伪影。