Geometry-Guided Temporal Allocation for Flow Matching Distillation

Flow matching and diffusion models achieve state-of-the-art image synthesis but require slow, iterative inference. While few-step distillation reduces this cost, these methods conventionally distribute training updates uniformly across time. We propose Complexity-Weighted Rectified Distillation (CWRD), a plug-and-play training wrapper that uses a frozen teacher-geometry profile to reweight timestep sampling. The profile is based on Local Complexity (LC), which measures how the dominant endpoint-sensitivity subspace varies across neighboring latent states. CWRD increases the relative sampling probability of high-LC timesteps, reallocating temporal capacity allocation while strictly preserving the host method's per-example loss, student architecture, and inference procedure. We evaluate CWRD on $\pi$-Flow, LADD, TDM, and MFD using high-resolution Stable Diffusion 3.5 and SANA teachers. Across the evaluated configurations, CWRD reduces teacher-referenced FID and improves semantic prompt alignment, with absolute data-referenced FID improvements dependent on the specific host architecture. Complementary trajectory diagnostics show that following a teacher-trajectory prefix before a single suffix update can preserve the teacher output better than uniform sampling at a matched evaluation budget. These findings suggest that endpoint-sensitivity variation provides a useful empirical signal for temporal allocation in few-step generative models.