CoVA: Composing VAE Representations for Generative Learning

Zhaowei Chen; Yikai Ge; Linze Li; Yao Tang; Xiang Li

cs.CVOct 9, 2026

CoVA: Composing VAE Representations for Generative Learning

Authors: Zhaowei Chen, Yikai Ge, Linze Li, Yao Tang, Xiang Li

Organizations: Jiiov Technology · Nankai University

Abstract

Representation alignment with pretrained visual encoders improves diffusion transformer training but requires auxiliary encoders and extra training overhead. Recent approaches instead exploit representations from the generative model's own variational autoencoder (VAE). However, reconstruction-oriented VAE features are not optimized for the semantic discrimination and spatial structure needed for generative learning. We introduce CoVA (Composed VAE Representations), a lightweight framework that composes VAE features online into stage-wise guidance for flow-matching training. A small latent composer is jointly trained with class supervision, masked latent reconstruction, and horizontal-flip consistency, yielding semantically discriminative, spatially structured dense features and sample-dependent class posteriors. CoVA aligns intermediate SiT features with the composed representations early in training, then removes the increasingly detrimental alignment and reuses detached composer posteriors for sample-dependent soft label conditioning. We further uncover a hidden confound in soft label conditioning: averaging class embeddings attenuates their magnitude, coupling semantic interpolation with conditioning strength. This attenuation alone can substantially improve FID while degrading Recall, revealing a non-semantic source of the apparent gain. CoVA corrects this confound by restoring the hard-label magnitude, isolating posterior-dependent directional semantics. CoVA adds only 8.5M training-only auxiliary parameters (1.3% of SiT-XL/2) and introduces no inference-time overhead. On ImageNet 256 $\times$ 256, CoVA reduces SiT-XL/2 FID from 17.2 to 7.4 (57%) in 400K steps. With classifier-free guidance, it achieves an FID of 2.03 in only 400 epochs, versus 2.06 after 1,400 epochs for vanilla SiT-XL/2.

Figures & tables

Figure 1: High-fidelity image generation with CoVA. Samples from SiT-XL/2 with CoVA on ImageNet at $256\times 256$ resolution. CoVA produces images with fine-grained details, coherent structures, and strong semantic consistency.

Figure 2: CoVA reorganizes raw VAE features into more spatially coherent representations. Each row shows one ImageNet image together with PCA visualizations and cosine-similarity maps of raw VAE latent features, CoVA-composed features, and DINOv2-B/14 features. For PCA visualization, patch features are L2-normalized and projected onto the first three principal components, which are mapped to RGB channels. Each cosine map shows the similarity between all patch tokens and the query token marked by the yellow box. Compared with raw VAE features, CoVA produces substantially more coherent and object-aligned spatial structure. Quantitatively, Spatial LDS( $\uparrow$ ) improves from $0.28$ for raw VAE features to $0.41$ with CoVA, comparable to the $0.40$ achieved by DINOv2.

Figure 3: Comparison between SRA2 and CoVA. (a) SRA2 aligns intermediate SiT features with VAE features through a lightweight projection layer and an auxiliary feature-alignment loss. (b) CoVA first composes clean VAE latents with an online composer trained by semantic and spatial objectives, and then reuses its dense features for early representation alignment and its detached class posterior for later soft conditioning. (c) CoVA’s stage-adaptive transition from representation alignment to soft conditioning enables faster training convergence and improved generation quality.

Table 4

Method	#EFP (M)	TS (Iter/s) $\uparrow$	GFLOPs $\downarrow$	Latency (ms) $\downarrow$
SiT-XL/2	0 + 0	8.1	114.46	18.05
+REPA	86 + 8	6.3 ( -22% )	138.50 ( +21% )	22.81 ( +26% )
+SRA2	0 + 18	7.2 ( -11% )	118.55 ( +4% )	19.20 ( +6% )
+CoVA	0 + 8.5	7.0 ( -14% )	122.69 ( +7.2% )	21.17 ( +17% )

Table 3: Training-time computational overhead on SiT-XL/2. EFP denotes additional feature-processing parameters, reported as external model parameters + trainable auxiliary parameters .

Figure 4: Stage-wise scheduling of representation alignment and conditioning. (a) Early stopping consistently outperforms full-training alignment, with stable performance across a broad stopping range. (b) Sample-adaptive SLC consistently outperforms fixed label smoothing across all replacement probabilities, with the advantage becoming more pronounced as the replacement probability increases. (c) Longer naive SLC warmup consistently improves performance across the tested schedules. For both (a) and (b), the MIM mask ratio is fixed to

0.5

$\mathcal{L}_{\mathrm{cls}}$	$\mathcal{L}_{\mathrm{mim}}$	$\mathcal{L}_{\mathrm{flip}}$	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
$\times$	$\times$	$\times$	33.02	6.46	43.71	0.53	0.63
$\times$	$\checkmark$	$\times$	31.92	6.26	48.14	0.54	0.64
$\checkmark$	$\times$	$\times$	30.85	6.83	50.94	0.54	0.63
$\checkmark$	$\checkmark$	$\times$	28.11	6.41	54.91	0.56	0.64
$\checkmark$	$\times$	$\checkmark$	28.06	6.23	54.77	0.56	0.62
$\checkmark$	$\checkmark$	$\checkmark$	26.83	6.32	57.69	0.57	0.64

Table 4: Ablation on composer training objectives. All variants use the same feature-alignment configuration without SLC or early stopping. When enabled, classification uses label smoothing

\epsilon=0.1

, and MIM uses a masking ratio of

0.7

. Prec. and Rec. denote precision and recall.

Table 8

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Backbone	SiT-B/2	SiT-L/2	SiT-XL/2
Architecture
# Params	130M	458M	675M
Input latent	$32\times 32\times 4$	$32\times 32\times 4$	$32\times 32\times 4$
# SiT blocks	12	24	28
Hidden dimension	768	1,024	1,152
# Attention heads	12	16	16

Appendix

Table 7: CoVA hyperparameter settings across SiT model scales. The alignment layer is indexed by the one-based SiT Transformer block depth. The same composer and stage-wise schedule are used across scales unless otherwise specified.

Feature / objective	Near $\uparrow$	Far $\downarrow$	LDS $\uparrow$
Raw VAE	0.44	0.16	0.28
CE	0.48	0.19	0.29
MIM only	0.92	0.86	0.06
CE+MIM	0.60	0.24	0.36
CE+flip	0.92	0.83	0.09
CE+MIM+flip	0.66	0.25	0.41

Appendix

Table 8: Spatial LDS objective comparisons at 200K. Near/far denote mean cosine similarity; LDS is near minus far. All values are rounded to two decimal places.

Figure 5: PCA and cosine-map comparison for the seven table conditions. Each panel uses the same panda image, $16\times 16$ token grid, and center anchor.

Configuration	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
SiT-B/2	33.02	6.46	0.53	0.63	43.71
+ Composed representation alignment	26.83	6.32	0.57	0.64	57.69
+ Alignment early stopping	24.87	5.96	0.58	0.64	61.19
+ Soft label conditioning	21.45	5.69	0.59	0.63	74.57

Appendix

Table 9: Cumulative ablation of CoVA on SiT-B/2. All models are trained for 400K iterations on ImageNet

256\times 256

without CFG under the same optimization and sampling protocol. Each row progressively adds one component to the preceding configuration.

Composer	Params	GFLOPs	GMACs	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
$(192,12,3,4)$	5.59M	3.34	1.67	27.66	6.32	0.57	0.64	55.28
$(192,8,3,4)$	3.81M	2.23	1.11	27.52	6.22	0.57	0.64	55.96
$(768,3,12,4)$	22.26M	11.54	5.77	28.19	6.25	0.56	0.64	54.84
$(768,4,12,4)$	29.35M	15.38	7.69	27.32	6.27	0.57	0.64	55.74

Appendix

Table 10: Composer capacity. Configurations are reported as

(d_{h},\mathrm{depth},\mathrm{heads},r_{\mathrm{MLP}})

. Params, GFLOPs, and GMACs refer to the auxiliary composer.

Setting	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
Label smoothing
$\eta=0.00$	31.91	6.78	0.54	0.64	48.64
$\eta=\mathbf{0.10}$	30.86	6.83	0.55	0.64	50.95
$\eta=0.20$	30.76	6.84	0.55	0.63	50.45
$\eta=0.30$	31.36	6.79	0.54	0.64	49.90
MIM masking ratio

Appendix

Table 11: Sensitivity to composer objective hyperparameters. We vary label smoothing

\eta

and MIM masking ratio

\rho

while keeping the remaining training settings fixed.

(a) Alignment block
Setting	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
Block $4$ (Default)	27.66	6.32	0.57	0.64	55.28
Block $6$	30.02	6.43	0.55	0.64	51.09
Block $8$	32.65	6.42	0.54	0.64	47.13
(b) Loss weights
Default $(1,1,1)$	27.66	6.32	0.57	0.64	55.28

Appendix

Table 12: Sensitivity to alignment placement and loss weights. Each sweep varies one factor while keeping the remaining settings fixed. The default alignment block is

4

, and all default loss weights are

1.0

Composer adaptation
Setting	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
Joint online composer	27.66	6.32	0.57	0.64	55.28
Pre-train then freeze	29.27	6.40	0.56	0.63	52.30
Alternative SLC schedules
Early activation ( $150$ K)	35.57	6.49	0.52	0.64	43.31
Middle activation	35.96	6.48	0.52	0.64	43.02

Appendix

Table 13: Online adaptation and alternative SLC schedules. These schedules are exploratory variants of the final late-stage SLC design.

Guidance-scale sensitivity
CFG scale	Interval	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
1.60	$[0,0.50]$	2.01	4.63	0.79	0.62	291.74
1.65	$[0,0.50]$	1.99	4.54	0.79	0.61	299.15
1.70	$[0,0.50]$	2.01	4.56	0.79	0.62	304.19
1.80	$[0,0.50]$	2.03	4.52	0.80	0.62	316.42
Guidance-interval sensitivity

Appendix

Table 14: CFG sensitivity of SLC-trained SiT-XL/2. Results are evaluated after

2

M training iterations using

50

K samples. The interval denotes the portion of the sampling trajectory where CFG is applied.

Inference-time null-centered scaling
Setting	FID $\downarrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	IS $\uparrow$
$\gamma=1.00$	24.87	5.95	0.58	0.64	61.19
$\gamma=1.10$	18.48	5.72	0.62	0.61	81.05
$\gamma=1.20$	13.73	5.55	0.66	0.59	104.42
$\gamma=1.30$	10.41	5.47	0.70	0.57	129.03
$\gamma=1.35$	9.19	5.45	0.71	0.56	141.10

Appendix

Table 15: Null-centered controls for soft label conditioning. Top: inference-time scaling of the hard-label displacement from the learned unconditional embedding using Eq. 22 ; all model parameters are fixed and the model is trained without SLC. Bottom: training-time SLC using the null-centered magnitude correction in Eq. 20 , applied during the final 50K iterations.

Figure 6: The visualization results of SiT-XL/2 + CoVA use CFG with w = 4.0.

Figure 7: The visualization results of SiT-XL/2 + CoVA use CFG with w = 4.0.

Figure 8: The visualization results of SiT-XL/2 + CoVA use CFG with w = 4.0.

Figure 9: Additional visualizations of spatial representations.

Diffusion-Based Image Generation

Diffusion-based image generation creates images by progressively denoising noise, with controls for content, style, and structure.

See all 328 papers

Representation Conditioning for Generation
Oct 14, 2026Saeed Razavi, Xingjian Leng, Liang ZhengDiffusion-Based Image Generation Diffusion Transformer
Australian National University
SpeeDiff: Scalable Pixel-Anchored End-to-End Latent Diffusion Model
Oct 26, 2026Bingliang Zhang, Wenda Chu, Yizhuo Li +5Latent Diffusion Models Diffusion Transformer
ByteDance Seed · California Institute of Technology · Meta
V-Co: A Closer Look at Visual Representation Alignment via Co-Denoising
Mar 17, 2026Han Lin, Xichen Pan, Zun Wang +4Diffusion Model Alignment Pixel-Space Diffusion Models
1UNC Chapel Hill · 2NYU · 3Meta +1

Oct 14, 2026J/K move · Enter open · S save

Representation Conditioning for Generation

Saeed Razavi, Xingjian Leng, Liang Zheng

Australian National University

Diffusion-Based Image Generation Diffusion Transformer Diffusion Model Acceleration Conditional Diffusion Models

Open paper AskPDF arXiv

Representation-alignment methods effectively accelerate the diffusion transformer training. REPA, for example, uses a pretrained vision encoder to supervise an intermediate layer through an MLP projection head which is discarded after training. It aims to import semantics from the vision encoder to the diffusion backbone. While this objective is achieved at early layers by back-propagating the alignment loss, there is no guarantee that later layers, supervised by the velocity prediction loss alone, can well access such semantics. We propose a simple solution, representation conditioning (RECON), to improve semantics in later layers. During training, we feed the pretrained vision encoder tokens into later layers; during inference, because we no longer have the clean image, we reuse the MLP head (as a proxy to the vision encoder) and project its output tokens to the later layers. Our method adds no parameters and leads to much faster convergence: RECON is

63\times

faster than SiT and

25\times

faster than REPA. Code is available at https://github.com/saeidrazavi/RECON

Abstract

Diffusion-Based Image Generation

Explore similar work

Representation Conditioning for Generation

SpeeDiff: Scalable Pixel-Anchored End-to-End Latent Diffusion Model

V-Co: A Closer Look at Visual Representation Alignment via Co-Denoising