Command Palette
Search for a command to run...
FuseReg:正则化层融合缓解表示自编码器中的重建-生成差距
FuseReg:正则化层融合缓解表示自编码器中的重建-生成差距
摘要
表示自编码器(RAE)将预训练视觉编码器的特征重用为重建和扩散潜变量,从而将强大的视觉表示融入图像生成。然而,RAE 仍需要决定哪些编码器层构成生成器和像素解码器的共享潜在空间。这一选择涉及权衡:较浅层往往能更好地保留细粒度像素细节,而较深层往往能带来更好的生成指标。因此,固定的启发式层融合会把两个从不同信息中受益的阶段耦合在一起。我们提出 FuseReg,用对编码器层随机子集进行训练的方式替代启发式特征选择。我们从理论上分析了其底层机制:子集采样显式地惩罚对跨层不一致的敏感性。在采用 DINOv3-L 的 ImageNet-256 上,单个 FuseReg 解码器无需重新训练即可从全量、稀疏和单层融合中进行重建,并获得比专用于固定融合的解码器更高的 PSNR。这种灵活性同样有益于生成:仅替换解码器即可在 RAEv2 DiT-XL 生成器保持不变的情况下将无引导 gFID 降低 27%。同样的正则化原理还可扩展到扩散训练,对两个阶段进行联合正则化在 DiT-Base 上将无引导 gFID 降低 29%。这些结果表明,针对层融合鲁棒性训练下游模型,可以在不修改预训练编码器的情况下缩小重建-生成差距。
一句话总结
南加州大学 PSI 实验室的研究人员与合作者提出 FuseReg,它通过在编码器层的随机子集上训练来正则化表示自编码器,以惩罚跨层不一致,从而使单个解码器无需重新训练即可从完整、稀疏和单层融合中重建,同时提升 PSNR,并将无引导 gFID 最多降低 29%。
核心贡献
- FuseReg 是一种层融合正则化器,它在随机采样的编码器层的归一化融合上训练表示自编码器解码器和生成器,取代固定的启发式层选择。
- 理论分析表明,随机子集采样保留了全层均值,同时使跨层不一致成为显式的正则化惩罚,并建立了与确定性全局融合方法的二阶分离。
- 在带 DINOv3-L 的 ImageNet-256 上,单个 FuseReg 解码器无需重新训练即可从完整、稀疏和单层融合中重建,并比专门用于固定融合的解码器获得更高 PSNR;仅替换解码器就在 k=23 时将无引导 gFID 从 3.01 降至 2.21,而联合正则化则将 DiT-Base 的 gFID 从 13.96 降至 9.93。
引言
作者研究表示自编码器(RAE),其中冻结的视觉编码器定义扩散模型的潜在空间,但其层级特征必须通过层融合压缩为单个潜在表示。这个接口很关键,因为解码器从生成的潜在表示重建像素,而生成器建模潜在分布;先前固定融合选择(如 RAEv2)会产生重建与生成之间的差距:重建偏好浅层像素对齐层,而扩散偏好更深语义层。例如,在 DINOv3-L 设置下,将融合子集从 k=7 扩展到 k=23 会使重建 PSNR 从 22.58 提升到 27.04 dB,但会使无引导 gFID 从 1.65 恶化到 3.01。作者提出 FuseReg,一种层融合正则化方法,在随机采样的归一化编码器层子集上训练解码器和生成器。这在期望上保留全层表示,同时减少对任何单一层组合的依赖,产生的单个解码器无需重新训练即可支持完整、稀疏和单层融合,并在不改变架构或增加额外计算成本的情况下提高生成质量。
方法
作者通过提出 FuseReg 来解决表示自编码器中固定层融合的局限性,该方法将随机层融合用作正则化技术。在传统设置中,固定融合规则或学习的全局门控只让下游模型在训练期间看到单一层组合。然而,学习到的门控通常未能有效利用深层语义特征,反而几乎完全依赖浅层像素对齐层。
如下图所示:
为了克服这种对固定捷径的依赖,作者用层子集上的均值保持分布取代确定性融合权重。给定图像和冻结编码器,来自 K 个选定层的层归一化 tokens 被映射到单个潜在表示。标准归一化融合使用固定权重 a,使得 za=∑k=1Kakhk 且 1⊤a=1,而 FuseReg 为每个训练样本采样单独的层丢弃掩码。对于丢弃率 p∈[0,1),该方法抽取独立同分布的 Bernoulli 掩码 m,并将其条件限制为非零。所得融合潜在表示计算为:
zm=∑kmk∑kmkhk这种归一化很关键,因为它确保采样潜在表示的期望值等于全层均值。因此,采样引入的变化严格针对编码器层不一致的方向,而不移动平均表示。
该框架将这种随机融合同时应用于像素解码器和扩散 Transformer,但为每个阶段采用不同的丢弃率,因为它们解决的是根本不同的预测问题。对于解码器,子集融合用于重建像素。解码器被训练为将采样潜在表示 zm 映射回图像,使用像素、感知和对抗目标的组合。这迫使解码器从多种层组合重建图像,减少对浅层捷径的依赖。
对于扩散 Transformer,生成器接收带噪声的子集融合,但被训练为预测全层聚合目标。使用由子集聚合构建的流匹配插值,模型通过时间加权预测损失回归全层目标。将丢弃率设为零可恢复标准全融合目标,而任何正丢弃率都会训练生成器将多个子集融合映射到相同的全层目标。
理论分析证实,这种归一化子集采样平均上保留所有层表示,同时在层不一致的分量上精确增加方差。此外,这些随机融合权重的二阶矩张成整个层对比子空间。由于确定性输入无关融合规则的二阶矩秩为一,没有固定的全局融合能同时匹配随机化方法的均值和方差特性,这使得 FuseReg 成为弥合重建与生成差距的独特有效正则化器。
实验
评估使用冻结的 DINOv3-L 编码器和在 ImageNet-256 上训练的 ViT 解码器,并使用匹配的 DiT-Base 和 DiT-XL 预算;评估通过在不同层融合上测试同一个解码器、在固定生成器和潜在表示时更换解码器、以及联合正则化两个阶段来隔离解码器、生成器和联合效应。实验表明,随机融合解码器在完整、子集和单层融合上都能稳健重建,而固定融合解码器则专门化并在偏移融合上性能下降;提高解码器稳健性也会在生成器固定的情况下改善生成,尤其是对于偏移融合。联合正则化在 DiT-Base 上带来互补增益,但在 DiT-XL 上表现更微妙且依赖指标,因此需要根据规模、评估指标和引导情况分别选择解码器和生成器正则化率。
实验比较了完整、子集和单层融合上的重建质量。固定融合 RAEv2 解码器仅在其训练融合上表现良好,转移到其他融合时急剧下降。单个 FuseReg 解码器在非平凡融合上训练,在所有评估融合类型上保持稳定强劲的 PSNR、SSIM 和 rFID。固定融合 RAEv2 解码器严重专注于其训练融合,在非训练融合上 PSNR 和 rFID 大幅下降。FuseReg 解码器在所有测试融合上达到或超过确定性解码器,每种情况下保持 rFID 低于 0.6 且 SSIM 高于 0.67。
在这次无引导 ImageNet-256 Inception Score 扫描中,解码器正则化具有非单调效应:小解码器率比固定融合基线有所改善,但较大解码器率会降低分数。生成器正则化通常降低 Inception Score,因此最强结果来自无或极少生成器正则化与适度解码器率的组合。最佳 Inception Score 出现在生成器率为零且解码器率较小的位置,相比复现的固定融合基线有显著改善。对于每个生成器率,分数往往在解码器率为 0.05 或 0.1 时达到峰值,然后随解码器正则化增加而下降。高生成器正则化与高解码器正则化结合产生扫描中的最低分数。
第一个实验评估完整、子集和单层融合上的重建质量,表明固定融合 RAEv2 解码器对其训练融合过拟合,并在转移时急剧下降,而在非平凡融合上训练的单个 FuseReg 解码器保持稳定强劲的 PSNR、SSIM 和 rFID。第二个实验在无引导 ImageNet-256 上扫描解码器和生成器正则化并使用 Inception Score,发现轻度解码器正则化比固定融合基线有所改善,但更大正则化量会造成损害,且生成器正则化通常降低分数。最佳设置使用无或极少生成器正则化与适度解码器率,而两侧高正则化产生最低分数。