File size: 3,462 Bytes
e0eb79a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Noise schedule functions for MDLM diffusion.
All functions operate on PyTorch tensors and are pure (no global state).
Convention: alpha(t) is the fraction of tokens that remain *unmasked*.
- alpha(0) = 1.0 (fully clean)
- alpha(1) = 0.0 (fully masked)
"""
from __future__ import annotations
import math
from collections.abc import Callable
import torch
from torch import Tensor
def linear_schedule(t: Tensor) -> Tensor:
"""Linear noise schedule: alpha(t) = 1 - t.
Args:
t: Diffusion time in [0, 1]. Any shape.
Returns:
Retention probability alpha_t, same shape as *t*.
"""
return 1.0 - t
def cosine_schedule(t: Tensor) -> Tensor:
"""Cosine noise schedule: alpha(t) = cos(pi/2 * t).
MDLM Appendix E.1 eq (92) ("Cosine"); the same function the craftax
repo names "cosine".
Args:
t: Diffusion time in [0, 1]. Any shape.
Returns:
Retention probability alpha_t, same shape as *t*.
"""
return torch.cos(t * (math.pi / 2.0))
def cosine_sq_schedule(t: Tensor) -> Tensor:
"""Cosine-squared noise schedule: alpha(t) = cos(pi/2 * t)^2.
MDLM Appendix E.1 eq (91) ("Cosine Squared", after Nichol & Dhariwal).
Previously registered under the name "cosine" in this repo; renamed so
the label "cosine" denotes the same function in both repos.
Args:
t: Diffusion time in [0, 1]. Any shape.
Returns:
Retention probability alpha_t, same shape as *t*.
"""
return torch.cos(t * (math.pi / 2.0)) ** 2
def linear_schedule_deriv(t: Tensor) -> Tensor:
"""Analytic d(alpha)/dt for the linear schedule."""
return torch.full_like(t, -1.0)
def cosine_schedule_deriv(t: Tensor) -> Tensor:
"""Analytic d(alpha)/dt for the cosine schedule."""
return -(math.pi / 2.0) * torch.sin(t * (math.pi / 2.0))
def cosine_sq_schedule_deriv(t: Tensor) -> Tensor:
"""Analytic d(alpha)/dt for the cosine-squared schedule."""
return -(math.pi / 2.0) * torch.sin(t * math.pi)
_SCHEDULE_MAP: dict[str, Callable[[Tensor], Tensor]] = {
"linear": linear_schedule,
"cosine": cosine_schedule,
"cosine_sq": cosine_sq_schedule,
}
_DERIV_BY_FN: dict[Callable[[Tensor], Tensor], Callable[[Tensor], Tensor]] = {
linear_schedule: linear_schedule_deriv,
cosine_schedule: cosine_schedule_deriv,
cosine_sq_schedule: cosine_sq_schedule_deriv,
}
def get_schedule_deriv_for(
schedule_fn: Callable[[Tensor], Tensor],
) -> Callable[[Tensor], Tensor]:
"""Analytic derivative for a registered schedule function.
The NELBO weight uses the analytic d(alpha)/dt as stated in
MDLM eq (10) / Shi eq (4).
Raises:
KeyError: If *schedule_fn* is not a registered schedule.
"""
if schedule_fn not in _DERIV_BY_FN:
raise KeyError(
"No analytic derivative registered for "
f"{getattr(schedule_fn, '__name__', schedule_fn)!r}"
)
return _DERIV_BY_FN[schedule_fn]
def get_schedule(name: str) -> Callable[[Tensor], Tensor]:
"""Look up a noise schedule by name.
Args:
name: One of ``"linear"``, ``"cosine"``, ``"cosine_sq"``.
Returns:
The schedule function ``alpha(t)``.
Raises:
KeyError: If *name* is not registered.
"""
if name not in _SCHEDULE_MAP:
raise KeyError(
f"Unknown schedule '{name}'. Available: {list(_SCHEDULE_MAP.keys())}"
)
return _SCHEDULE_MAP[name]
|