FuseReg: Regularizing Layer Fusion Mitigates the Reconstruction-Generation Gap in Representation Autoencoders

What it claims

Randomly sampling which frozen-encoder layers are averaged into the latent during training makes both the RAE pixel decoder and the DiT robust to layer composition. This shrinks the conflict where shallow layers help reconstruction and deep layers help generation. The paper quantifies that conflict for the RAEv2 baseline (DINOv3-L, sum of the last k layers):

So the fusion that reconstructs best is the worst to generate from. FuseReg's headline absolute numbers:

• Decoder swap only, unchanged RAEv2 DiT-XL at k=23: unguided gFID drops from 3.01 to 2.21.
• Joint decoder plus DiT regularization, DiT-Base: gFID drops from 13.96 to 9.93.

The authors claim this adds no compute and requires no architectural change.

How it works

• Layer normalization. Each of K=23 candidate DINOv3-L layers is layer-normalized.
• Per-example mask. Each training example gets its own Bernoulli(1−p) keep mask over layers, resampled if empty.
• Normalized fusion. The latent is the mean of the kept layers, not the sum. The normalization is essential: it keeps the expected latent equal to the full-layer mean. Inference uses all layers by default.
• Decoder (rate p_dec). A ViT decoder, trained from scratch, reconstructs pixels from subset fusions. The loss is L2 + LPIPS + GAN.
• DiT (rate p_dit). The flow-matching input is built from a subset fusion, but the regression target is always the full-layer mean. It uses x-prediction with weight ρ(t)=max(1−t, 0.05)⁻² and a shifted logit-normal time law. Setting p_dit=0 recovers standard training.
• Separate rates. The two rates are deliberately decoupled, because the stages solve different prediction problems.

Motivating observation: a learnable softmax gate over layers L11–L23 collapses onto the shallowest layer (L11). The collapse accelerates once the GAN loss turns on. The authors therefore argue that learned global weights do not solve the problem.

The theory is exact for linear squared-loss surrogates:

• Subset sampling preserves the mean and adds variance only along cross-layer disagreement, scaled by a coefficient c_K(p).
• This is equivalent to an explicit penalty tr(DΓDᵀ) on the decoder, and a t²-scaled version of that penalty on the DiT.
• No fixed or learned input-independent fusion can match it. The random weights' second moment has rank K, while any fixed fusion's has rank ≤1.

Evidence

Setup:

• ImageNet-256 with a frozen DINOv3-L encoder.
• Decoder variants are matched in architecture, parameter count, hyperparameters and step budget. The same controls apply within each DiT scale (Base and XL).
• Reconstruction is scored with PSNR, SSIM and rFID on 50k images.
• Generation is scored with gFID and IS on 50k class-conditional samples using 50 Euler steps.

Reconstruction claim: one FuseReg decoder reconstructs from full (k=23), sparse (k=7) and single-layer (ℓ11) fusions with higher PSNR than RAEv2 decoders trained specifically on k=7 or k=23. Qualitatively, RAEv2's decoder feature maps degrade away from its training fusion, while FuseReg's stay stable.

How to use it

The page links a website, GitHub and Hugging Face but gives no URLs, commands, checkpoint names or licence in the fetched text. Adopting it means changing two things:

1. Swap fixed layer sums for normalized random-subset means.
2. Change the DiT target to the full-layer mean.

The paper says the decoder can be dropped into an existing RAEv2 DiT-XL pipeline unchanged.

Limits

• Truncated page. The fetched text cuts off at Table 1, so per-fusion PSNR, SSIM and rFID values and all ablations are not visible here.
• Missing settings. The chosen p_dec and p_dit values are not stated in the visible text.
• No guided result. No guided gFID is reported for FuseReg in the visible text.
• Weak DiT-Base baseline. The joint-regularization result is shown only at DiT-Base, where the baseline gFID of 13.96 is far from state of the art.
• Narrow coverage. Everything is tested on one encoder (DINOv3-L), one dataset (ImageNet-256) and class-conditional generation only.
• Theory scope. The guarantees hold for linear surrogates; the nonlinear behaviour is supported only empirically.