FuseReg trains a Representation Autoencoder (RAE) decoder and diffusion generator on randomly sampled subsets of frozen encoder layers, so both stages stop over-relying on one fixed layer combination. This closes the gap where shallow layers help reconstruction but deep layers help generation.
When you train a modern image diffusion model, you don’t diffuse on raw pixels. You diffuse in the latent space of an autoencoder. A recent line of work (Representation Autoencoder (RAE), and its follow-up RAEv2) replaces the usual VAE with a frozen pretrained vision encoder like DINOv3, plus a learned decoder that maps encoder features back to pixels. The generator then models the distribution of those encoder features.
One awkward choice remains: a vision transformer has ~23 layers, each producing features. Which layer(s) do you feed into the decoder and the diffusion model? RAEv2 sums the last k layers and picks k by hand. The paper shows this creates a real dilemma on ImageNet-256: with k=7 deep layers you get low unguided gFID (generation FID) of 1.65 but weak reconstruction PSNR of 22.58 dB; with k=23 (all layers) reconstruction jumps to 27.04 dB but gFID worsens to 3.01. Reconstruction wants shallow, pixel-aligned features; generation wants deeper, semantic ones. A single fixed fusion forces both stages to share the wrong compromise. The obvious fix, a learned gate that softmax-weights layers, collapses to the shallowest layer because the pixel loss rewards that shortcut (Figure 1).
The intuition: instead of picking one layer combination, train as if the combination is random. Each training step, sample a random nonempty subset of the K encoder layers, average their (layer-normalized) features, and feed that to the decoder. The decoder has to reconstruct pixels no matter which subset it sees, so it learns to pull reconstruction signal from every depth rather than leaning on one shallow shortcut.
Crucially, the average is size-normalized: divide by the number of kept layers, not by K. This preserves an important property. In expectation, the sampled latent equals the full all-layer mean, which is what the model will actually see at inference. So subset sampling doesn’t shift the deployment target; it only adds variance around it. The paper proves (for a linear squared-loss surrogate) that this variance lands exactly along cross-layer disagreement directions, the directions where different layers say different things. Training penalizes a predictor’s sensitivity to those directions. That is the regularization.
The same trick applies to the diffusion transformer (Diffusion Transformer (DiT)). It gets a noisy subset-mean as input but is asked to predict the full-layer mean as target. Two separate drop rates control the two stages: p_dec for the decoder, p_dit for the generator, since they solve different prediction problems.
# One training step, per example
K = 23 # candidate encoder layers
h = [layernorm(encoder[l](x)) for l in range(K)]
while True: # rejection-sample nonempty mask
m = [bernoulli(1 - p) for _ in range(K)]
if sum(m) > 0: break
z_subset = sum(m[k] * h[k] for k in range(K)) / sum(m)
# Decoder: reconstruct pixels from the subset mean
loss_dec = pixel_and_lpips_and_gan(decoder(z_subset), x)
# DiT: predict the FULL-layer mean from a noisy subset input
z_full = sum(h) / K
z_t = t * z_subset + (1 - t) * noise
loss_dit = weight(t) * mse(dit(z_t, t, class_c), z_full)
No new parameters, no architecture changes, no extra compute per step.
One decoder, many fusions. With DINOv3-L on ImageNet-256, the FuseReg decoder (trained with p_dec=0.95) hits 23.77 / 27.52 / 25.13 dB PSNR on the k=7, k=23, and single-layer-11 inputs respectively. The two fixed-fusion RAEv2 baselines each dominate their own training fusion but collapse elsewhere: RAEv2 trained on k=23 drops to 12.51 dB on the k=7 input. So the same decoder now supports full, sparse, and single-layer inputs without retraining.
Decoder swap alone improves generation. Take an unchanged RAEv2 DiT-XL generator (so the sampled latents are identical), and just swap in a FuseReg decoder to render them. Unguided gFID on the native k=23 fusion drops from 3.01 to 2.21 (~27% relative). On the shifted k=7 fusion, it drops from 27.73 to 1.92. Under RAEv2’s internal guidance at scale 1.78, guided gFID stays at 1.25 (matched) on k=23 but improves from 16.35 to 1.42 on k=7. Interpretation: better readout of a fixed representation improves the final images even when the generator and its latent samples are held constant.
Joint regularization compounds on DiT-Base. Regularizing only the decoder (p_dec=0.9) takes unguided gFID from 13.96 → 12.96; only the generator (p_dit=0.7) gives 13.96 → 12.09. Both together reach 9.93, more than the sum of the individual gains at the reported point. On the larger DiT-XL (baseline 2.91), the pattern shifts: strong decoder regularization still helps (down to 2.38 at p_dec=0.95, p_dit=0), but generator regularization alone doesn’t help gFID at that scale, though it does help Inception Score (IS). Preferred rates depend on model scale, metric, and whether guidance is on. Gains reproduce on SigLIP2-L and EUPE-B encoders.
One nuance the authors flag: gFID along the decoder-rate axis is non-monotonic. Small p_dec can briefly hurt before larger values recover and surpass baseline. The theory says a higher drop rate amplifies the disagreement penalty coefficient, not the amount of signal destroyed, because the mean is preserved.
•
If you’re already using an RAE-style pipeline (frozen vision encoder + learned decoder + latent diffusion), FuseReg is close to a free swap. Replace the fixed “sum the last k layers” step with size-normalized random-subset averaging during training. No architecture, parameter-count, or compute changes. The paper’s evidence is strongest for improving reconstruction robustness and for decoder-swap gains on a fixed generator.
•
Tune the two drop rates separately, per setting. The paper explicitly warns against assuming one rate fits both stages or both model scales. DiT-Base liked joint regularization; DiT-XL preferred strong decoder regularization with p_dit=0. Guidance changes preferences again. Budget a small grid.
•
Worth testing if your problem has a similar shape: multiple frozen feature sources that must collapse into one latent, where downstream consumers have conflicting preferences over which source. The mechanism (mean-preserving random subset averaging penalizes sensitivity to source disagreement) isn’t specific to vision layers. The paper only evaluates image RAEs though, so treat this as a hypothesis, not a claim.
•
What FuseReg does NOT do: it does not change the encoder, it does not improve any single-fusion PSNR beyond what a specialist decoder achieves on its own training fusion (RAEv2 at k=23 still wins on k=23 PSNR by a small margin in Table 1), and it does not always help every generation metric. On DiT-XL under internal guidance, the highest IS still sits at the unregularized baseline.
Code and checkpoints: GitHub, Hugging Face.
All experiments are ImageNet-256; higher resolutions, other domains, and longer training are not evaluated. The theoretical results (mean preservation, disagreement-covariance penalty, separation from any deterministic global fusion) are exact only for the specified subset distribution and for homogeneous linear predictors under squared loss. For the actual nonlinear ViT decoder and DiT, the theory is mechanism-level intuition, not a quantitative prediction, and the paper is explicit that it doesn’t predict the sign of the joint-rate effect or guarantee super-additivity. The reported best (p_dec, p_dit) combinations are point estimates from a grid, not statistical interaction tests. Finally, the internal-guidance analysis surfaces a real complication: guidance subtracts one predictor from another, and reducing one predictor’s sensitivity to layer disagreement can either help or hurt depending on how the two predictors’ disagreement responses align.