Skip to content

README.md

Latest commit

 

History

History
692 lines (539 loc) · 16 KB

README.md

File metadata and controls

692 lines (539 loc) · 16 KB
Raw

Reanchored Auxiliary-Free Data-Free Step Distillation

This repo is a PyTorch/DDP research implementation for data-free flow-map distillation experiments inspired by FreeFlow, with modular switches for:

  • baseline prior-anchored FreeFlow-style prediction;
  • original auxiliary correction baseline (g_psi);
  • FSF auxiliary-free correction;
  • self-denoising fake endpoint training;
  • reanchored/on-policy intermediate states;
  • hybrid aux→FSF correction;
  • optional CFG-conditioned teacher queries.

The main research path is:

ImageNet-1k from Hugging Face
→ SD-VAE ImageNet-256 latents [4, 32, 32]
→ DiT latent teacher/student
→ data-free student distillation
→ FID/KID + teacher/student paired metrics logged to W&B

CIFAR/UNet is only for smoke/debug. Do not use CIFAR/UNet as the main paper path.

Important teacher checkpoint clarification

FreeFlow's public Hugging Face checkpoints are distilled student checkpoints, not the original teacher velocity model needed to train a new data-free student. A FreeFlow-style distillation run needs a teacher velocity model u_T(x,t,y).

So there are two cases:

Case A — you already have a compatible teacher velocity checkpoint

Skip teacher training. Set:

export TEACHER=/path/to/compatible_teacher_checkpoint.pt

and go directly to student ablations.

A compatible checkpoint means this repo can load it as a DiT velocity model using freeflowx.train_distill. If you have an official SiT/FreeFlow teacher checkpoint in a different format, it may need a converter/loader first.

Case B — you do not have a compatible teacher velocity checkpoint

Use train_teacher.py to train a local ImageNet latent rectified-flow DiT teacher. This is a fallback so the end-to-end pipeline runs from a fresh server. It is not claiming to reproduce the official FreeFlow teacher.

Do not confuse this with the released FreeFlow student weights. You cannot use the released one-step student checkpoint as the teacher for the original FreeFlow objective unless you intentionally switch to a different experiment: redistilling/compressing FreeFlow itself.

0. Fresh cloud server assumptions

Recommended:

  • Linux cloud GPU box;
  • CUDA-capable PyTorch environment;
  • 2+ GPUs preferred for DDP;
  • large disk for ImageNet + cached latents;
  • Hugging Face account with access to ILSVRC/imagenet-1k;
  • W&B account for logging.

You must accept the ImageNet terms on Hugging Face once in a browser before the CLI download works.

1. Unpack repo

unzip freeflowx-imagenet-hf-readme-fixed.zip
cd freeflowx_v2

2. Conda environment + editable install

conda env create -f environment-cuda.yml
conda activate freeflowx
pip install -e .

Check CUDA:

python - <<'PY'
import torch
print('torch:', torch.__version__)
print('cuda available:', torch.cuda.is_available())
print('num gpus:', torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
    print(i, torch.cuda.get_device_name(i))
PY

Run unit tests:

pytest -q

3. Login for W&B and Hugging Face

wandb login
huggingface-cli login

Alternative non-interactive HF login:

export HF_TOKEN=<your_huggingface_token>
HF_TOKEN=$HF_TOKEN bash scripts/setup_imagenet_hf.sh

4. Download/check ImageNet-1k from Hugging Face

This checks both validation and train splits through the Hugging Face dataset loader.

bash scripts/setup_imagenet_hf.sh

Expected dataset interface:

image: PIL image
label: int in 0..999

5. Precompute ImageNet-256 SD-VAE latents

FreeFlow-style latent DiT models train on SD-VAE latents rather than raw pixels. This command encodes ImageNet images to scaled latents:

x_latent: [4, 32, 32]
y: int class label

Run:

BATCH=128 WORKERS=8 SHARD_SIZE=10000 VAE=ema bash scripts/prepare_imagenet256_latents.sh

Outputs:

data/imagenet256_latents/train/shard_00000.pt
data/imagenet256_latents/train/shard_00001.pt
...
data/imagenet256_latents/validation/shard_00000.pt
...

Verify the latent dataset:

python -m freeflowx.data --dataset imagenet256_latents --data-dir data --split train
python -m freeflowx.data --dataset imagenet256_latents --data-dir data --split validation

6. Optional smoke test

This only checks that the code runs. It is not a research experiment.

bash scripts/smoke.sh

7. Teacher setup

Option A: use an existing compatible teacher checkpoint

If you already have a compatible ImageNet latent DiT teacher checkpoint:

export TEACHER=/path/to/compatible_teacher_checkpoint.pt

Then skip to Section 8.

Option B: train a local teacher fallback

Only do this if you do not have a compatible teacher velocity checkpoint. This trains a local class-conditional rectified-flow DiT teacher on cached ImageNet-256 latents.

NGPU=2 \
STEPS=400000 \
BATCH=256 \
WANDB_PROJECT=freeflowx \
bash scripts/train_teacher_imagenet_latents_dit.sh

Default output:

runs/teacher_imagenet256_latents_dit_b/checkpoint.pt

Then set:

export TEACHER=runs/teacher_imagenet256_latents_dit_b/checkpoint.pt

Default teacher architecture:

backbone: DiT
hidden: 768
channels: 768
depth: 12
heads: 12
patch size: 2
input: [4, 32, 32]
classes: 1000

8. Train ablations individually first

Run the individual ablations before the combined method. This is the clean experimental order.

Set common variables:

export DATA_DIR=data
export NGPU=2
export STEPS=200000
export BATCH=256
export WANDB_PROJECT=freeflowx
export FID_SAMPLES=10000
export EVAL_EVERY=5000

8.1 Prediction-only baseline

No correction. No self-denoise. No reanchor.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_pred \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction none \
  --lambda-pred 1.0 \
  --lambda-corr 0 \
  --lambda-self-denoise 0 \
  --lambda-reanchor 0 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_pred

8.2 Original FreeFlow-style auxiliary correction baseline

Uses auxiliary noising-velocity model g_psi.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_aux \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --aux-channels 512 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction aux \
  --lambda-pred 1.0 \
  --lambda-corr 0.05 \
  --lambda-aux 1.0 \
  --lambda-self-denoise 0 \
  --lambda-reanchor 0 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_aux

8.3 Reanchored/on-policy states only, with auxiliary correction

Turns on reanchoring, but keeps original auxiliary correction. This isolates the effect of on-policy intermediate states.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_aux_reanchor \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --aux-channels 512 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction aux \
  --lambda-pred 1.0 \
  --lambda-corr 0.05 \
  --lambda-aux 1.0 \
  --lambda-self-denoise 0 \
  --lambda-reanchor 0.5 \
  --teacher-local-steps 2 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_aux_reanchor

8.4 FSF auxiliary-free correction + self-denoise

Turns on FSF and self-denoising fake endpoint training. No reanchor yet.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_fsf \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction fsf \
  --lambda-pred 1.0 \
  --lambda-corr 0.025 \
  --lambda-self-denoise 1.0 \
  --lambda-reanchor 0 \
  --ema-decay 0.999 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_fsf

9. Train the combined main candidate

After the individual ablations have run without exploding, train the combined method:

FreeFlow prediction
+ FSF auxiliary-free correction
+ self-denoising fake endpoint training
+ reanchored/on-policy intermediate states

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_fsf_reanchor \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction fsf \
  --lambda-pred 1.0 \
  --lambda-corr 0.025 \
  --lambda-self-denoise 1.0 \
  --lambda-reanchor 0.5 \
  --teacher-local-steps 2 \
  --ema-decay 0.999 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_fsf_reanchor

10. Optional hybrid auxiliary + FSF correction

This is the safer/stabler version if pure FSF is noisy. It starts with auxiliary correction and ramps toward FSF.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_hybrid \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --aux-channels 512 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --correction hybrid \
  --hybrid-start 0.0 \
  --hybrid-end 1.0 \
  --hybrid-ramp-steps 50000 \
  --lambda-pred 1.0 \
  --lambda-corr 0.05 \
  --lambda-aux 1.0 \
  --lambda-self-denoise 1.0 \
  --lambda-reanchor 0.5 \
  --teacher-local-steps 2 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_hybrid

11. Optional: one-command ablation matrix

After the setup is verified, this runs the ablation matrix in the correct order.

TEACHER=$TEACHER \
DATA_DIR=data \
NGPU=2 \
STEPS=200000 \
BATCH=256 \
FID_SAMPLES=10000 \
WANDB_PROJECT=freeflowx \
bash scripts/train_distill_matrix_imagenet_latents.sh

12. Optional CFG-conditioned run

This tests the teacher-data mismatch story under guided teachers. The teacher is queried with randomly sampled CFG scale, and the student receives the same scalar gw conditioning.

torchrun --nproc_per_node=$NGPU -m freeflowx.train_distill \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir $DATA_DIR \
  --outdir runs/student_imagenet256_latents_cfg_fsf_reanchor \
  --backbone dit \
  --teacher-backbone dit \
  --channels 768 \
  --teacher-channels 768 \
  --dit-hidden 768 \
  --dit-depth 12 \
  --dit-heads 12 \
  --patch-size 2 \
  --guidance cfg \
  --guidance-min 1.0 \
  --guidance-max 5.0 \
  --correction fsf \
  --lambda-pred 1.0 \
  --lambda-corr 0.025 \
  --lambda-self-denoise 1.0 \
  --lambda-reanchor 0.5 \
  --teacher-local-steps 2 \
  --ema-decay 0.999 \
  --steps $STEPS \
  --batch-size $BATCH \
  --eval-every $EVAL_EVERY \
  --fid-samples $FID_SAMPLES \
  --workers 8 \
  --wandb --wandb-project $WANDB_PROJECT --wandb-group imagenet256_latents_cfg_fsf_reanchor

13. Final paper-grade evaluation

During training, FID is logged with --fid-samples 10000 by default because full FID-50K is expensive. For paper tables, run final FID-50K.

Student FID-50K + teacher/student paired metrics

python -m freeflowx.eval \
  --student-ckpt runs/student_imagenet256_latents_fsf_reanchor/checkpoint.pt \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir data \
  --outdir runs/eval_fsf_reanchor_50k \
  --fid-samples 50000 \
  --batch-size 128 \
  --teacher-steps 50 \
  --guidance-w 1.0 \
  --wandb --wandb-project freeflowx --wandb-group final_eval

Teacher FID-50K baseline

python -m freeflowx.eval \
  --teacher-ckpt $TEACHER \
  --dataset imagenet256_latents \
  --data-dir data \
  --outdir runs/eval_teacher_50k \
  --fid-samples 50000 \
  --batch-size 128 \
  --teacher-steps 50 \
  --eval-teacher \
  --wandb --wandb-project freeflowx --wandb-group final_eval

14. How to turn each technique on/off

FSF auxiliary-free correction

Off:

--correction none

or original auxiliary baseline:

--correction aux --lambda-aux 1.0 --lambda-corr 0.05

On:

--correction fsf --lambda-corr 0.025 --ema-decay 0.999

Hybrid aux→FSF:

--correction hybrid --lambda-aux 1.0 --lambda-corr 0.05 --hybrid-start 0.0 --hybrid-end 1.0 --hybrid-ramp-steps 50000

Self-denoising fake endpoint training

Off:

--lambda-self-denoise 0

On:

--lambda-self-denoise 1.0

Reanchored/on-policy intermediate states

Off:

--lambda-reanchor 0

On:

--lambda-reanchor 0.5 --teacher-local-steps 2

Recommended order:

1. prediction-only
2. auxiliary correction
3. auxiliary + reanchor
4. FSF + self-denoise
5. FSF + self-denoise + reanchor
6. hybrid if pure FSF is unstable

15. Metrics to watch in W&B

Training losses:

train/loss
train/loss_pred
train/loss_corr
train/loss_aux
train/loss_reanchor
train/loss_self_denoise
train/grad_norm
train/lr

Paper-facing eval metrics:

eval/fid_student_1step
eval/kid_student_1step
eval/pair_mse_teacher_student
eval/path_mse
eval/student_mean_abs
eval/student_std

Interpretation:

prediction-only should be weakest.
aux correction is the FreeFlow-style baseline.
aux + reanchor isolates the benefit of on-policy intermediate states.
FSF + self-denoise tests whether auxiliary correction can be removed.
FSF + self-denoise + reanchor is the main candidate.
hybrid is the safer/stabler version if pure FSF is noisy.

Main paper table:

Method FID-50K ↓ KID ↓ Pair MSE ↓ Path MSE ↓ Extra aux net? Notes
Prediction-only no drift baseline
FreeFlow-style aux yes main baseline
Aux + reanchor yes on-policy states
FSF + self-denoise no auxiliary-free correction
FSF + self-denoise + reanchor no main candidate
Hybrid yes during training stable transition

16. Caveats

  • The official FreeFlow public repo provides model definitions, sampling/evaluation code, and pretrained student weights, but not the original training code.
  • This repo is a PyTorch/DDP training reimplementation for ImageNet latent DiT experiments.
  • The main research path is imagenet256_latents + --backbone dit.
  • train_teacher.py is a fallback for creating a compatible teacher checkpoint when no compatible public teacher checkpoint is available.
  • Hugging Face ImageNet access is gated. Accept the terms before running setup.