| """Noise schedules for masked discrete diffusion. |
| |
| alpha(t) is the retention probability: alpha(0)=1 (clean), alpha(1)=0 (fully masked). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Callable |
|
|
| import jax.numpy as jnp |
|
|
| ScheduleFn = Callable[[jnp.ndarray], jnp.ndarray] |
|
|
|
|
| def linear_schedule(t: jnp.ndarray) -> jnp.ndarray: |
| """alpha(t) = 1 - t. Default in MDLM / ReMDM.""" |
| return 1.0 - t |
|
|
|
|
| def linear_schedule_deriv(t: jnp.ndarray) -> jnp.ndarray: |
| return jnp.full_like(t, -1.0) |
|
|
|
|
| def cosine_schedule(t: jnp.ndarray) -> jnp.ndarray: |
| """alpha(t) = cos(pi * t / 2).""" |
| return jnp.cos(t * jnp.pi / 2.0) |
|
|
|
|
| def cosine_schedule_deriv(t: jnp.ndarray) -> jnp.ndarray: |
| return -(jnp.pi / 2.0) * jnp.sin(t * jnp.pi / 2.0) |
|
|
|
|
| SCHEDULE_MAP: dict[str, tuple[ScheduleFn, ScheduleFn]] = { |
| "linear": (linear_schedule, linear_schedule_deriv), |
| "cosine": (cosine_schedule, cosine_schedule_deriv), |
| } |
|
|
