File size: 7,702 Bytes
278b5e7 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """
Noising schedule for SAD (Soft-Ancestor Skip Diffusion).
We implement a practical discrete-time level-mixture noising process:
- Sample t ~ Uniform[eps, 1-eps]
- Convert t to categorical level weights rho_0(t), ..., rho_L(t)
- For each token independently, sample a corruption level h ~ rho(t)
- Build noisy input using the state representation at level h
Two schedules are provided:
1. `CategoricalLevelSchedule`: rho_l(t) is a piecewise linear function of t.
rho_0(t) = 1-t (stay clean),
rho_l(t) proportional to max(0, t - (l-1)/L) - max(0, t - l/L) for l=1..L
This is a uniform "smear" across levels as t increases.
Easy to implement and works well in practice.
2. `AdjacentLevelSchedule`: heavier mass on adjacent-level transitions.
More principled; TODO: connect to exact CTMC in a future version.
APPROXIMATION NOTE:
The exact CTMC ELBO for a continuous-time multi-level process requires computing
per-token transition rates between non-adjacent states, which is mathematically
involved for learned soft ancestors. The schedules below are practical discrete-time
surrogates. Hooks for a stricter CTMC ELBO are marked with # TODO(CTMC).
"""
from abc import ABC, abstractmethod
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
class LevelNoiseSchedule(nn.Module, ABC):
"""
Base class for level-mixture noise schedules.
Level numbering convention:
level 0 : clean leaf tokens (vocabulary)
level 1 .. L-1 : intermediate prototype levels (from level_sizes[1..L-1])
level L (= num_levels) : MASK state (implicit; uses mask_token_id)
So `num_levels` = len(level_sizes) counts only the real vocabulary levels.
The mask state lives one step *above* the coarsest prototype level and is
NOT a prototype — it is represented by mask_token_id directly.
`total_states` = num_levels + 1 (includes the mask level).
"""
def __init__(self, num_levels: int, vocab_size: int, mask_token_id: int):
"""
Args:
num_levels: number of real vocabulary levels (leaf + intermediate prototypes).
level_sizes has length num_levels.
Mask is the implicit (num_levels)-th level.
vocab_size: |V|, leaf vocabulary size.
mask_token_id: id of the [MASK] token.
"""
super().__init__()
self.num_levels = num_levels # number of real vocab levels (NOT counting mask)
self.total_states = num_levels + 1 # num_levels real levels + 1 mask level
self.vocab_size = vocab_size
self.mask_token_id = mask_token_id
@abstractmethod
def level_weights(self, t: torch.Tensor) -> torch.Tensor:
"""
Return per-level probability weights rho(t).
Args:
t: [B] timesteps in [0, 1]
Returns:
rho: [B, total_states] where total_states = num_levels + 1.
rho[:, 0] = weight for level 0 (clean leaf)
rho[:, 1..L-1] = weight for intermediate prototype levels
rho[:, num_levels] = weight for MASK level
"""
raise NotImplementedError
def sample_levels(self, t: torch.Tensor, seq_len: int) -> torch.Tensor:
"""
Sample independent corruption levels for each (batch, position).
Returns levels in [0, total_states-1].
Level num_levels means MASK.
Args:
t: [B]
seq_len: S
Returns:
levels: [B, S] int64, values in [0, total_states-1]
value == num_levels means MASK
"""
rho = self.level_weights(t) # [B, total_states]
rho_expanded = rho.unsqueeze(1).expand(-1, seq_len, -1) # [B, S, total_states]
levels = torch.multinomial(
rho_expanded.reshape(-1, self.total_states), 1
).reshape(-1, seq_len) # [B, S]
return levels
class CategoricalLevelSchedule(LevelNoiseSchedule):
"""
Simple piecewise-linear level-mixture schedule.
States: 0=leaf, 1..num_levels-1=intermediate prototypes, num_levels=MASK
total_states = num_levels + 1
At t=0: all weight on state 0 (clean leaf).
At t=1: all weight on state num_levels (MASK).
Intermediate states peak in between with triangular tent functions.
"""
def level_weights(self, t: torch.Tensor) -> torch.Tensor:
# t: [B]
B = t.shape[0]
N = self.total_states # num_levels + 1
device = t.device
dtype = t.dtype
# Band boundaries: [0, 1/N, 2/N, ..., 1] (N+1 boundaries, N bands)
boundaries = torch.linspace(0.0, 1.0, N + 1, device=device, dtype=dtype)
rho = torch.zeros(B, N, device=device, dtype=dtype)
# State 0: rho_0 = max(0, 1 - t / b1)
rho[:, 0] = (1.0 - t / boundaries[1].clamp(min=1e-8)).clamp(0.0, 1.0)
# Intermediate states 1..N-2: triangular tent
for l in range(1, N - 1):
lo_l = boundaries[l]
hi_l = boundaries[l + 1]
w_l = (hi_l - lo_l).clamp(min=1e-8)
rise = ((t - lo_l) / w_l).clamp(0.0, 1.0)
fall = ((t - hi_l) / w_l).clamp(max=0.0).abs()
rho[:, l] = (rise - fall).clamp(0.0)
# Last state N-1 (MASK): rho_mask = max(0, (t - b_{N-1}) / w_{N-1})
lo_last = boundaries[N - 1]
w_last = (boundaries[N] - lo_last).clamp(min=1e-8)
rho[:, -1] = ((t - lo_last) / w_last).clamp(0.0, 1.0)
# Normalize
rho = rho / rho.sum(dim=1, keepdim=True).clamp(min=1e-8)
return rho # [B, total_states]
class AdjacentLevelSchedule(LevelNoiseSchedule):
"""
Adjacent-level schedule: at each t, mass splits between two adjacent states.
States: 0=leaf .. num_levels-1=coarsest prototype, num_levels=MASK
total_states = num_levels + 1
mu(t) = t * num_levels maps t in [0,1] to mean state in [0, num_levels].
# TODO(CTMC): Connect to exact continuous-time rates.
"""
def level_weights(self, t: torch.Tensor) -> torch.Tensor:
B = t.shape[0]
N = self.total_states # num_levels + 1
device = t.device
dtype = t.dtype
mu = t * (N - 1) # [B] in [0, N-1]
lo = mu.floor().long().clamp(0, N - 2)
hi = (lo + 1).clamp(0, N - 1)
frac_hi = (mu - lo.float()).clamp(0.0, 1.0)
frac_lo = 1.0 - frac_hi
rho = torch.zeros(B, N, device=device, dtype=dtype)
rho.scatter_(1, lo.unsqueeze(1), frac_lo.unsqueeze(1))
rho.scatter_add_(1, hi.unsqueeze(1), frac_hi.unsqueeze(1))
return rho # [B, total_states]
def sample_t(batch_size: int, eps: float = 1e-4,
low_discrepancy: bool = False,
device=None) -> torch.Tensor:
"""Sample timesteps t ~ Uniform[eps, 1-eps]."""
if low_discrepancy:
t = torch.arange(batch_size, device=device).float() / batch_size
t = (t + torch.rand(1, device=device)).fmod(1.0)
else:
t = torch.rand(batch_size, device=device)
t = (1 - 2 * eps) * t + eps
return t
def get_schedule(schedule_type: str, num_levels: int,
vocab_size: int, mask_token_id: int) -> LevelNoiseSchedule:
if schedule_type == "categorical":
return CategoricalLevelSchedule(num_levels, vocab_size, mask_token_id)
elif schedule_type == "adjacent":
return AdjacentLevelSchedule(num_levels, vocab_size, mask_token_id)
else:
raise ValueError(f"Unknown schedule type: {schedule_type}")
|