Skip to main content
Back to timeline
arXivSource publication:

Diffusion training drops the external vision encoder: reusing VAE features cuts SiT-XL/2 FID from 17.2 to 7.4

Synopsis

CoVA recomposes the generative model's own VAE reconstruction features online into supervision with stronger semantics and spatial structure, cutting SiT-XL/2 FID on ImageNet 256x256 from 17.2 to 7.4 without an external vision encoder.

Source-provided article image: CoVA: Composing VAE Representations for Generative Learning
Figure 1 · arXiv

Interpretation

Diffusion transformer training can stop relying on external vision encoders such as DINO or CLIP and instead reuse features from the generative model's own VAE as intermediate supervision. Methods such as REPA align intermediate generative features to representations from a pretrained visual encoder, which requires an additional visual backbone throughout training and adds compute and memory overhead. On ImageNet 256x256 under the same 400K-step budget, CoVA reduces FID from 33.0 to 21.4 for SiT-B/2, from 18.8 to 9.3 for SiT-L/2, and from 17.2 to 7.4 for SiT-XL/2.

A lightweight online composer reorganizes reconstruction-oriented VAE latents into dense features that carry both class semantics and spatial structure, then uses them in stages. Raw VAE features are learned primarily for reconstruction, so their semantic discrimination and spatial structure are not necessarily suited to supervising a generative transformer. The composer is trained jointly with class supervision, masked latent reconstruction, and horizontal-flip consistency, adding only 8.5M training-only auxiliary parameters, 1.3% of SiT-XL/2, and is discarded at inference.

Representation alignment is treated as a stage-dependent signal: composed features align an intermediate SiT layer early, then alignment is removed and the composer's class posterior is reused for sample-dependent soft label conditioning. Maintaining alignment late in training degrades generation, and naively averaging class embeddings for a soft condition changes both the semantic direction and the magnitude of the condition. On ImageNet 256x256, CoVA reaches FID 2.03 and IS 316.4 in only 400 epochs with classifier-free guidance, versus FID 2.06 and IS 270.3 for vanilla SiT-XL/2 after 1,400 epochs.

Perspective

The report covers class-conditional image generation on ImageNet 256x256 with SiT-B/2, SiT-L/2, and SiT-XL/2, where images are encoded into latents by a pretrained Stable Diffusion VAE. The composer and alignment projector are used only during training and discarded at inference, and sampling follows the standard SiT pipeline. It is aimed at researchers and engineers working on diffusion transformer training efficiency and representation alignment.

Varying composer capacity from smallest to largest changes FID only marginally, indicating the gains do not rely on auxiliary model capacity. Alignment becomes detrimental late in training, and soft label conditioning without magnitude correction changes conditioning strength and creates a fidelity-coverage trade-off. These conclusions come from the ImageNet 256x256 experimental setting.