PaperlayerView exact revision on arXiv

Source-linked AI summary

RepFlow: Reciprocal Supervision Improves Generation and Representation in Flow Models

Weili Zeng, Feng Tian, Shengqi Liu, Yichao Yan

arXiv:2609.33217v1cs.CV

TL;DR

Generator states vary with noise level and network depth, complicating the extraction of stable visual representations. RepFlow trains a timestep-free encoder to predict those states from masked clean images, then feeds its representation back to guide generation. On ImageNet, this improves multi-step generation and linear probing with SiT and JiT, and the representation also supports one-step JiT post-training.

  • Problem

    Generator states are entangled with noise level and network depth, complicating the extraction of a stable visual representation from the generator.

    Source: §1 · p. 1

  • Method

    RepFlow trains a separate timestep-free encoder to predict generator states across depths and noise levels from masked clean images, then uses its full-image features to guide the generator.

    Source: §1 · p. 2

  • Results

    On ImageNet, unconditional FID falls from 32.83 to 19.43 on SiT and from 35.58 to 28.76 on JiT, while linear-probe accuracy improves 10.04 and 6.80 points over each model’s best searched raw features.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Under matched budgets, the frozen RepFlow representation also guides one-step JiT distribution matching, lowering FID from 6.34 to 3.14 against a pretrained MAE feature.

    Source: §1 · pp. 2–3

Abstract

from arXiv · show

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.

1 INTRODUCTION

RepFlow trains a clean-image representation from a generator’s evolving states, then feeds that representation back to guide generation. On ImageNet, it improves both multi-step generation and frozen linear probing, and the learned representation also supports one-step JiT post-training.

1 INTRODUCTION

  • Generator states vary with depth and noise level, but offer related observations of the same image that can supervise one clean-image representation.The challenge is that a feature useful at one layer and timestep may behave differently elsewhere in the trajectory.Source: §1 · pp. 1–2
  • RepFlow closes a feedback loop: a separate encoder predicts generator states from masked clean images, and its full-image features guide generator updates that change later prediction targets.The encoder receives neither the reference noise nor masked-out content, and timestep conditioning is confined to the predictor.Source: §1 · p. 2
  • FID falls from 32.83 to 19.43 on SiT and from 35.58 to 28.76 on JiT; linear-probe accuracy rises 10.04 and 6.80 points over each model’s best searched raw features.The results are for unconditional ImageNet experiments; class-conditional results are competitive with REPA and SRA.Source: §1 · p. 2
  • In matched-budget one-step JiT post-training, replacing a pretrained MAE feature with the RepFlow student lowers FID from 6.34 to 3.14.This extends the learned representation from instance-level alignment during reciprocal learning to distribution-level supervision afterward.Source: §1 · pp. 2–3
  • RepFlow addresses the shared visual-structure objective of generation and representation learning without an externally pretrained representation teacher.The paper distinguishes generators’ need to retain instance-specific detail from representation learners’ incentive to discard nuisance variation.Source: §1 · pp. 1, 3

2 PRELIMINARIES

The preliminaries set out flow-based generation and masked predictive learning. RepFlow applies a timestep-conditioned predictor to generator-state targets while keeping the encoder’s input independent of the generative stage.

2.1 FLOW-BASED GENERATIVE MODELING

  • Flow modeling interpolates between clean data and independent noise, with t=0 denoting data and t=1 noise; SiT predicts velocity, while JiT predicts clean data.For JiT, the experiments use a linear path and a weighted reconstruction form of its native velocity-space loss.Source: §2.1 · p. 3

2.2 PREDICTIVE REPRESENTATION LEARNING

  • Masked predictive learning trains an encoder to predict targets from partial image observations, using an optional conditioning variable in the predictor.The prediction objective stops gradients through the target; in RepFlow, generator states supply the target and the generative timestep conditions only the predictor.Source: §2.2 · p. 3
  • Generator states across depth and noise level serve as prediction targets for a clean-image encoder, whose input remains independent of the target’s generative stage.The same image induces different feature maps at different stages of computation.Source: §2.2 · p. 3, Figure 2 · p. 4

3 METHOD

RepFlow trains a separate image encoder to predict generator states across depths and noise levels, then uses its full-image features to guide the generator. The resulting generator updates change the states that supervise subsequent representation learning, and the learned feature space is also reused for one-step JiT distribution matching.

3.1 FROM GENERATIVE STATES TO PREDICTIVE TARGETS

  • Generator states at different depths and noise levels provide related prediction tasks about the same image, enabling a single clean-image representation to learn from their shared structure.The squared-error analysis attributes irreducible prediction error to target variation unresolved by timestep conditioning; the implemented objective uses Smooth L1.Source: §3.1 · p. 4

3.2 LEARNING THE SUPERVISORY REPRESENTATION

  • A timestep-conditioned predictor trains the encoder to recover normalized generator features from selected depths at masked positions, using a masked clean-image view.The encoder receives no class label, reference noise, or timestep; the loss updates only the encoder and predictor.Source: §3.2 · pp. 4–5
  • The encoder must support predictions across several depths and sampled noise levels from image context alone, while the predictor handles timestep-dependent decoding.As the generator evolves under its native objective, the student’s prediction targets evolve with it.Source: §3.2 · p. 5

3.3 RECIPROCAL TRAINING

  • Full-image EMA-encoder features guide an early generator layer through token-wise cosine alignment, while stopped targets separate the two gradient paths.Aligning at an early layer lets updates affect states processed by subsequent generator blocks; reciprocity occurs as each model’s updates shape supervision available to the other at later updates.Source: §3.3 · pp. 5–6

3.4 REPRESENTATION-GUIDED ONE-STEP POST-TRAINING

  • RepFlow reuses a frozen EMA encoder’s feature space for one-step JiT distribution matching, extending the representation’s role beyond instance-level alignment.Real statistics and historical generated moments are fixed during backpropagation; gradients through current generated-image features update the generator while the encoder remains frozen.Source: §3.4 · p. 6

4 EXPERIMENTS

The experiments test RepFlow’s generation and representation quality under matched ImageNet evaluations, then examine reciprocal-supervision choices, class conditioning, and one-step post-training.

4.1 EXPERIMENTAL SETUP

  • Evaluations use matched native/RepFlow backbones, data, epoch budgets, samplers, and evaluators, reporting FID-50K and frozen ImageNet top-1 accuracy.The experiments also test class-conditional generation and one-step post-training; RepFlow adds computational cost per update.Source: §4.1 · pp. 6–7, §4 · p. 6

4.2 UNCONDITIONAL GENERATION AND REPRESENTATION

  • After 600 epochs, RepFlow-L reduces unconditional SiT FID from 32.83 to 19.43 (40.8%) and JiT FID from 35.58 to 28.76 (19.2%).RepFlow-L reaches the native 600-epoch FID after approximately 143 epochs on SiT and 114 epochs on JiT.Source: §4.2 · p. 7, Figure 4 · p. 7
  • RepFlow-L improves frozen linear-probe accuracy over the best searched raw generator features by 10.04 percentage points on SiT and 6.80 on JiT.RepFlow uses the final-layer, full-view EMA student feature; increasing student capacity improves probing on both backbones, while JiT FID changes little between sizes.Source: §4.1 · p. 6, §4.2 · p. 7, Figure 5 · p. 7

4.3 ANALYZING RECIPROCAL SUPERVISION

  • Feedback lowers FID by 13.40 on SiT and 6.82 on JiT, while increasing student accuracy by a further 2.71 and 2.40 percentage points, respectively.R2G changes later generator states used as G2R targets; it supplies no direct gradient to the student.Source: §4.3 · p. 7, Table 1 · p. 7
  • With the student, masking ratio, feedback branch, data, and training budget retained, predicting generator states outperforms masked-pixel reconstruction on both FID and linear-probe accuracy.Multi-level targets lower FID by 1.28 on SiT and 1.68 on JiT and improve linear-probe accuracy by 2.80 and 1.42 percentage points versus the strongest evaluated single-depth target.Source: §4.3 · p. 8

4.4 CLASS-CONDITIONAL GENERATION

  • At 600 epochs with class conditioning, RepFlow reduces SiT FID from 9.62 to 6.99 (27.3%) and JiT FID from 16.80 to 12.65 (24.7%).Against SRA, RepFlow has slightly higher SiT FID but higher IS and recall; on JiT it has lower FID and sFID and higher IS and recall. REPA retains lower FID on both backbones, with gaps of 0.97 on SiT and 0.64 on JiT, while RepFlow achieves higher recall.Source: §4.4 · p. 8, Table 2 · p. 8

4.5 REPRESENTATION-GUIDED ONE-STEP POST-TRAINING

  • In matched one-step JiT post-training, replacing the MAE feature with the RepFlow student lowers FID from 6.34 to 3.14, a 50.5% reduction.FDr6 decreases from 9.33 to 8.06; IS and precision increase, while recall remains 0.56. The student supplies distribution-level supervision without using any of the six evaluation encoders in its objective.Source: §4.5 · pp. 8–9, Table 3 · p. 9

5 CONCLUSION

RepFlow learns representation guidance from a generator’s evolving computation and feeds it back to improve generation. On ImageNet, it improves both multi-step generation and frozen linear probing with latent-space SiT and pixel-space JiT, including one-step JiT distribution matching.

5 CONCLUSION

  • A timestep-free encoder predicts generator states across depths and noise levels from masked clean images, and its feedback changes the states that supervise later representation learning.The supervisory signal is thus shaped by the same training process it helps guide.Source: §5 · p. 9
  • RepFlow improves multi-step generation and frozen linear probing with latent-space SiT and pixel-space JiT, without an externally pretrained representation teacher.The frozen student also supports one-step JiT distribution matching, extending its use beyond joint training.Source: §5 · p. 9

6 AI USE STATEMENT.

The appendix discloses how generative AI tools were used in manuscript preparation and situates RepFlow among predictive representation learning, generator features, representation-guided generation, and one-step flow models.

6 AI USE STATEMENT.

  • Generative AI tools assisted with drafting, readability, and formatting, but not data work, model development, analysis, or interpretation; the authors reviewed the work and take responsibility.The disclosure specifies that tools were not used to generate or clean data or translate research content.Source: §6 · p. 10

A.1 PREDICTIVE REPRESENTATION LEARNING AND GENERATIVE FEATURES

  • RepFlow takes a predictive-learning approach, using generator internal computation as targets rather than another clean-view encoder.Unlike DIFT and CleanDIFT, its student learns jointly from multiple generator depths and returns features to the generator.Source: §A.1 · p. 13

A.2 REPRESENTATION-GUIDED GENERATIVE MODELING

  • RepFlow learns its guiding encoder concurrently from generator states, avoiding external representation pretraining but retaining the cost of an additional student.The separate student–generator interface is studied with both latent-space SiT and pixel-space JiT.Source: §A.2 · pp. 13–14

A.3 FLOW MODELS AND ONE-STEP GENERATION

  • RepFlow keeps joint training multi-step and uses the learned student only as the metric for a short JiT post-training stage.This differs from methods that directly learn finite-time maps or one-step generation objectives.Source: §A.3 · p. 14

B OBJECTIVE AND OPTIMIZATION ANALYSIS

The analysis characterizes what generator-state prediction rewards and how stopped-gradient paths create feedback across optimization steps. The implementation instantiates the separate student and generator interface across ImageNet, SiT, and JiT experiments.

B.1 G2R AS EXPLAINED TEACHER VARIATION

  • G2R rewards a student for explaining generator-state variation beyond timestep alone, while a Bayes-optimal code need only preserve the masked-view conditional prediction profile.Components depending on reference noise or unobserved image content remain irreducible from the masked view.Source: §B.1 · p. 15
  • The objective does not guarantee semantic features; masking and a shared timestep-free student provide inductive bias, and ImageNet linear probing evaluates class separation.Experiments use Smooth L1 rather than the squared-loss surrogate used for the variance identities.Source: §B.1 · p. 15

B.2 GRADIENT ROUTING AND RECIPROCAL DYNAMICS

  • Stop-gradient removes direct G2R-to-generator and EMA-target-to-student gradients, yet reciprocal feedback arises across updates as each path changes later targets.At an idealized stationary fixed point, the student summarizes predictable variation in the current generator while the generator is stationary under its native objective and student alignment.Source: §B.2 · p. 15

C IMPLEMENTATION AND EVALUATION DETAILS

  • Experiments use ImageNet-1K at 256×256 resolution, with a global batch of 1024 and eight NVIDIA H200 GPUs; the implementation reports timesteps in τ = 1 − t.The main text parameterizes the path from data to noise, while the implementation uses the reversed timestep coordinate.Source: C IMPLEMENTATION AND EVALUATION DETAILS · p. 16, §C.1 · p. 16

C.1 DATA AND GENERATIVE BACKBONES

  • SiT uses frozen-VAE 32×32×4 latents, while JiT operates directly on RGB patches; both backbones have 28 transformer blocks.The SiT student receives the corresponding RGB image, and JiT predicts x0 using an unweighted velocity MSE objective.Source: §C.1 · p. 16

C.2 REPFLOW INSTANTIATION

  • The student is an RGB-only, timestep-free encoder; G2R predicts detached, normalized states from blocks {5, 12, 19, 26}, while R2G aligns block 7 to an EMA full-view student.G2R masks 75% of 256 tokens; the predictor receives timestep conditioning, and R2G uses per-token cosine distance with EMA decay 0.9999.Source: §C.2 · pp. 16–17

C.3 BASELINES AND EVALUATION PROTOCOLS

The evaluation compares RepFlow with native and representation-guided generation controls, evaluates generation and frozen linear probing, and tests pixel-reconstruction and single-depth alternatives.

C.3 BASELINES AND EVALUATION PROTOCOLS

  • The unconditional generator control retains G2R but sets λ2 = 0, making its generator update identical to native training; R2G is the only additional generator gradient in the comparisons.The class-conditional native baseline omits the student.Source: §C.3 · p. 17
  • Generation evaluation reports FID-50K and other metrics from 50,000 EMA samples; one-step JiT uses a single terminal-noise forward pass.SiT uses the adaptive ODE sampler, JiT uses Heun, and all evaluations use CFG = 1.Source: §C.3 · p. 17
  • Frozen probing averages the EMA student’s 256 final-layer tokens, while the raw-generator baseline selects layer 21 at τ = 0.7 for SiT and τ = 0.5 for JiT.The raw-feature selection is used only for its baseline; RepFlow evaluates one full-view student feature.Source: §C.3 · p. 17
  • The pixel-target ablation replaces detached multi-level generator targets with masked RGB-patch reconstruction, while the target-depth ablation predicts only block 19 for SiT or block 20 for JiT.The multi-level variant predicts blocks {5, 12, 19, 26}.Source: §C.3 · p. 17

C.4 ONE-STEP POST-TRAINING

The class-conditional one-step experiment updates a frozen-checkpoint generator using distributional statistics in a frozen feature space, with an MAE-feature control and several feature-space evaluation metrics.

C.4 ONE-STEP POST-TRAINING

  • The experiment starts from a 400-epoch RepFlow-JiT checkpoint with its ViT-L student frozen, tracks generated feature moments over an effective population of 50,000, and updates only the generator.Real statistics use layer-normalized, mean-pooled ImageNet features; generated moments use EMA decay 0.999, and evaluation uses CFG = 1.Source: §C.4 · p. 17
  • The matched control changes only the frozen feature from RepFlow ViT-L/16 to ImageNet-pretrained MAE ViT-L/16, keeping initialization, samples, optimizer, schedule, duration, and real-statistics protocol fixed.Source: §C.4 · p. 17
  • Post-training uses AdamW at learning rate 10^-5 with cosine decay, 6,250 warmup steps, global batch size 1024, and 30 additional epochs (37,500 steps).Source: §C.4 · p. 17
  • Evaluation averages normalized Fréchet distances across six feature spaces, while RepFlow-student FD uses a feature space distinct from all six.FD-SIM optimizes SigLIP, Inception, and MAE distributions, overlapping FID and FDr6 evaluation.Source: §C.4 · p. 18

C.5 TRAINING COST AND STATISTICAL SCOPE

RepFlow's generative gain persists when matched to Native wall-clock compute, while the reported training costs show additional time for RepFlow at 600 epochs.

C.5 TRAINING COST AND STATISTICAL SCOPE

  • At 600 epochs, SiT takes about 100, 128, and 141 hours for Native, RepFlow-L, and conditional RepFlow-L; JiT takes 66 and 99–105 hours for Native and RepFlow-L.Throughput excludes sampling, and all runs use seed 0.Source: §C.5 · p. 18
  • Table 5 reports measured system cost and the RepFlow checkpoint matched to a 600-epoch Native run, computed as 600 × sRF/sNative.Measurements use global batch size 1024 on eight H200 GPUs; SiT memory was not logged.Source: Table 5 · p. 18
  • At all four matched checkpoints, RepFlow remains below the 600-epoch Native endpoint, so its generative gain persists under matched wall-clock compute rather than only equal data exposure.Source: §C.5 · p. 18

D UNCURATED ONE-STEP SAMPLES

The appendix presents random, uncurated one-step ImageNet samples as generated, without manual selection or post hoc filtering, across the class ranges shown in Figures 8–12.

D UNCURATED ONE-STEP SAMPLES

Loading 2609.33217v1…