FuseReg: Regularizing Layer Fusion Mitigates the Reconstruction-Generation Gap in Representation Autoencoders
发表时间:2026-09-25
arXiv · HF 讨论 · alphaXiv① 一句话定位
本文提出 FuseReg,一种面向表示自编码器(RAE)层融合接口的训练正则化方法。它不寻找“更好的固定融合”,而是让像素解码器与扩散生成器在训练中看到随机采样的编码器层子集,并以归一化子集均值作为隐变量。文章用线性平方误差代理目标证明该扰动保持全层均值、把跨层分歧转化为显式惩罚,并论证其在权重二阶矩上无法被任何确定性全局融合(含启发式与可学习全局门)等价替代。总体定位是训练策略与理论分析型工作,贡献集中在层融合鲁棒性,而非新架构。
② 背景与动机
表示自编码器使用冻结视觉编码器的特征作为重建和扩散的隐空间,从而把强视觉表征引入图像生成。关键问题是:编码器各层内容随深度变化,层融合定义了生成器与解码器之间的接口。浅层更保留像素细节,深层更利于生成建模;RAEv2 采用“最后 k 层求和”的固定启发式融合,导致同一隐变量必须同时服务两个偏好不同的消费者。已有工作或改变编码器、隐空间几何、训练配方,或学习一个确定性的融合规则(如 DRoRAE、注意力式多层探针),但仍收敛到单一组合。随机配置训练(dropout、nested dropout、stochastic depth、LayerDrop、slimmable networks、Matryoshka)通常服务于推理期提取更小子模型,而本文把同一机制迁移到“哪些编码器层进入解码器/生成器”这一轴,目的是覆盖所有非空层子集、缓解重建-生成差距。
③ 方法机制
FuseReg 的核心是把固定层融合替换为训练期的随机子集融合。设冻结编码器 E 对图像 x 产生 K 个候选层的归一化 token:h_k = LN(E_{l_k}(x)) ∈ R^{N×d}。固定融合写作 z_a = Σ_{k=1}^K a_k h_k,且权重满足 1^T a = 1。固定单层、固定前缀和可学习全局门都共享这一形式,因而都只暴露单一层组合给下游模型。 FuseReg 的采样分布为:对每个训练样本独立抽取原始掩码 m̃_k ~ Bernoulli(1-p),条件在其非零,即 m ~ M_p = Law(m̃ | 1^T m̃ > 0),隐变量为 z_m = (Σ_k m_k h_k)/(Σ_k m_k)。对 0 < p < 1,支撑集是所有 2^K - 1 个非空子集;推理时默认保留全部层,z_m 退化为全层均值 h̄ = K^{-1}Σ_k h_k。归一化是关键:按线性期望,E_m[z_m | x] = h̄,改变 p 只改变围绕部署隐变量的分歧方差,不系统性移动其均值。 两个阶段共享机制但使用不同丢弃率。解码器侧以 p_dec 采样 m,训练 ViT 解码器 D 重建 x̂ = D(z_m),损失为 L2 加感知项和对抗项:L_dec = ||x̂ - x||_2 + λ_p L_LPIPS(x̂,x) + λ_g L_GAN(x̂)。生成器侧以 p_dit 采样 m,保留共同预测目标为全层聚合 z_full = K^{-1}Σ_k h_k。流匹配插值为 z_t = t z_m + (1-t)ε,其中 m ~ M_{p_dit}、t ~ ν、ε ~ N(0,I),并以类别 c 为条件;DiT 回归全层目标,使用时间加权 x 预测损失:L_DiT = E ρ(t)||x_θ(z_t,t,c) - z_full||_2^2,其中 ρ(t)=max(1-t,0.05)^{-2},ν 是移位 logit-normal 时间律。p_dit=0 恢复标准全融合目标;p_dit>0 强制生成器把多个子集融合映射到同一全层目标。 理论部分先分析采样本身。定义层均值 h̄、层偏差 e_k = h_k - h̄、分歧协方差 V(x) = K^{-1}Σ e_k e_k^T、总体 Γ = E_x[V(x)]。设子集大小 R=|S|,系数 c_K(p) = E[(K-R)/(R(K-1))],期望在条件 R>0 的 Binomial(K,1-p) 下取。可得精确协方差恒等式:E_S[(z_S - h̄)(z_S - h̄)^T | x] = c_K(p) V(x)。因此子集采样在期望上保留 h̄,同时把方差精确注入层间分歧方向。 将该扰动放入线性平方误差消费者,可得到显式正则化。对解码器,展开 E_S||D z_S - y||^2 后交叉项因 E[δ_S|x]=0 消失,得到 ||D h̄ - y||^2 + c_K(p) tr(D V(x) D^T)。平均后惩罚项为 c_K(p) tr(D Γ D^T),等价于 c_K(p)/K Σ_k E||D(h_k - h̄)||^2,即惩罚解码器对层偏差的敏感性。闭式最优为 D_p* = A(Σ_h̄ + c_K(p)Γ)^{-1},其中 Σ_h̄ = E[h̄h̄^T]、A = E[y h̄^T];这会在层间分歧大的方向上抑制 D 的响应。对 DiT,z_{t,S} = t(h̄+δ_S)+(1-t)ε = z̄_t + tδ_S,展开后多出 t^2 c_K(p) tr(W V(x) W^T) 项,再乘 ρ(t);因此生成器侧分歧惩罚有时间依赖强度,且时间权重保持 RAEv2 速度空间损失的尺度。 文章还论证固定融合不等价:令随机融合权重 w_S = 1_S/|S|,P = I - K^{-1} 11^T 为层对比子空间投影。FuseReg 保持均值、并在每个层对比方向上产生方差;任何输入无关的确定性全局融合都无法匹配这些一、二阶权重矩,因此会诱导不同的线性目标。内部引导的交互也被显式化:x̂_guided = x̂_full + w(x̂_full - x̂_repa),其分歧惩罚 Ω(W_guided) = (1+w)^2 Ω(W_full) + w^2 Ω(W_repa) - 2w(1+w) tr(W_full Γ W_repa^T)。两个预测器对同一层分歧方向的响应若一致会在相减中抵消,若相反则相互放大,因此仅降低全预测器敏感性不足以控制引导组合的敏感性。 实验配置:冻结 DINOv3-L 编码器,ImageNet-256 分辨率,候选 K=23 个 transformer 层;ViT 解码器从头训练,架构、参数量、优化超参与训练步预算均匹配。解码器扫描为 16 个 epoch,使用像素、感知、对抗和分类器代理项;DiT 网格为 40 个 epoch,DiT-XL 隐藏宽度 1440。解码器丢弃率集合为 {0,0.05,0.1,0.3,0.5,0.7,0.9,0.95},DiT-Base 网格为 {0,...,0.9},DiT-XL 的 p_dit 集合为 {0,0.5,0.7,0.9}。主比较融合包括全 k=23、k=7(层 {ℓ11,ℓ13,...,ℓ23})以及单层 ℓ11;FuseReg 按保留层数归一化,固定融合基线保留原尺度约定。
④ 实验与证据
实验分三块分离解码器与生成器效应:其一,让一个解码器在全层、稀疏子集和单层融合上重建,检验其跨融合稳健性;其二,固定生成器与采样隐变量,仅替换解码器,考察重建训练对生成的影响;其三,在 DiT-Base 与 DiT-XL 上联合扫描两个丢弃率,并分别报告无引导与内部引导结果。总体发现:随机子集训练的解码器无需重新训练即可覆盖多种融合,固定融合解码器只在训练配置上表现好;解码器稳健性可转化为生成增益,即使生成器不变;两阶段正则化在 DiT-Base 上有互补收益,但在 DiT-XL 上最优速率随指标与引导设置变化;丢弃率响应非单调,说明机制放大的层间分歧惩罚而非信号损失。具体数值见下方 claim–quote 对。
-
扩大固定融合子集从 k=7 到 k=23 会提高重建 PSNR,但同时恶化无引导与引导 gFID,直接展示重建-生成差距。
expanding the fused subset from k = 7 k{=}7 to k = 23 k{=}23 increases reconstruction PSNR from 22.58 22.58 to 27.04 27.04 dB but also increases unguided gFID from 1.65 1.65 to 3.01 3.01 and guided gFID from 1.06 1.06 to 1.25 1.25
-
RAEv2 固定 K=23 解码器在其训练融合上达到较高 PSNR 和较低 rFID,但在非训练融合上 PSNR 大幅下降。
RAEv2 K=23 reaches 27.04 27.04 dB and 0.18 0.18 rFID on its training fusion but drops to 12.51 12.51 – 14.31 14.31 dB on the others; RAEv2 K=7 shows the same pattern.
-
单个 FuseReg 解码器在 k=7、k=23 与单层 ℓ11 三种融合上都保持较强重建,且 rFID 低于 0.6。
A single FuseReg decoder remains strong across all three ( 23.77 / 27.52 / 25.13 23.77/27.52/25.13 dB), with rFID below 0.6 0.6 .
-
在原生 k=23 融合上仅替换解码器即可降低无引导 gFID;在 k=7 融合上替换解码器同样大幅降低无引导 gFID。
On the native k = 23 k{=}23 fusion, swapping the plain RAEv2 decoder for a FuseReg decoder reduces gFID from 3.01 3.01 to 2.21 2.21 ; under the shifted k = 7 k{=}7 fusion, it reduces gFID from 27.73 27.73 to 1.92 1.92 .
-
在 k=7 融合下,无引导与引导生成都随解码器替换显著改善。
Under shifted k = 7 k{=}7 , both unguided and guided generation improve substantially, from 27.73 27.73 to 1.92 1.92 and from 16.35 16.35 to 1.42 1.42 , respectively.
-
原生 k=23 融合的引导 gFID 在中等丢弃率出现暂时回退,在强正则化下恢复到基线水平。
guided gFID temporarily regresses from 1.25 1.25 to as high as 1.49 1.49 at intermediate rates, then returns to 1.25 1.25 at p dec = 0.95 p_{\mathrm{dec}}{=}0.95 .
-
DiT-Base 上单独正则化解码器或生成器均改善 gFID,联合正则化达到更低的 gFID。
regularizing either the decoder ( p dec = 0.9 p_{\mathrm{dec}}{=}0.9 , p dit = 0 p_{\mathrm{dit}}{=}0 ) or the generator ( p dec = 0 p_{\mathrm{dec}}{=}0 , p dit = 0.7 p_{\mathrm{dit}}{=}0.7 ) improves gFID from the 13.96 13.96 baseline to 12.96 12.96 and 12.09 12.09 , respectively. Joint regularization reaches 9.93 9.93 at ( p dec = 0.9 p_{\mathrm{dec}}{=}0.9 , p dit = 0.7 p_{\mathrm{dit}}{=}0.7 ).
-
DiT-Base 无引导 gFID 在低解码器丢弃率时先退化,高丢弃率才超过基线,说明速率响应非单调。
at p dit = 0 p_{\mathrm{dit}}{=}0 , gFID initially degrades (e.g., 17.28 17.28 at p dec = 0.05 p_{\mathrm{dec}}{=}0.05 ) before outperforming the baseline at higher rates.
-
论文报告 DiT-XL 基线无引导 gFID,并声称生成器单独正则化无改善,而联合正则化解锁进一步收益。
When scaling to the stronger DiT-XL baseline (gFID 2.91 2.91 ), the generator’s response to regularization becomes nuanced ( Table 2(c) ). For gFID, regularizing the generator alone fails to yield improvements; however, joint regularization with the decoder successfully unlocks further gFID gains.
-
在 DiT-XL 无引导 gFID 表中,p_dit=0 行从 p_dec=0 的 2.91 降至 p_dec=0.95 的 2.38,显示强解码器正则化且生成器不丢弃时最低。
0 | 2.91 | 3.21 | 3.22 | 3.00 | 2.91 | 2.76 | 2.50 | 2.38
-
内部引导下 DiT-XL 最低引导 gFID 出现在强解码器正则化且生成器不丢弃的配置,优于全零配置。
The lowest guided gFID is 1.50 1.50 at ( p dec , p dit ) = ( 0.9 , 0 ) (p_{\mathrm{dec}},p_{\mathrm{dit}})=(0.9,0) , compared with 1.57 1.57 at ( 0 , 0 ) (0,0) .
-
解码器丢弃率扫描显示全融合 PSNR 在轻微扰动时达到峰值,而稀疏与单层输入随丢弃率持续改善,强正则化解码器覆盖三种融合。
Full-fusion PSNR peaks at p dec = 0.05 p_{\mathrm{dec}}{=}0.05 , while the k = 7 k{=}7 and ℓ 11 \ell_{11} inputs improve throughout the sweep. At p dec = 0.95 p_{\mathrm{dec}}{=}0.95 , one decoder reaches 23.77 23.77 , 27.52 27.52 , and 25.13 25.13 dB across the three fusions.
⑤ 审稿人视角
Novelty:创新点在于把层融合从“固定或可学习权重”重新定义为“训练分布”,并用归一化子集均值保持部署隐变量无偏。这与搜索更优固定融合的路线形成明确对比。随机丢弃、LayerDrop、Matryoshka 等机制已有,但本文将其用于层融合轴、服务鲁棒接口而非推理裁剪,并给出“分歧协方差惩罚”和“与确定性全局融合二阶分离”的理论解释。总体新颖性中等偏上:核心操作简单,概念转换和理论刻画是主要增量。 Soundness:受控实验较严谨:冻结编码器、固定生成器与采样隐变量以分离解码器效应,匹配架构、参数量与训练预算,并用线性平方损失代理目标给出可证明机制,同时明确其假设边界。但结论稳健性仍有保留:证据限于单一编码器、数据域和两个 DiT 尺度;部分叙述与表格不完全一致(如 DiT-XL 的“联合解锁”说法,而表中最优值对应生成器不丢弃);结果多为点估计,未报告重复实验或置信区间;内部引导的交互机制只得到部分线性解释。因此内部有效性与可复现性较好,外部有效性和统计可靠性尚未充分证明。 Significance:对 RAE 子领域有实际意义:在不修改冻结编码器的前提下,缓解重建与生成对层深度的偏好冲突,并让单个解码器跨融合可用。它把“层融合选择”从一次性设计变为可训练鲁棒性问题,可能推广到任何冻结层级特征的下游接口。不过当前证据集中在 DINOv3-L/ImageNet-256 与 DiT-Base/XL,距离更广泛的编码器、分辨率和任务仍有距离,因此意义目前更多是方法性和方向性,而非已被广泛验证的通用改进。
⑥ 局限与保留意见
作者自陈与本文判断一致的主要局限:实验仅覆盖 ImageNet-256、冻结 DINOv3-L 编码器、DiT-Base 与 DiT-XL,未跨编码器层级、分辨率、数据域和更长训练周期验证。两个阶段对正则化的响应不同,最优丢弃率依赖模型尺度、评价指标和引导设置,因此新场景需要分别选率,难以给出通用速率处方。理论中的采样恒等式与二阶分离是精确的,但目标分解依赖齐次线性预测器和平方损失,未覆盖混合重建损失与完整引导采样轨迹。此外,引导与融合正则化存在耦合,当前只在线性代理下解释,未系统解决;实验结果以点估计呈现,缺少不确定性量化。
⑦ 值得跟进
后续工作可沿几条线展开。第一,跨编码器与任务泛化:在 DINOv2、CLIP、SigLIP、MAE 等编码器以及更高分辨率、视频 RAE 上检验 FuseReg,判断分歧惩罚是否普遍有效。第二,速率自适应:研究逐层丢弃率、基于输入或训练进度的调度,或与 DRoRAE 类可学习融合结合,避免手工网格搜索。第三,深化引导交互:在完整采样轨迹上刻画全预测器与 REPA 预测器对层分歧的响应,显式校准引导强度与 p_dit。第四,理论与统计强化:把线性平方损失分解扩展到混合损失和非线性流匹配,并补充多次运行、置信区间与更长训练预算,澄清 DiT-XL 上生成器正则化的真实收益。第五,考察极端 K 与计算开销,以及子集鲁棒性是否对下游可控生成、编辑和条件控制带来额外益处。