Results on Dit-B

[nicolas-dufour]

Hi,
I'm trying to reproduce on DiT-B but i don't see any notable improvements. At 400k steps my FID goes from 33 for the SIT-B baseline to 29 for self flow.

have you made some experiments at smaller scales? Does Self-flow shows it's full potential only after a certain scale?
Would be great if you have some DiT-B results you can share!

[hila-chefer]

Hi @nicolas-dufour, thanks for reaching out and sharing your results!

While we haven't experimented directly with DiT-B, a 4-point drop in FID (from 33 to 29) is actually a solid improvement (depending on the model of course). Have you had a chance to look at the resulting generations or the stratified diffusion loss to see the qualitative differences?

Regarding your question about scale, Self-Flow doesn't require massive scale to show its potential (see Fig. 6), but there are a few architectural and dataset considerations that can help:

1. Possible mismatches in training (e.g., noise scheduling): As reported in our limitations section, noise scheduling determines the masking timesteps and can have an impact on Self-Flow. More optimal noise scheduling choices often lead to a more significant boost.

2. Layer selection: We use a fixed ratio for layer selection ($\ell_\theta = 0.3D$, $\ell_{\theta'} = 0.7D$). Have you adapted the student/ teacher layers accordingly?

3. ImageNet bias: As we point out in the paper, ImageNet experiments are inherently biased in favor of DINO, since it has been trained extensively on ImageNet and additional data collected specifically to match ImageNet classes. Because of this, seeing improvements on ImageNet over baselines like REPA takes more training steps than our other experiments (T2I, T2V, T2A). I generally recommend evaluating on more than just one dataset, and specifically trying to rely less exclusively on ImageNet results to gauge full performance.

Hope this helps!

[nicolas-dufour]

Hey @hila-chefer,

Thanks for the complete response!

My feeling that 29 was not much comes from "TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training", Krause et al., ICCV 2025 (https://arxiv.org/abs/2501.04765) which achieves FID of 15 with the same SiT B/2 config, and with better compute efficiency since they drop 50% of the tokens.

I do ajust the layer selection so it's around the 30% layer matching the 70% layer.

I'm going to try sweeping noise schedulers! Thanks for the tip. In your experience do you want to noise more or less than the standard diffusion scheduler?

I totally agree on your last point! I'm all for having an end-to-end solution without external knowledge injected!

[hila-chefer]

A few thoughts:

TREAD and Self-Flow operate on different axes. TREAD modifies how computation flows through the network during training. Self-Flow/ REPA/..., on the other hand, keep the full computation intact and instead modify what the model is asked to learn via an SSL objective.
TREAD's own results suggest these methods are complementary, not competing. Looking at the table below in their paper (DiT-B/2 at 400K): REPA alone gets 34.63, TREAD alone gets 21.34, and TREAD+REPA together get 20.06. The same pattern holds at 800K.

The table below is actually very relevant to your question specifically. TREAD's gains vary across architectures, for instance, it improves DiT-XL by ~62% but RWKV by only ~10%. Self-Flow/ REPA, in contrast don't target architectural optimizations but the training objective itself. Notice that the objective also differs between the two where Self Flow employs rectified flow whereas TREAD employs EDM.
Rather than competing with TREAD, these are complementary: TREAD improves training efficiency through routing, while Self-Flow improves what the model learns through the objective.

My suggestion: either combine Self-Flow with TREAD (for smaller models/ DiT & SiT architectures, as done in their paper with REPA), or compare apples to apples with feature alignment.

Noise scheduler choice is very important not just for Self-Flow, but for any diffusion model :) we give more details on that in Appendix A.