FuseReg: Regularizing Layer Fusion Mitigates the Reconstruction-Generation Gap in Representation Autoencoders
Hongyang Du, Yunfei Xie, Junjie Ye, Jiawei Yang, Xiaoyan Cong, Haodong Zhang, Yongchao Huang, Haiyu Wu, Zongxia Li, Shihang Gui, Dawei Liu, Runhao Li, Jingcheng Ni, Chen Wei, Randall Balestriero, Yue Wang
cs.CV
2026-09-26
Training RAE decoder and DiT on random encoder-layer subsets drops ImageNet-256 unguided gFID from 3.01 to 2.21 by swapping the decoder alone.
Representation autoencoders (RAEs) reuse a frozen visual encoder as the latent for both a diffusion generator and a pixel decoder. The remaining design choice is which encoder layers become that shared latent.
Shallower layers keep pixel detail. Deeper layers are easier to model generatively. RAEv2 commits to a fixed heuristic: sum the last k layers. On ImageNet-256 with DINOv3-L, expanding k from 7 to 23 lifts reconstruction PSNR from 22.58 dB to 27.04 dB, but unguided gFID worsens from 1.65 to 3.01 and guided gFID from 1.06 to 1.25. A learnable global gate collapses onto the shallowest layer, faster once the GAN turns on. Reconstruction wants shallow features. Diffusion wants deep ones. A single fixed fusion ties both stages to the same compromise.
FuseReg does not search for a better fixed fusion. During training it samples a nonempty random subset of encoder layers and uses their normalized mean as the latent. The expected latent stays the full-layer mean, which is what inference uses. The extra variance sits on directions where layers disagree. Under linear squared loss that variance is an explicit disagreement penalty. No deterministic global fusion, including a learned input-independent gate, can match those second-order weight moments.
The same sampler is applied at two stages with separate drop rates:
Pixel readout and latent prediction are different problems, so pdec and pdit are chosen independently. The encoder stays frozen. Architecture and compute barely change.
One FuseReg decoder reconstructs from full, sparse, and single-layer fusions. Fixed-fusion decoders collapse off their training subset.
| Decoder | k=7 PSNR / rFID | k=23 PSNR / rFID | layer ℓ11 PSNR / rFID |
| RAEv2 trained at k=7 | 22.58 / 0.30 | 18.37 / 1.62 | 17.50 / 6.35 |
| RAEv2 trained at k=23 | 12.51 / 16.10 | 27.04 / 0.18 | 14.31 / 11.21 |
| FuseReg p=0.95 | 23.77 / 0.60 | 27.52 / 0.42 | 25.13 / 0.45 |
Generation is the stricter test. Holding an RAEv2 DiT and its sampled latents fixed, swapping only the decoder cuts unguided gFID from 3.01 to 2.21 at k=23 (about 27%). Guided gFID stays at 1.25. Under the shifted k=7 fusion, unguided gFID falls from 27.73 to 1.92 and guided from 16.35 to 1.42.
Joint regularization on DiT-Base takes unguided gFID from 13.96 to 9.93 (about 29%), more than adding the single-axis gains. On DiT-XL the lowest unguided gFID is 2.38 with strong decoder regularization and pdit=0, from a 2.91 baseline. Generator-side dropout helps Inception Score more than gFID at that scale. Rate curves are not monotone: small drop rates can hurt before high rates beat the baseline.
This is an interface regularizer, not a new tokenizer. If an RAE/RAEv2 pipeline is already running, a fusion-robust decoder can improve generation with the generator weights untouched. Random subset sampling barely changes the architecture. The claim is narrow: read out information the frozen encoder already has, instead of wasting it on a brittle layer recipe.
All numbers are ImageNet-256, frozen DINOv3-L, DiT-Base/XL, matched step budgets. Other encoder stacks, resolutions, domains, and longer schedules are untested. Preferred drop rates move with model scale, metric, and internal guidance, so a new setting needs a two-axis sweep. The linear squared-loss decomposition is an analysis tool; nonlinear gains are empirical. How layer regularization interacts with RAEv2 internal guidance, which amplifies the gap between the full head and the REPA head, is still open.