Rabiul Awal

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 FeCo loss features agree
The two views are the same image at two points on one noise-to-data path. A flow map should land in the same place with one long jump or with two short jumps through the cleaner view. FeCo also asks an early block on the noisier view to predict a deep block on the cleaner view.

TL;DR: A flow map can supervise its own features. Its composition rule gives two views of an image and a teacher with no extra forward pass, and no external encoder.

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, that is, flow maps, I noticed 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.
FeCo (previous version, without SIGReg)
MeanFlow with FeCo (previous version, 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.
FeCo (previous version, without SIGReg)
Shortcut with FeCo (previous version, without SIGReg), one-step samples with the same classes and noise as the baseline.

Uncurated. Each row is one class; both grids in a pair use the same classes and noise. Models trained for 1.2M steps. These samples come from an earlier version of FeCo (feature loss only, without SIGReg), not the final method.

Method

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

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

Why it is free

Every flow-map family builds its training target by moving along the path to a cleaner point. FeCo reads the teacher features at that point, so it needs no extra forward pass.

FamilyTeacher viewAlready there because
Shortcut ModelsEMA half stepits target chains two half steps
MeanFlowtangent from its JVPits target runs a JVP
LSDlanding point of its target forwardits target runs that forward

What remains 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, FeCo alone 51.8, FeCo plus SIGReg 48.9. Baseline: 57.3Baseline57.3 REPA (DINOv2): 51.2REPA (DINOv2)51.2 SIGReg alone: 56.3SIGReg alone56.3 FeCo alone: 51.8FeCo alone51.8 FeCo + SIGReg: 48.9FeCo + SIGReg48.9 48 52 56 60 1-step FID (lower is better)
MeanFlow SiT-B/2, one-step FID, lower is better. SIGReg alone barely helps; on top of FeCo it adds 2.9 and passes REPA. The faint dotted line marks REPA. Preliminary, ImageNet-100 speedrun.
FamilyBaselineFeCo+ SIGRegREPA
MeanFlow (SiT-B/2)57.351.848.951.2
Shortcut (DiT-B/2)139.9135.4116.8–
LSD (SiT-B/2)–72.764.6–

Preliminary, ImageNet-100 speedrun: one-step FID, mean of 2 seeds, 39.5k training steps, not the paper budget. REPA aligns with a frozen DINOv2 encoder. No LSD baseline in this study.

Main results

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

BibTeX

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