FuseReg：以层融合正则化缓解表示自编码器中的重建-生成鸿沟 FuseReg: Regularizing Layer Fusion Mitigates the Reconstruction-Generation Gap in Representation Autoencoders
训练时随机采样编码器层子集做融合，让解码器与生成器对层组合鲁棒，弥合重建-生成差距。
前置知识
表示自编码器（RAE）
用冻结的预训练视觉编码器（如 DINOv3）定义扩散模型的潜空间，再训练一个解码器把编码器特征映射回像素。与重建导向的 VAE 不同，RAE 的潜空间天然带有语义结构，让扩散 Transformer 更容易训练，同时保持像素可重建。
本文的全部问题都发生在 RAE 的"编码器特征 → 潜空间 → 解码器/生成器"接口上，理解 RAE 的两阶段结构是读懂动机的前提。
层融合（Layer Fusion）
视觉 Transformer 各层特征内容随深度变化：浅层贴近像素细节，深层偏语义。层融合指把选定的若干层特征（如 RAEv2 对最后 k 层求和）按权重组合成一个单一潜变量，供解码器重建或生成器建模。融合权重 $a_k$ 满足归一化约束 $\mathbf{1}^\top a = 1$。
层融合是解码器和生成器共享的接口定义，选哪些层直接决定潜空间的性质，是本文要正则化的对象。
重建与生成指标（PSNR/rFID 与 gFID/IS）
PSNR 衡量重建像素保真度（越高越好），rFID 衡量重建图像与真实图像的分布距离；gFID 和 IS 则在 50k 个类条件生成样本上评估生成质量。本文用 50 Euler 步采样，无引导与内部引导（scale 1.78）两种设置分别报告。
重建-生成鸿沟正是通过这组指标的背离来定义的：k 越大 PSNR 越高但 gFID 越差，读者需要理解这组指标才能看懂所有实验表。
Flow Matching 与 DiT
扩散 Transformer（DiT）用 flow matching 插值 $z_t = t z_m + (1-t)\epsilon$ 从噪声走向数据，训练目标是时间加权的 x-prediction 损失，权重 $\rho(t) = \max(1-t, 0.05)^{-2}$，时间服从 shift $a=8$ 的 logit-normal 分布。
FuseReg 把生成器的训练目标改成"从带噪子集融合预测全层融合"，公式推导和 $t^2$ 因子都依赖 flow matching 的记号。
Dropout 与 Nested Dropout
Dropout 随机置零单元做正则化；nested dropout 随机保留一个前缀以获得有序（PCA 式）表示；LayerDrop、slimmable networks、Matryoshka 表示学习则把随机子配置训练用于推理提效。共同点是抽取的子配置都比完整模型小。
FuseReg 借用了随机子集训练这一机制家族，但把它搬到"编码器层融合"这个新轴上且目的不同，理解这一脉络才能分清它与先前工作的本质差异。
内部引导（Internal Guidance / REPA）
REPA 在 DiT 中间层用投影头对齐预训练编码器表示。RAEv2 的内部引导把完整 DiT 输出与 REPA 头预测组合：$\hat{x}_{guided} = \hat{x}_{full} + w(\hat{x}_{full} - \hat{x}_{repa})$，引导强度 $w$ 会放大两个预测对层子集响应的差异。
有引导与无引导下最优丢弃率不同的实验现象，正是通过引导组合中 $\mathrm{tr}(W_{full}\Gamma W_{repa}^\top)$ 交叉项解释的，是讨论部分的关键。
研究动机
表示自编码器（RAE）用冻结的 DINOv3-L 等视觉编码器定义扩散模型的潜空间，但编码器产出一整层级的特征，必须决定哪些层组成解码器与生成器共享的潜空间。现有 RAEv2 采用固定启发式：对最后 $k$ 层求和。这个 $k$ 实际上是一条重建-生成的帕累托前沿：在 ImageNet-256 + DINOv3-L 上，$k=7$ 时重建 PSNR 只有 22.58 dB，但无引导 gFID 低至 1.65、有引导 gFID 1.06；把融合扩展到 $k=23$ 后 PSNR 升到 27.04 dB，无引导 gFID 却恶化到 3.01、有引导 gFID 升到 1.25。浅层特征贴合像素、直接喂饱 $\ell_2$ 与 LPIPS 重建损失，深层语义特征则更容易被扩散模型建模，而一个固定融合把预测目标完全不同的解码器和生成器绑死在同一个选择上。作者还实验了可学习的全局 softmax 门控（Figure 1），结果门控权重在训练中坍缩到最浅层、深层权重衰减到 0，且 GAN 损失开启后坍缩加速——像素重建损失天然偏好浅层捷径，连"学一个融合"也靠不住。
本文的目标是本文的目标不是再去寻找另一个更好的固定融合，而是让下游模型对任意层融合配置都保持有效，且完全不修改冻结的预训练编码器。具体拆成三个可验证的目标：其一，训练出的单个解码器应当能在全层融合、稀疏子集融合和单层融合下都可靠重建，而不需要为每种配置分别重训；其二，这种解码器鲁棒性应当转化为生成质量的提升——在不重新训练生成器、不改变已采样潜变量的前提下，仅替换解码器就能降低 gFID；其三，把同一正则化原则扩展到扩散训练本身，并给出理论刻画：随机化融合到底改变了什么训练信号、为什么它无法被任何确定性融合复制。最终在 DiT-Base 和 DiT-XL 两个规模上量化对 PSNR/SSIM/rFID 与 gFID/IS 的影响。
与已有工作不同的是，与本文最相关的两条线索各有局限：多层融合方法（RAEv2 的固定前 $k$ 层求和、DRoRAE 的深度路由融合模块、注意力多层探测）虽然在"选哪些层、怎么加权"上做文章，但都在收敛到某一个确定性的组合规则；而 dropout/LayerDrop/Matryoshka 等随机配置训练虽然引入随机性，目的却是正则化或推理提效，且抽取的子配置总比完整模型小。本文的独特切入是把层融合本身当作一个训练分布：每个样本随机采样层子集并做归一化（除以保留层数），使期望严格保持全层均值 $\bar{h}$ 不变，随机扰动只落在层间分歧方向上。理论上证明（Proposition 2）随机子集融合权重的二阶矩秩为 $K$，而任何确定性全局融合（包括可学习门控）的二阶矩秩至多为 1，二者在线性平方损失目标函数上本质不可互替——这从机制上解释了为什么可学习门控会坍缩到浅层捷径，而随机化天然覆盖全部 $2^K-1$ 个非空子集。
核心方法
直觉上，与其教模型只用一种层组合，不如让它训练时见过所有层组合。技术上 FuseReg 非常轻：给定冻结编码器 $K=23$ 层的 layer-normalized token 特征 $h_k = \mathrm{LN}(E_{l_k}(x))$，训练时对每个样本独立抽取 Bernoulli 掩码 $m_k \sim \mathrm{Bernoulli}(1-p)$ 并条件化于非空，构造归一化子集均值 $z_m = \sum_k m_k h_k / \sum_k m_k$。由于线性期望 $\mathbb{E}[z_m \mid x] = \bar{h}$（全层均值），部署时默认的全层融合不会因训练分布而偏移；随机性只沿层间分歧方向注入方差，强度由丢弃率 $p$ 控制。这个分布被施加到 RAE 的两个预测阶段：解码器以率 $p_{dec}$ 从随机子集融合重建像素（像素 + LPIPS + GAN 损失）；DiT 以率 $p_{dit}$ 从带噪子集融合预测全层融合目标 $z_{full}$（flow matching 的 x-prediction 损失）。两个阶段 solving 的问题不同，所以丢弃率分开设置。整个过程零额外推理开销、零架构改动。
核心创新是把"对特定层组合的依赖"本身识别为病灶并直接正则化，这与已有方法的本质区别有三层。第一，不搜索最优固定融合，而是让训练分布覆盖全部 $2^K-1$ 个非空子集（Corollary 1 给出概率加权覆盖保证：平均损失小则每个子集的损失都被上界控制）；第二，归一化使采样无偏——Lemma 1 证明 $\mathbb{E}[z_S \mid x] = \bar{h}$ 且协方差恰为 $c_K(p)V(x)$，其中 $V(x) = \frac{1}{K}\sum_k e_k e_k^\top$ 只含层间分歧方向——因此 $p_{dec}=0.95$ 绝不意味着"摧毁 95% 的信号"，而是把分歧协方差放大到接近上限；第三，Proposition 1 证明在线性平方损失下，子集采样严格等价于给"输出随层间分歧变化"的预测器加显式惩罚 $c_K(p)\,\mathrm{tr}(D\Gamma D^\top)$，而 Proposition 2 证明随机子集权重的二阶矩秩为 $K$、任何确定性全局融合至多秩 1，二者目标函数不可互替。可学习门控会坍缩到浅层捷径，随机化从数学结构上排除了这种坍缩。
方法步骤详情
完整流程分四步。第一步，特征抽取：冻结 DINOv3-L，取 $K=23$ 个候选 Transformer 层，每层做 layer normalization 得到 $h_k \in \mathbb{R}^{N \times d}$。第二步，随机融合：每个训练样本抽取 Bernoulli 掩码 $m_k \sim \mathrm{Bernoulli}(1-p)$，条件化于至少保留一层，输出归一化子集均值 $z_m$。第三步，解码器训练（率 $p_{dec}$）：ViT 解码器从零训练 16 epochs，损失 $\mathcal{L}_{dec} = \|\hat{x}-x\|^2 + \lambda_p \mathcal{L}_{LPIPS} + \lambda_g \mathcal{L}_{GAN}$，输入是随机子集融合，输出重建像素，另含分类器代理项。第四步，生成器训练（率 $p_{dit}$）：DiT-Base/XL 训练 40 epochs，flow matching 插值 $z_t = t z_m + (1-t)\epsilon$，时间从 shift $a=8$ 的 logit-normal 抽取，训练目标为时间加权的 x-prediction 损失 $\mathcal{L}_{DiT} = \mathbb{E}\,\rho(t)\|x_\theta(z_t,t,c) - z_{full}\|^2$，权重 $\rho(t)=\max(1-t,0.05)^{-2}$，full 头与 REPA base 头共用该损失。推理时默认全层融合（$z_m = \bar{h}$），再在 $(p_{dec}, p_{dit}) \in \{0, 0.05, ..., 0.95\}^2$ 网格上评估：50k 图像测 PSNR/SSIM/rFID，50k 样本（50 Euler 步）测 gFID/IS。
技术新颖性
技术新颖性体现在四处。（1）轴的迁移：把随机子集机制从单元/层块/维度迁移到"冻结编码器的层融合"这一全新轴，且目的不是缩小计算图而是让下游模型覆盖所有信息视角——抽取的子配置不是更小的计算路径，而是不同的表征内容。（2）精确的统计刻画而非启发式：Lemma 1 的均值保持与分歧协方差恒等式、$c_K(p) \in [0,1]$ 单调有界、Proposition 2 的二阶矩秩分离，都是不依赖线性假设的精确结果；Proposition 4 还构造性证明"全层融合零损失的解码器可以在只差一层的单子集上损失任意大"，论证了固定融合不可迁移。（3）双阶段差异化丢弃率：$(p_{dec}, p_{dit})$ 是每层独立率 relaxation 的结构化二维切片，实验证明两轴互补且超可加（DiT-Base 上联合改善 4.03，超出单轴之和 2.87 达 1.16）。（4）对内部引导的耦合分析：引导组合 $W_{guided} = (1+w)W_{full} - w W_{repa}$ 的分歧惩罚含交叉项 $\mathrm{tr}(W_{full}\Gamma W_{repa}^\top)$，首次给出"有引导/无引导下最优率不同"的机制性解释。
实验结果
论文用三个递进实验逐一验证。（1）单解码器跨融合鲁棒（Table 1/Table 4）：固定融合解码器在偏离训练融合时崩溃——RAEv2-K=23 在 $k=7$ 和单层 $\ell_{11}$ 上分别掉到 12.51 和 14.31 dB、rFID 恶化到 16.10 和 11.21；而一个 FuseReg（$p_{dec}=0.95$）解码器在 $k=7/k=23/\ell_{11}$ 三种融合下分别达到 23.77/27.52/25.13 dB，rFID 全部低于 0.61，且在 $k=23$ 上以 27.52 dB 超过专用解码器的 27.04 dB。Figure 2 的余弦相似度图显示 FuseReg 在所有融合下保持干净的空间结构。（2）仅换解码器提升生成（Table 5/Figure 3）：固定同一 DiT 与同一批 50k 潜变量，$k=23$ 无引导 gFID 从 3.01 降至 2.21（降 27%），有引导 gFID 保持 1.25（中间率一度回退到 1.49）；偏移到 $k=7$ 融合时无引导 gFID 从 27.73 降至 1.92、有引导从 16.35 降至 1.42，IS 从 97.7 升至 221.8。（3）联合正则化跨规模（Table 2）：DiT-Base 基线 gFID 13.96，解码器单独（$p_{dec}=0.9$）到 12.96、生成器单独（$p_{dit}=0.7$）到 12.09、联合（0.9, 0.7）达 9.93，超出可加外推 1.16，IS 从 107.5 升至 125.7；DiT-XL 基线 gFID 2.91，最佳 2.38 出现在（$p_{dit}=0, p_{dec}=0.95$），生成器正则单独无助于 gFID，但 IS 最高 221.1（基线 207.4）需要双高率。（4）有引导下 DiT-XL 最低 gFID 1.50 在（0.9, 0），而最高 IS 262.6 反而在（0,0）——最优率随规模、指标与引导设置而变。（5）机制诊断（Figure 4/6）：FuseReg 的 leave-one-out 敏感度曲线更平坦，单层+ℓ23 代理读出下全层 PSNR 27.93 dB vs RAEv2 的 27.36 dB，Shapley 排序大体保留但各层不再不可或缺，小子集读出更强、更早饱和。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-256 跨融合重建（k=23 全层融合） | PSNR | 27.52 dB（FuseReg, $p_{dec}$=0.95） | 27.04 dB（RAEv2-K=23 专用解码器） | +0.48 dB |
| ImageNet-256 跨融合重建（单层 ℓ11） | PSNR | 25.13 dB | 17.50 dB（RAEv2-K=7）/ 14.31 dB（RAEv2-K=23） | +7.6 至 +10.8 dB |
| 类条件生成（仅替换解码器，k=23，无引导） | gFID | 2.21 | 3.01（RAEv2 解码器） | −26.6% |
| 类条件生成（解码器偏移到 k=7 融合，无引导） | gFID | 1.92 | 27.73 | −93.1% |
| DiT-Base 联合正则化（无引导） | gFID | 9.93（$p_{dec}$=0.9, $p_{dit}$=0.7） | 13.96（复现 RAEv2 基线） | −28.9% |
| DiT-XL 无引导生成 | gFID | 2.38（$p_{dec}$=0.95, $p_{dit}$=0） | 2.91 | −18.2% |
局限与改进
作者明确承认的局限：实验范围限于 ImageNet-256 + 冻结 DINOv3-L + DiT-Base/XL 且训练预算匹配，向其他编码器层级结构、分辨率、数据域和更长训练调度的迁移留待验证；解码器与生成器对正则化的响应不同，最优丢弃率依赖模型规模、评估指标和引导设置，在新场景应用时必须分别率选；理论分解只在齐次线性预测器 + 平方损失下精确，对非线性模型仅提供机制层面直觉，混合重建损失和完整引导采样轨迹的刻画仍未解决。我自己的观察还有四点：其一，率网格搜索成本高——DiT-XL 需要 $4 \times 8 = 32$ 个配置各训 40 epochs，论文未提供廉价率选择方法；其二，有引导下中间率会出现暂时回退（1.25→1.49），引导与正则化的交互只有固定 $t$ 的线性交叉项解释，缺乏系统校准；其三，最强结论多在自建复现基线上得出，未与 DRoRAE 等同期方法在同一协议下直接对比；其四，DiT-XL 上 gFID 绝对改善较小（2.91→2.38），且全层 rFID 0.42 仍逊于专用解码器的 0.18，鲁棒性与峰值保真之间存在残余权衡。
独立分析的弱点
弱点一：率选择昂贵。每个 $(p_{dec}, p_{dit})$ 组合都要从头训练，完整网格超过 60 次训练；改进方向是让 $p$ 随训练动态退火，或利用线性理论的闭式最优预测器（式 24/29）在线估计分歧协方差 $\Gamma$ 来预测最优率，避免全网格扫描。弱点二：率参数化受限。实现使用交换性率 $r_k = 1-p$，只是 $2^K$ 维率空间的一维切片，无法表达"前 7 层必保 + 深层随机"这类结构化分布；改进方向是按层深度设置分层保留率，或与 nested dropout 的前缀采样混合以利用层的序结构。弱点三：引导交互刻画不完整。对内部引导的分析固定在单个 $t$ 的线性代理上，而实际引导跨越整条采样轨迹；改进方向是把 REPA 头一并纳入随机化并做三路联合率选，或把引导差值方向的分歧惩罚直接写入损失。弱点四：验证域单一。只在类条件 ImageNet-256 验证，文本到图像、视频 RAE（VideoRAE/V-RAE）等场景未测；改进方向是在 T2I 潜空间的时间/空间层融合上复用同一正则。弱点五：峰值保真让步。全层 rFID 0.42 对比专用解码器 0.18，对需要极高保真编辑的应用，鲁棒解码器未必是最优选择，可以考虑推理时按融合类型切换或对 $p_{dec}$ 做轻量微调。
未来方向
作者提出的方向：把理论刻画扩展到混合重建损失与完整引导采样轨迹；在其他编码器层级、分辨率、数据域和更长训练调度上评估；澄清生成器容量与收敛如何影响率偏好（DiT-Base 偏好双正则、DiT-XL 的 gFID 偏好 $p_{dit}=0$）。基于成果可自然延伸的方向包括：（1）把 FuseReg 推广到视频 RAE 的时空层融合和文本到图像生成，检验"分歧方向正则化"是否跨模态成立；（2）设计自适应率调度，例如按 $c_K(p)$ 从 0 退火到目标值，或与 REPA/内部引导联合率选，解决有引导场景的率校准问题；（3）研究输入相关的层路由器——论文明确指出它落在独立选择规则模型之外，是现成的理论空白；（4）把层间分歧惩罚思想迁移到 LLM 的多层隐状态融合与层选择蒸馏；（5）探究 FuseReg 潜空间是否改善编辑、插值等下游应用，因为更平坦的层依赖可能意味着更线性、更可操控的潜空间几何；（6）联合校准引导强度 $w$ 与融合正则率，利用交叉项 $\mathrm{tr}(W_{full}\Gamma W_{repa}^\top)$ 的符号做诊断。
复现评估
复现条件较好。作者承诺公开全部源码、模型权重、训练脚本以及复现论文每一张表格的脚本，并声明已有独立研究者在各自计算环境中验证过所有表格数值，这在同类工作中相当罕见。依赖资源：DINOv3-L 权重（Meta 开源）、ImageNet-256、RAEv2 官方检查点（Table 1 的固定融合基线直接取自官方）。算力估算：单解码器 16 epochs、DiT-Base 40 epochs，但完整网格（解码器 8 个率 + DiT-Base $7\times7$ + DiT-XL $4\times8$ 配置）合计数十次训练，属于学术实验室可承受但并非轻量。建议的低成本复现路径：先复现 Table 4 的解码器率扫描（每次 16 epochs，最便宜，足以验证跨融合鲁棒性），再用公开的 RAEv2 DiT 做 Table 5 的解码器替换实验（无需训练生成器），最后视算力决定是否进入联合网格。理论部分（Lemma 1、Propositions 1/2）证明完整且自包含于附录 E，无需实验即可核对。
论文图表
在 DINOv3 的 7 个层上训练 softmax 可学习门控，绘制逐层门控权重随训练步数（0 到 15k）的演化。最浅层（L11）的权重持续增长并主导融合，深层（L15–L23）权重衰减到 0，且 GAN 损失开启后坍缩加速；虚线标注均匀权重 1/7 的参考线。
这张图否定了"学习一个更好的固定融合"这条看似自然的路线：像素重建损失会驱动门控坍缩到浅层捷径。它是引出"与其选一个融合，不如对所有融合鲁棒"这一核心动机的直接证据。