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×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×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( ↑ ) 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) ↑
GFLOPs ↓
Latency (ms) ↓
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 .
Lcls
Lmim
Lflip
FID ↓
sFID ↓
IS ↑
Prec. ↑
Rec. ↑
×
×
×
33.02
6.46
43.71
0.53
0.63
×
✓
×
31.92
6.26
48.14
0.54
0.64
✓
×
×
30.85
6.83
50.94
0.54
0.63
✓
✓
×
28.11
6.41
54.91
0.56
0.64
✓
×
✓
28.06
6.23
54.77
0.56
0.62
✓
✓
✓
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 ϵ=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×32×4
32×32×4
32×32×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 ↑
Far ↓
LDS ↑
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×16 token grid, and center anchor.
Configuration
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
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×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 ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
(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 (dh,depth,heads,rMLP) . Params, GFLOPs, and GMACs refer to the auxiliary composer.
Setting
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
Label smoothing
η=0.00
31.91
6.78
0.54
0.64
48.64
η=0.10
30.86
6.83
0.55
0.64
50.95
η=0.20
30.76
6.84
0.55
0.63
50.45
η=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 η and MIM masking ratio ρ while keeping the remaining training settings fixed.
Setting
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
(a) Alignment block
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 .
Setting
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
Composer adaptation
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.
CFG scale
Interval
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
Guidance-scale sensitivity
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.
Setting
FID ↓
sFID ↓
Prec. ↑
Rec. ↑
IS ↑
Inference-time null-centered scaling
γ=1.00
24.87
5.95
0.58
0.64
61.19
γ=1.10
18.48
5.72
0.62
0.61
81.05
γ=1.20
13.73
5.55
0.66
0.59
104.42
γ=1.30
10.41
5.47
0.70
0.57
129.03
γ=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.
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× faster than SiT and 25× faster than REPA. Code is available at https://github.com/saeidrazavi/RECON
We present Scalable Pixel-anchored End-to-end Diffusion (SpeeDiff), a latent diffusion method that jointly trains the VAE and the diffusion model from scratch. In principle, joint training allows the diffusion loss gradient to directly guide the VAE encoder, encouraging the formation of a generation-friendly latent space and potentially yielding faster convergence than the conventional two-stage approach with a pretrained frozen VAE. However, a naive end-to-end implementation severely degrades performance, as unrestricted backpropagation of the diffusion loss leads to latent space collapse. Our main technical contribution is a simple yet effective Tweedie Pixel Reconstruction (TPR) loss, which provides additional pixel-level feedback by decoding a predicted clean latent from an intermediate noisy state using Tweedie's formula, thereby alleviating collapse. Furthermore, our method enables jointly scaling a fully transformer-based architecture and enhances representation alignment within the end-to-end framework. Our SpeeDiff-XL model achieves over 140x and 61x faster training compared to Vanilla SiT and REPA, respectively, while attaining an FID of 1.50 without guidance on ImageNet 256x256 generation. With a more efficient 32x compressed VAE, our model further reaches an FID of 1.53 without guidance on ImageNet 512x512 generation.
Bingliang Zhang, Wenda Chu, Yizhuo Li +5
ByteDance Seed · California Institute of Technology · Meta
Pixel-space diffusion has recently re-emerged as a strong alternative to latent diffusion, enabling high-quality generation without pretrained autoencoders. However, standard pixel-space diffusion models receive relatively weak semantic supervision and are not explicitly designed to capture high-level visual structure. Recent representation-alignment methods (e.g., REPA) suggest that pretrained visual features can substantially improve diffusion training, and visual co-denoising has emerged as a promising direction for incorporating such features into the generative process. However, existing co-denoising approaches often entangle multiple design choices, making it unclear which are truly essential. We therefore present V-Co, a systematic study of visual co-denoising in a unified JiT-based framework. This controlled setting allows us to isolate the ingredients that make visual co-denoising effective. Our study reveals two main ingredients. First, co-denoising benefits from preserving feature-specific computation while enabling flexible cross-stream interaction, which leads to a fully dual-stream architecture together with a structurally defined unconditional prediction for classifier-free guidance. Second, it requires both stronger semantic supervision and proper cross-stream calibration, which we realize through a perceptual-drifting hybrid loss and RMS-based feature rescaling. Together, these findings yield a simple recipe for visual co-denoising. Experiments on ImageNet-256 show that, at comparable model sizes, V-Co outperforms the underlying pixel-space diffusion baseline and strong prior pixel-diffusion methods while using fewer training epochs, offering practical guidance for future representation-aligned generative models.