FuseReg: Mitigating the Reconstruction–Generation Gap in RAEs

Representation Autoencoders fuse features from multiple encoder layers into a shared latent space.

However, reconstruction and generation prefer different parts of the hierarchy:

• Decoder: prefers shallow, pixel-rich features.
• DiT: prefers deeper, more structured semantic features.

This creates a reconstruction–generation mismatch.

FuseReg

FuseReg replaces fixed layer fusion with random layer-subset sampling during training.

For a sampled subset (S),

$$z_S=\frac{1}{\mathrm{\mid }S\mathrm{\mid }}\sum _{k\in S}{h}_k\mathrm{.}z\_S\; =\; \backslash frac\{1\}\{|S|\}\backslash sum\_\{k\backslash in\; S\}\; h\_k.$$

The normalized sampling keeps the latent mean unchanged:

$$\mathbb{E}[z_S\mid x]=\bar{h}\mathrm{.}\backslash mathbb\{E\}[z\_S\; \backslash mid\; x]\; =\; \backslash bar\; h.$$

Under a linear squared-loss surrogate, we further show that random subset sampling explicitly regularizes cross-layer disagreement.

In short, FuseReg discourages models from relying too strongly on any particular layer composition.

Robust across layer fusions

A single FuseReg decoder remains effective across full, sparse, and even single-layer inputs, while fixed-fusion decoders degrade strongly outside their training fusion.

Generation improvements

FuseReg can be applied independently to the decoder and DiT.

• Decoder replacement only: unguided gFID 3.01 → 2.21
• DiT-Base, joint FuseReg: unguided gFID 13.96 → 9.93
• DiT-XL: unguided gFID 2.91 → 2.38

The same trend also holds across different pretrained encoder families.

One decoder, many layer combinations

The same FuseReg decoder can reconstruct from different layer subsets without retraining, including sparse and single-layer inputs.

What changes?

FuseReg distributes reconstruction information more evenly across encoder depth:

• lower dependence on individual layers;
• stronger reconstruction from single layers;
• more stable representations under changing layer compositions.

Instead of searching for one optimal fusion, FuseReg trains the model to remain robust across a distribution of layer fusions.