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

2026-09-25

AI Summary

This study addresses the inherent conflict between reconstruction and generation tasks in representation autoencoders, where divergent encoder layer selections lead to performance gaps. To resolve this, we propose FuseReg, which replaces heuristic feature selection with random subset sampling regularization, theoretically penalizing cross-layer inconsistency sensitivity. By jointly training with DINOv3-L and a DiT architecture, a single decoder accommodates diverse layer fusion strategies without retraining. On ImageNet-256, our method achieves superior PSNR compared to dedicated decoders. Replacing only the decoder reduces gFID by 27%, while joint regularization further decreases it to 29%, effectively unifying reconstruction and generation optimization within a single framework.

Abstract

Representation autoencoders (RAEs) reuse features from a pretrained visual encoder as reconstruction and diffusion latents, integrating strong visual representations into image generation. However, RAEs still need to decide which encoder layers form the shared latent space for the generator and pixel decoder. This choice involves a trade-off. Shallower layers tend to preserve fine pixel details better, while deeper layers tend to yield better generation metrics. A fixed heuristic layer fusion therefore couples two stages that benefit from different information. We introduce FuseReg, which replaces heuristic feature selection with training over random subsets of encoder layers. We theoretically analyze the underlying mechanism: subset sampling explicitly penalizes sensitivity to cross-layer disagreement. On ImageNet-256 with DINOv3-L, a single FuseReg decoder reconstructs from full, sparse, and single-layer fusions without retraining, achieving higher PSNR than decoders specialized to fixed fusions. This flexibility also benefits generation: decoder replacement alone reduces unguided gFID by 27% with an unchanged RAEv2 DiT-XL generator. The same regularization principle extends to diffusion training, with joint regularization of both stages reducing unguided gFID by 29% on DiT-Base. These results show that training downstream models for layer-fusion robustness narrows the reconstruction-generation gap without modifying the pretrained encoder.

Problem

• Representation Autoencoders
• Layer Fusion
• Reconstruction-Generation Gap
• Latent Space
• Image Generation

Innovation

• Representation Autoencoder
• Layer Fusion Regularization
• Random Subset Sampling
• Reconstruction-Generation Gap
• Diffusion Latents