Built and trained a flow-matching model on MNIST from scratch. The project begins with a single-step UNet denoiser, then adds time conditioning for iterative generation and class conditioning for controllable digit synthesis.

The model architecture, training setup, and all results below follow the original technical report.

I implemented the prescribed UNet as a contracting encoder and expanding decoder connected with skip connections. Conv layers use 3×3 kernels; stride-2 downsampling maps 28×28 MNIST images to a 7×7 bottleneck; transposed convolutions restore spatial resolution.

The unconditional model uses hidden dimension 128 and is trained with Adam (lr 1e-4), MSE loss, batch size 256, σ=0.5 input noise, and five epochs. It recovers digits well at the trained noise level. At σ=0.8 and 1.0, outputs degrade because the model was never trained for those distributions.

FCBlocks map scalar t to channel-wise signals at two decoder locations: after unflattening and after the first UpBlock. The model is trained on random timesteps with x_t = (1−t)x₀ + tx₁, predicting velocity x₁ − x₀.

The first tile is the training-loss curve; the next three are samples generated from pure noise after epochs 1, 5, and 10. The time-conditioned UNet uses hidden dimension 64, batch size 64, an initial 1e-2 learning rate with exponential decay, and 50 Euler steps at inference. Samples become legible by epoch 10, though they remain softer than class-conditioned outputs.

Two further FCBlocks consume one-hot class labels and combine their output with time conditioning at the same decoder stages. With p_uncond=0.1, labels are dropped during training so the model learns both conditional and unconditional predictions.

The first tile is the class-conditioned loss curve. The next three grids show four samples for every digit class after epochs 1, 5, and 10. During sampling, classifier-free guidance with γ=5 combines conditional and unconditional flows: u = u_uncond + γ(u_cond − u_uncond). By epoch 10, each row consistently corresponds to its target digit and is sharper than time-only generation.