Time dependent loss reweighting for flow matching and diffusion models is theoretically justified

TL;DR

This paper provides a theoretical foundation for time-dependent loss weighting in flow matching and diffusion models, showing that such schemes are justified and can simplify model construction.

Contribution

It demonstrates that time-dependent loss functions and generator parameterizations are theoretically valid within Generator Matching and Edit Flows frameworks.

Findings

Time-dependent loss weighting schemes are theoretically justified.

Dependence of loss and generator on time $t$ can simplify model design.

Examples show the practical benefits of time-dependent parameterizations.

Abstract

This brief note clarifies that, in Generator Matching (which subsumes a large family of flow matching and diffusion models over continuous, manifold, and discrete spaces), both the Bregman divergence loss and the linear parameterization of the generator can depend on both the current state $X_t$ and the time $t$, and we show that the expectation over time in the loss can be taken with respect to a broad class of time distributions. We also show this for Edit Flows, which falls outside of Generator Matching. That the loss can depend on $t$ clarifies that time-dependent loss weighting schemes, often used in practice to stabilize training, are theoretically justified when the specific flow or diffusion scheme is a special case of Generator Matching (or Edit Flows). It also often simplifies the construction of $X_1$-predictor schemes, which are sometimes preferred for model-related reasons. We show examples that rely upon the dependence of linear parameterizations, and of the Bregman divergence loss, on $t$ and $X_t$.