RepFlow: Reciprocal Supervision Improves Generation and Representation in Flow Models

Weili Zeng; Feng Tian; Shengqi Liu; Yichao Yan

cs.CVSep 27, 2026

RepFlow: Reciprocal Supervision Improves Generation and Representation in Flow Models

Authors: Weili Zeng, Feng Tian, Shengqi Liu, Yichao Yan

Organizations: Shanghai Jiao Tong University

Abstract

Generative models learn visual structure through denoising, yet their internal states are entangled with both noise level and network depth, making it difficult to obtain a stable visual representation from the generator itself. We introduce RepFlow, which learns such a representation from the generator's evolving computation and uses it to guide generation. Specifically, a separate timestep-free encoder is trained, through a timestep-conditioned predictor, to recover generator states across depths and noise levels from masked clean images. By excluding the reference noise and masked-out content from the encoder's input, we encourage the encoder to distill visual information that is recoverable from the visible image context and predictive of generator states. The representation learned from the generator's evolving states is then fed back to guide and improve the generator, whose updated states provide supervision for further representation learning, forming a reciprocal learning process. This reciprocal process improves multi-step generation and representation quality, as measured by frozen linear probing on ImageNet with latent-space SiT and pixel-space JiT, without an externally pretrained representation teacher. Across the two unconditional settings, FID decreases by $19.2$ -- $40.8\%$ relative to native training, while linear-probe accuracy improves by 6.80--10.04 percentage points over the best searched raw generator features. The learned representation also serves as a distributional metric for one-step JiT post-training, extending its role from instance-level alignment to distribution-level supervision.

Figures & tables

Figure 1: Connections between generation and representation learning. Two common strategies extract features from a pretrained generator or use a pretrained representation to guide generation. RepFlow learns a clean-image representation from the generator’s evolving states and uses it to shape subsequent generative learning.

Figure 2: Generator states across network depth and noise level. The same image induces different feature maps at different stages of generative computation. RepFlow uses these states as prediction targets for a clean-image encoder.

Figure 3: RepFlow. Generator states supervise masked prediction through G2R, while full-image features from an EMA encoder guide an early generator state through R2G. Stopped targets separate the gradient paths, and the models interact through the targets used in subsequent updates. The native generator objective is $\mathcal{L}_{\mathrm{FM}}$ for SiT and $\mathcal{L}_{\mathrm{JiT}}$ for JiT. Flames indicate trainable modules; double slashes indicate stopped gradients.

		SiT		JiT
Component	Variant	FID $\downarrow$	LP $\uparrow$	FID $\downarrow$	LP $\uparrow$
Reciprocal path	Native / raw state	32.83	58.02	35.58	61.79
	G2R only	32.83	65.35	35.58	66.19
	G2R+R2G	19.43	68.06	28.76	68.59
G2R target	Pixel reconstruction	25.23	57.23	31.56	59.88
	Generator states	19.43	68.06	28.76	68.59

Table 1: Unconditional ImageNet ablations at 600 epochs. LP denotes linear-probe accuracy (%).

Method	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Pre. $\uparrow$	Rec. $\uparrow$
SiT	9.62	5.31	123.0	0.67	0.68
SiT + SRA	6.77	4.91	142.0	0.69	0.68
SiT + RepFlow	6.99	5.17	147.9	0.68	0.70
SiT + REPA	6.02	5.05	159.3	0.70	0.69
JiT	16.80	6.93	93.7	0.61	0.67
JiT + SRA	14.24	7.04	101.8	0.63	0.67

Table 2: Class-conditional ImageNet at 600 epochs. Bold/underline denote the best/second-best results per backbone; ties share formatting.

Panel of Figure 6 — Figure 6: Class-conditional ImageNet FID-50K across training epochs for native training, RepFlow, SRA, and REPA. RepFlow learns its representation target concurrently with the generator.

One-step generative models
Method	Space	Epoch	NFE	Ext. Enc.	FID $\downarrow$	$\mathrm{FDr}^{6}$ $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
iCT-XL/2 ( Song and Dhariwal, 2024 )	latent	–	1	No	34.24	–	–	–	–
Shortcut-XL/2 ( Frans et al., 2025 )	latent	250	1	No	10.60	–	–	–	–
MeanFlow ( Geng et al., 2025 )	latent	240	1	No	3.43	–	–	–	–
iMF-XL ( Geng et al., 2026 )	latent	800	1	No	1.82	8.39	278.9	0.78	0.63
Drift-L (latent) ( Deng et al., 2026 )	latent	1280	1	latent-MAE	1.53	10.92	257.2	0.79	0.63

Table 3: Class-conditional one-step ImageNet generation. The final two rows share the generator checkpoint and post-training protocol. Ext. Enc. lists externally pretrained representation encoders.

\mathrm{FDr}^{6}

( Yang et al., 2026 ) averages normalized Fréchet distances in Inception ( Szegedy et al., 2016 ) , ConvNeXt ( Liu et al., 2022 ) , DINOv2 ( Oquab et al., 2023 ) , MAE ( He et al., 2022 ) , SigLIP ( Zhai et al., 2023 ) , and CLIP ( Radford et al., 2021 ) spaces.

\dagger

denotes FD-SIM training in SigLIP, Inception, and MAE spaces. RepFlow-student FD uses its own learned encoder.

Figure 7: Selected one-step samples from RepFlow-JiT-XL with RepFlow-student FD on ImageNet at $256\times 256$ resolution.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

	SiT-XL/2	JiT-XL/16
AdamW $\beta$	$(0.9,0.999)$	$(0.9,0.95)$
Learning rate	$1{\times}10^{-4}$	$2{\times}10^{-4}$
Schedule	constant	cosine, $5$ warmup epochs
Minimum learning rate	–	$1{\times}10^{-6}$
Weight decay	$0$	$0$
EMA decay	$0.9999$	$0.9999$

Appendix

Table 4: Optimization and sampling settings. The EMA decay applies to the generator and, in RepFlow, to the full-view student target.

Panel of Figure 8 — Figure 8: Random, uncurated one-step ImageNet samples, classes 15–341.

Panel of Figure 9 — Figure 9: Random, uncurated one-step ImageNet samples, classes 353–551.

Panel of Figure 10 — Figure 10: Random, uncurated one-step ImageNet samples, classes 553–666.

Panel of Figure 11 — Figure 11: Random, uncurated one-step ImageNet samples, classes 676–804.

Panel of Figure 12 — Figure 12: Random, uncurated one-step ImageNet samples, classes 834–961.

Image Generation

Image generation studies models that create or edit images from text, examples, layouts, and other structured controls.

See all 463 papers

From Pixels, Without Pre-training: Joint Generative and Self-Supervised Representation Learning in One Model
Oct 5, 2026Vicente Balmaseda, Ching-Long Lin, Tianbao YangRepresentation Learning Image Generation
Department of Computer Science and Engineering Texas A&M University College Station, TX, USA · Mechanical Engineering University of Iowa Iowa City, IA, USA
LoopREG: Representation Supervision Enables Test-Time Recurrent Depth in Diffusion Transformers
Oct 13, 2026Zipeng Wang, Anping Xiong, Jing Zhang +17Diffusion-Based Image Generation Diffusion Transformer
Chongqing University of Posts and Telecommunications · China United Network Communications Group Co., Ltd. · University of Electronic Science and Technology of China +1
Training-Free Hidden-State Refinement for Flow-Matching Image Generators
Aug 29, 2026Yuanyi Yan, Xinzhe Rao, Canyu Shen +5Flow-Based Generative Modeling Flow Matching
Huazhong University of Science and Technology · Tongji University · King Abdullah University of Science and Technology +2

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

From Pixels, Without Pre-training: Joint Generative and Self-Supervised Representation Learning in One Model

Vicente Balmaseda, Ching-Long Lin, Tianbao Yang

Department of Computer Science and Engineering Texas A&M University College Station, TX, USA · Mechanical Engineering University of Iowa Iowa City, IA, USA

Representation Learning Image Generation Self-Supervised Visual Representation Learning

Open paper AskPDF arXiv

Strong image generation models are conditioned on class labels, aligned to frozen pretrained encoders, or built on separately trained autoencoders. While effective, generation then depends on supervision or pretraining: labels must be annotated, and encoders or autoencoders pretrained for the target domain. We study joint generative and self-supervised representation learning in a single model, enabling self-conditioned generation without labels or pretrained models. This is challenging because the objectives are mismatched: contrastive learning consumes clean augmented views and favors coarse, invariant semantics, while flow matching consumes noisy images and must preserve the fine detail and spatial layout that contrastive learning discards. We propose SCION (Self-conditioned Generation on Self-supervised representation), whose core is a single pixel-space encoder conditioned on the flow timestep and an embedding. For representation learning, this conditioning embedding is a learned global vector shared across images, with the encoder's [CLS] token yielding the semantic representation trained by the contrastive loss. For generative training, the conditioning embedding is the image's own [CLS] representation, while patch tokens pass through a decoder to predict the image. To sample without a reference image at inference, we jointly learn a prior over the embedding. Gradient-norm balancing and stop-gradient mechanisms enable joint optimization in one run. SCION is self-supervised and self-contained, with no labels or pretrained models. On ImageNet 256x256, with the JiT-B recipe and no representation guidance, SCION reaches 8.92 FID, surpassing class-unconditional iREPA, which aligns to pretrained DINOv2 (46.44), and RCG, which conditions on it (14.27). With JiT-L, SCION achieves 5.89 FID without guidance and 3.47 with representation guidance, outperforming RCG with the ADM recipe (6.24).

Abstract

Image Generation

Explore similar work

From Pixels, Without Pre-training: Joint Generative and Self-Supervised Representation Learning in One Model

LoopREG: Representation Supervision Enables Test-Time Recurrent Depth in Diffusion Transformers

Training-Free Hidden-State Refinement for Flow-Matching Image Generators