Build a complete rectified-flow / linear interpolant flow-matching generative model in pure PyTorch. Instead of a long discrete reverse chain that predicts noise, you learn a velocity field that transports Gaussian noise to data along straight paths, then sample by integrating a simple ODE with Euler and Heun steps — including the namesake reflow pairing.
python scaffold.py- 1. interpolate_linear
- 2. target_velocity
- 3. sample_uniform_time
- 4. sample_gaussian_noise
- 5. make_flow_batch
- 6. flow_matching_loss
- 7. sinusoidal_time_embedding
- 8. init_velocity_mlp
- 9. velocity_mlp_forward
- 10. make_mixture_dataset
- 11. flow_train_step
- 12. train_rectified_flow
- 13. linspace_timesteps
- 14. euler_step
- 15. euler_sample
- 16. heun_step
- 17. heun_sample
- 18. sample_quality_mse
- 19. make_reflow_pairs
- 20. rectified_flow_experiment
steps: 200
loss: 3.4820 -> 1.7328
noise baseline MSE: 2.4872
trained sample MSE: 0.8022
improvement (noise - sample): 1.6850
reflow pair shapes: (16, 2) (16, 2)
experiment sample/baseline: (0.5652052164077759, 2.4827353954315186)