File size: 6,184 Bytes
7f316fe | 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 | # Copyright (c) Meta Platforms, Inc.
# All rights reserved.
from typing import Callable, List, Sequence, Tuple, Optional
import torch
import torch.nn.functional as F
from torch import Tensor
class EditFlowsEulerSolver:
"""
Ragged EditFlows CTMC solver with Euler thinning (≤1 jump per step).
At each step h:
- query model for (λ_ins[slots], Q_ins[slots], λ_del[tokens], λ_sub[tokens], Q_sub[tokens])
- total intensity Λ = sum(λ_ins) + sum(λ_del) + sum(λ_sub)
- with prob 1 - exp(-h Λ), take ONE jump:
* sample which event from concatenated rates
* if ins/sub: sample token from the Q at that position/slot
* apply edit to the ragged sequence (insert/delete/replace)
- advance t ← t + h
This is faithful to the CTMC in the small-h regime used by Euler schemes.
"""
def __init__(
self,
model, # callable: (x_list, t_vec) -> ragged heads
scheduler: Callable[[Tensor], object], # path.scheduler(t) providing alpha_t, d_alpha_t, sigma_t if needed
vocab_size: int,
dtype_categorical: torch.dtype = torch.float64,
generator: Optional[torch.Generator] = None,
) -> None:
self.model = model
self.scheduler = scheduler
self.vocab_size = vocab_size
self.dtype_cat = dtype_categorical
self.gen = generator
@torch.no_grad()
def _sample_categorical_logits(self, logits: Tensor) -> Tensor:
# logits: (V,) -> sample id
probs = F.softmax(logits, dim=-1, dtype=self.dtype_cat)
return torch.multinomial(probs, num_samples=1, replacement=True, generator=self.gen).squeeze(0)
@torch.no_grad()
def _maybe_one_jump(
self,
x: Tensor, # (n,)
lam_ins: Tensor, # (n+1,)
logits_ins: Tensor, # (n+1, V)
lam_del: Tensor, # (n,)
lam_sub: Tensor, # (n,)
logits_sub: Tensor, # (n, V)
h: float,
device: torch.device,
) -> Tensor:
"""
Apply at most one edit event to x with step size h.
Returns possibly modified x (1-D LongTensor).
"""
n = x.numel()
# Total intensity
sum_ins = lam_ins.sum() if lam_ins.numel() else torch.tensor(0.0, device=device)
sum_del = lam_del.sum() if lam_del.numel() else torch.tensor(0.0, device=device)
sum_sub = lam_sub.sum() if lam_sub.numel() else torch.tensor(0.0, device=device)
total = (sum_ins + sum_del + sum_sub).clamp(min=0.0)
if total.item() <= 0:
return x # no possible jumps
# Bernoulli: does a jump occur?
p_jump = 1.0 - torch.exp(torch.tensor(-h, device=device) * total)
if torch.rand((), device=device, generator=self.gen) >= p_jump:
return x # no jump this step
# Build event vector: [ins(0..n), del(0..n-1), sub(0..n-1)]
parts = []
if lam_ins.numel(): parts.append(lam_ins)
if lam_del.numel(): parts.append(lam_del)
if lam_sub.numel(): parts.append(lam_sub)
lam_cat = torch.cat(parts, dim=0) # (n+1 + n + n,)
probs = (lam_cat / total).to(self.dtype_cat)
# Sample which event
e_idx = torch.multinomial(probs, num_samples=1, generator=self.gen).item()
# Decode which branch
offset = 0
if e_idx < lam_ins.numel():
# INSERT at slot s
s = e_idx
# sample token from slot distribution
y = self._sample_categorical_logits(logits_ins[s]) # id ∈ [0..V-1]
# perform insertion at slot s: x' = [x[:s], y, x[s:]]
if s == 0:
x_new = torch.cat([y.view(1), x], dim=0)
elif s == n:
x_new = torch.cat([x, y.view(1)], dim=0)
else:
x_new = torch.cat([x[:s], y.view(1), x[s:]], dim=0)
return x_new
e_idx -= lam_ins.numel()
if e_idx < lam_del.numel():
# DELETE token j
j = e_idx
if n == 0:
return x
if j == 0:
return x[1:]
elif j == n - 1:
return x[:-1]
else:
return torch.cat([x[:j], x[j + 1 :]], dim=0)
e_idx -= lam_del.numel()
# SUBSTITUTE at token j
j = e_idx
if n == 0:
return x
y = self._sample_categorical_logits(logits_sub[j])
x_new = x.clone()
x_new[j] = y
return x_new
@torch.no_grad()
def sample(
self,
x_list: Sequence[Tensor], # ragged: list[(n_i,)]
n_steps: int,
t0: float = 0.0,
t1: float = 1.0,
verbose: bool = False,
) -> List[Tensor]:
"""
Evolve all sequences independently with shared time grid t_k.
"""
assert n_steps >= 1
B = len(x_list)
device = x_list[0].device if B > 0 else torch.device("cpu")
x = [xi.clone() for xi in x_list]
t = torch.full((B,), float(t0), device=device)
h = (t1 - t0) / float(n_steps)
for k in range(n_steps):
# Query model once at time t_k (vector of size B)
lam_ins, logits_ins, lam_del, lam_sub, logits_sub = self.model(x, t)
# Apply at most one jump per sample
for i in range(B):
x[i] = self._maybe_one_jump(
x=x[i],
lam_ins=lam_ins[i].squeeze(0) if lam_ins[i].dim() == 2 else lam_ins[i],
logits_ins=logits_ins[i].squeeze(0) if logits_ins[i].dim() == 3 else logits_ins[i],
lam_del=lam_del[i].squeeze(0) if lam_del[i].dim() == 2 else lam_del[i],
lam_sub=lam_sub[i].squeeze(0) if lam_sub[i].dim() == 2 else lam_sub[i],
logits_sub=logits_sub[i].squeeze(0) if logits_sub[i].dim() == 3 else logits_sub[i],
h=h,
device=device,
)
# advance time
t = t + h
return x
|