The Cosine Schedule is Fisher-Rao-Optimal for Masked Discrete Diffusion Models
This provides a theoretical justification for an existing heuristic in diffusion models, which is incremental but clarifies a key design choice for practitioners.
The authors tackled the problem of selecting a discretization schedule for sampling from masked discrete diffusion models, showing that the optimal schedule under Fisher-Rao geometry matches the widely-used cosine schedule.
In this work, we study the problem of choosing the discretisation schedule for sampling from masked discrete diffusion models in terms of the information geometry of the induced probability path. Specifically, we show that the optimal schedule under the Fisher-Rao geometry recovers the popularly-used cosine schedule.