Transformer自编码器在视觉任务中的设计与优化
1. 基于Transformer的自编码器架构解析在计算机视觉领域自编码器(RAE)作为一种无监督学习框架通过编码器-解码器结构实现数据的高效表示学习。近年来随着Vision Transformer(ViT)架构的兴起基于Transformer的自编码器展现出比传统CNN架构更强大的特征提取能力。我们采用的DINOv2 with Registers架构在多个视觉任务中表现出色。1.1 编码器归一化处理细节对于任何预训练好的表示编码器我们首先丢弃编码器输出的[CLS]或[REG]等特殊token仅保留所有的patch tokens。随后对每个token独立进行层归一化(LayerNorm)确保各token在通道维度上具有零均值和单位方差。这里需要特别说明的是我们使用的所有表示编码器都采用标准ViT架构其最后一个transformer block后已经应用了层归一化。因此我们只需要取消表示编码器中层归一化的仿射参数(affine parameters)。这一操作不会影响编码器的表示质量因为它只是一个线性变换。实际应用中发现使用DINOv2时需要注意其默认patch大小为14。为了生成256×256分辨率图像我们将输入图像插值到224×224但设置目标patch尺寸为16。这种配置确保模型仍能产生256个token(224/1416→16×16256)同时可以重建256×256图像。1.2 解码器架构设计要点解码器的核心任务是将编码器产生的token嵌入重建回像素空间。我们采用与编码器相同的patch大小因此可以生成与编码器输入相同空间分辨率的图像。借鉴He et al.(2021)的设计我们在解码器输入序列前添加一个可学习的[CLS]token并在解码完成后将其丢弃。解码器训练采用ImageNet-1K数据集大多数实验在256×256分辨率下进行。对于512分辨率合成(不使用解码器上采样的情况)我们直接在512×512图像上训练解码器。判别器架构选择我们主要遵循StyleGAN-T的设计选择但做了两个关键改进使用冻结的Dino-S/8而非Dino-S/16作为判别器。实验表明Dino-S/8能稳定训练避免解码器生成对抗性patch。移除了StyleGAN-T中的虚拟批归一化(virtual batch norm)改用标准批归一化。所有判别器输入都会先插值到224×224分辨率。下表展示了解码器和判别器的训练配置组件解码器判别器优化器AdamAdam最大学习率2×10⁻⁴2×10⁻⁴最小学习率2×10⁻⁵2×10⁻⁵学习率调度cosine衰减cosine衰减优化器betas(0.5, 0.9)(0.5, 0.9)权重衰减0.00.0批大小512512预热1 epoch1 epoch2. 噪声增强的解码器训练策略2.1 复合损失函数设计解码器训练采用L1损失、LPIPS损失和对抗损失的组合L_rec(x) ω_L LPIPS(x̂, x) L1(x̂, x) ω_G λ GAN(x̂, x)其中ω_L1.0ω_G0.75。对抗损失采用Esser et al.(2021)提出的自适应权重λ来平衡重建损失和对抗损失的尺度λ ∥∇_{x̂} L_rec∥ / (∥∇_{x̂} GAN(x̂, x)∥ ϵ)这种设计确保了在训练过程中模型不会过度偏向于某一种损失保持了生成图像的视觉质量和细节保真度之间的平衡。2.2 数据增强方案我们的数据增强流程包含两个关键步骤先将输入图像调整到384×384然后随机裁剪到256×256在输入判别器前应用Zhao et al.(2020)提出的可微分增强使用默认超参数这种组合增强策略有效提升了模型的泛化能力避免了过拟合问题。2.3 噪声增强的实证效果通过系统实验我们发现噪声增强训练(noise-augmented training)能显著提升生成质量。下表展示了不同编码器在使用噪声增强前后的性能对比模型gFID(原始)gFID(噪声增强)rFID(原始)rFID(噪声增强)DINOv2-B4.814.280.490.57SigLIP2-B6.694.930.530.82MAE-B16.148.380.160.28可以看到虽然噪声增强轻微降低了重建质量(rFID上升)但它显著改善了生成质量(gFID下降)。这种现象在性能较弱的编码器(如MAE-B)上尤为明显。这表明噪声增强使解码器学会了更平滑的映射能够更好地处理不完美的潜在表示。3. DiTDH架构实现细节3.1 模型配置我们默认使用LightningDiT作为扩散模型的骨干架构。采用连续时间形式的flow matching将时间步输入限制在[0,1]区间。借鉴Song et al.(2021)的做法我们用高斯傅里叶嵌入层替代时间步嵌入并在输入token中除了RoPE外还添加了绝对位置嵌入(APE)尽管实验表明APE对性能影响不大。对于DiTDH架构我们基本遵循DiT的设计但不为DDT头输入重新应用APE。当DiT和DDT头的维度不匹配时使用线性层进行映射。下表展示了不同规模的模型配置模型维度头数深度S384612B7681212L10241624XL11521628XXL12801632H15361632G20481640T268821403.2 优化策略对于DiT模型我们严格遵循LightningDiT的优化策略使用AdamW优化器(恒定学习率2.0×10⁻⁴)、批大小1024、EMA权重0.9999。对于DiTDH我们发现直接应用这一策略会导致后期训练出现损失尖峰因此调整为学习率从2.0×10⁻⁴线性衰减到2.0×10⁻⁵恒定预热40个epochEMA权重从0.9999调整为0.9995添加梯度裁剪(阈值1.0)所有模型训练80个epoch除非另有说明。我们只报告EMA模型的性能。3.3 DDT头设计原则实验表明DDT头的设计有两个关键原则宽度优先于深度宽而浅的头比窄而深的头表现更好。例如2层2048维(G)的头部优于6层1152维(XL)的头部尽管它们的计算量相似。与编码器规模匹配RAE编码器越大DDT头的最佳宽度也越大。对于DINOv2-S和DINOv2-B性能在DDT头宽度2048(G)时收敛而对于DINOv2-L2688(T)的头部仍能带来性能提升。基于这些发现我们默认使用2层2048维的DDT头。4. 采样与评估方法论4.1 采样策略我们默认使用欧拉采样器的标准ODE采样50步。实验表明超过50步后性能基本收敛。对于DiT和DiTDH使用相同的采样超参数。在FID-50K评估中我们遵循(Tian et al., 2024)等工作的协议从每个类别采样50张图像共50,000张。参考统计量取自ADM预计算的完整ImageNet数据集统计量。值得注意的是另一种常见的采样策略是均匀采样类别标签50K次。虽然随机采样在理论上会收敛到每类采样版本但实践中两种策略得到的FID分数略有不同。为确保公平比较我们用类平衡采样重新评估了之前的最先进方法。4.2 评估指标详解我们严格遵循ADM的评估设置使用相同的参考批次。主要评估指标包括FID评估真实图像和生成图像特征分布之间的距离。使用Inception-v3网络假设两个分布都遵循多元高斯分布。IS同样使用Inception-v3但直接评估logits。测量边缘标签分布与softmax归一化后条件标签分布之间的KL散度。精确率和召回率精确率反映生成图像中看起来真实的比例召回率反映生成样本覆盖的训练数据流形的比例。4.3 无条件生成结果我们也研究了RAE在无条件生成中的表现。训练时将所有标签设为null生成时使用相同的null标签。虽然在这种情况下无法使用分类器无关指导(CFG)但自动指导(AutoGuidance)仍然适用。下表展示了无条件生成的结果对比方法gFID ↓IS ↑DiT-XL VAE30.6832.73DiTDH-XL DINOv2-B4.96123.12RCG DiT-XL4.89143.2我们的方法比在VAE潜在空间上训练的DiT-XL表现好得多与专门为无条件生成设计的RCG方法相比也展现出竞争力同时实现更加简单直接。5. 实际应用中的经验总结在项目实施过程中我们积累了一些关键经验值得与同行分享编码器选择虽然MAE在重建任务上表现出色(rFID低)但其生成性能(gFID)明显不如DINOv2。这表明低rFID并不一定意味着好的图像tokenizer。DINOv2-B在各方面表现均衡是我们的默认选择。训练稳定性使用Dino-S/8作为判别器比Dino-S/16更稳定能有效防止解码器生成对抗性patch。这是通过大量对比实验得出的重要发现。计算资源分配对于基于RAE的模型我们使用patch大小为1对于VAE和像素输入的基线实验则分别使用patch大小2和16。在所有256×256实验中扩散模型处理的token序列长度保持256不变因此DiT骨干的计算成本在不同设置下几乎相同。指导方法选择AutoGuidance比CFG更易调节且性能更好。我们发现较弱的基模型和早期训练检查点能产生更强的指导效果。实践中我们使用最小的DiTDH-S变体作为指导模型通常采用早期检查点。可视化分析通过重建样本的可视化比较(图8)可以直观评估不同RAE的性能差异。这种定性分析往往能揭示定量指标无法反映的细节特征。