OptimalTransportScheduler

class OptimalTransportScheduler

Bases: Scheduler[DDTensor]

Optimal transport scheduler which is commonly used to train flow matching models.

Schedule: .. math:

\alpha_t = t, \quad \beta_t = 1 - t, \quad \dot{\alpha}_t = 1, \quad \dot{\beta}_t = -1.

alpha(x: DDTensor, t: Tensor) → DDTensor

\(\alpha_t = t\).

alpha_dot(x: DDTensor, t: Tensor) → DDTensor

\(\dot{\alpha}_t = 1\).