Rabiul Awal

A Self-Supervised Free Lunch for Flow Map Training

Rabiul Awal

Mila, Université de Montréal

Work in progress

A noise-to-data path with four images: noise, a noisier view, a cleaner view and data. One long jump goes from the noisier view to data; two short jumps go through the cleaner view. Below, an early block on the noisier view predicts a deep block on the cleaner view. one long jump two short jumps noise noisier view cleaner view data online model EMA model early block deep block predicts Flow-map loss outputs agree Consistency loss features agree A noise-to-data path with four images: noise, a noisier view, a cleaner view and data. One long jump goes from the noisier view to data; two short jumps go through the cleaner view. Below, an early block on the noisier view predicts a deep block on the cleaner view. Flow-map loss outputs agree one long jump two short jumps noise noisier view cleaner view data online model EMA model early block deep block predicts Consistency loss features agree
A flow map should land in the same place with one long jump from the noisier view or two short jumps through the cleaner view. iFeCo applies this rule to features and asks an early block on the noisier view to predict a deep block on the cleaner view.

TL;DR: Flow-map training already distills a teacher on a cleaner view into a student on a noisier view. iFeCo extends this distillation to features. It needs no external encoder and no extra forward pass.

Why this matters

Diffusion training seems difficult partly because of the semantic structure that has to come out of it. We have seen representation autoencoders, representation generation and representation alignment. Much of this work says the same thing: diffusion training needs some form of self-supervised signal to train faster and generate better.

Since I was mostly working on few-step models (flow maps), I noticed that they are, in principle, a knowledge distillation algorithm: a student tries to distill from a multi-step teacher. Looked at closely, the student and the teacher see two views of the same data. The teacher gets the cleaner one and the student the noisier one. This gives another dimension of consistency, in feature space, without any additional forward pass. That is the first free lunch for flow map training.

Then I realized that even when you align features across these free views, their quality is still below what self-supervised learning gives. A distributional regularizer like SIGReg can also be applied to diffusion to shape its representations. That is the second free lunch.

One-step samples, ImageNet 256×256

MeanFlow SiT-B/2, 1 step

Baseline
MeanFlow baseline, one-step samples: four classes, four samples each.
Consistency only (without SIGReg)
MeanFlow with the consistency loss only (without SIGReg), one-step samples with the same classes and noise as the baseline.

Shortcut DiT-B/2, 1 step

Baseline
Shortcut baseline, one-step samples: four classes, four samples each.
Consistency only (without SIGReg)
Shortcut with the consistency loss only (without SIGReg), one-step samples with the same classes and noise as the baseline.

Uncurated, after 1.2M training steps. Each row is one class, and both grids in a pair share classes and noise. iFeCo samples (with SIGReg) will follow.

Method

The online model reads the noisier view, and a predictor h on its early block feeds the consistency loss. The EMA model reads the cleaner view, and its deep block gives the target through a stop-gradient. SIGReg acts on the online model's features at the teacher's depth. online model noisier view SIGReg pushes features toward an isotropic Gaussian flow-map loss EMA copy predictor h predicts consistency loss stop-gradient cleaner view EMA model early block deep block The online model reads the noisier view, and a predictor h on its early block feeds the consistency loss. The EMA model reads the cleaner view, and its deep block gives the target through a stop-gradient. SIGReg acts on the online model's features at the teacher's depth. pushes features toward an isotropic Gaussian SIGReg online model noisier view flow-map loss EMA copy early block predictor h predicts consistency loss stop-gradient cleaner view EMA model deep block

Both terms are added to the usual flow-map loss.

Why it is free

Each flow-map family already passes through the cleaner view when it builds its training target. iFeCo reads the teacher from that pass, so it needs no extra forward pass.

Shortcut target: two EMA half steps student jump EMA half steps noise data teacher: EMA at the start of the second half step MeanFlow target: a JVP along the path student jump first-order step noise data teacher: model features moved along the path by the JVP LSD target: EMA velocity at landing student jump EMA velocity noise data teacher: EMA at the landing point
Black dot: the noisier view, where the student starts. Orange ring: the cleaner view, where the teacher reads its deep features. Dashed: steps the training target already takes. Shortcut Models, MeanFlow, LSD.

The only added cost is the small predictor and the SIGReg statistic.

Preliminary results (ImageNet-100 speedrun)

One-step FID for MeanFlow SiT-B/2 on the ImageNet-100 speedrun: baseline 57.3, REPA 51.2, SIGReg alone 56.3, consistency alone 51.8, iFeCo (consistency plus SIGReg) 48.9. Baseline: 57.3Baseline57.3 REPA (DINOv2): 51.2REPA (DINOv2)51.2 SIGReg alone: 56.3SIGReg alone56.3 Consistency alone: 51.8Consistency alone51.8 iFeCo: 48.9iFeCo (both)48.9 48 52 56 60 1-step FID (lower is better)
SIGReg alone barely helps, but on top of the consistency loss it lowers FID by 2.9 and passes REPA (dotted line). MeanFlow SiT-B/2, one-step FID, preliminary ImageNet-100 speedrun.
FamilyBaselineConsistencyiFeCo (+ SIGReg)REPA
MeanFlow (SiT-B/2)57.351.848.951.2
Shortcut (DiT-B/2)139.9135.4116.8137.3
LSD (SiT-B/2)76.872.764.674.6

One-step FID (lower is better), mean of 2 seeds after 39.5k steps on the preliminary ImageNet-100 speedrun, shorter than the paper budget. REPA aligns features with a frozen DINOv2 encoder.

Main results

ImageNet-1k 256×256 results at standard training budgets are coming soon.

BibTeX

@article{awal2026feco,
  title  = {A Self-Supervised Free Lunch for Flow Map Training},
  author = {Awal, Rabiul and others},
  year   = {2026},
  note   = {Preprint coming soon}
}