Hugging Face's logo

ChatterjeeLab
/
CIS6270

Model card Files
xet
 Community
CIS6270 / lecture_4 /planner.py
pranamanam's picture
pranamanam
Upload 87 files
0ba9d09 verified
Raw History Blame Contribute Delete
14.7 kB
"""Choosing the order of the reveals in masked diffusion on DNA.
Every masked position leaves at the same total rate w(t), so we are free to choose which positions
a step reveals, and the marginal reverse rates stay where they were. A planner spends that freedom
on the positions the denoiser is most sure about, and a corrector returns the least supported
visible tokens to the mask. We can also reweight the loss toward the positions the sampler reaches
first.
Chapter 4, "Choosing the order of the reveals", "One sampling step with a planner" and "Bringing
the planner into training". Time runs from t = 0 clean to t = 1 all masks, the survival probability
is alpha_t = 1 - t, and the rates R_t(x, y) act on probability rows, p_dot = p R. Trained on the
synthetic DNA of cis6270.data. Prints the worked selection, the worked PAPL weights, and the
sample statistics under both samplers.
python lecture_4/planner.py # train with planner weights, sample both ways
python lecture_4/planner.py --remask 0 # confidence order with no corrector
python lecture_4/planner.py --steps 20 # a few seconds, for a first look
"""
# %% 01. Imports and configuration
import torch # Tensors, gradients and the categorical draws.
import torch.nn.functional as F # Cross-entropy at every position.
from cis6270.data import DNA_SIZE, decode_dna, dna_splits # ACGT, K = 4 and the shared sequences.
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters # The denoiser d_theta(. | z, t).
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
MASK = DNA_SIZE # A, C, G, T are 0 to 3, and token 4 is the mask m.
VOCAB = DNA_SIZE + 1 # The masked alphabet V_m has K + 1 = 5 symbols.
TIME_FLOOR = 1e-3 # The reveal rate w(t) = 1/t is unbounded at t = 0.
# %% 02. The planner score and the selection
#
# The score condenses the distribution at a position into one number. The probability of the most
# likely clean token is high where the denoiser's distribution is concentrated and low where it is
# spread out, and the negative entropy of the same distribution is the other natural choice.
def planner_score(probabilities: torch.Tensor, kind: str = "confidence") -> torch.Tensor:
"""One number per position from d_theta(. | z, t). probabilities: [B, L, K] -> [B, L]."""
if kind == "entropy":
return (probabilities * probabilities.clamp(min=1e-20).log()).sum(-1) # Negative entropy.
return probabilities.max(-1).values # The probability of the most likely clean token.
def topk_where(score: torch.Tensor, eligible: torch.Tensor, count: int) -> torch.Tensor:
"""A boolean mask selecting the ``count`` highest-scoring eligible positions in each row."""
picked = torch.zeros_like(eligible) # [B, L] nothing selected yet.
if count <= 0:
return picked # Taking none recovers the plain sampler.
ranked = score.masked_fill(~eligible, float("-inf")) # Ineligible positions never win.
index = ranked.topk(min(count, score.shape[1]), dim=1).indices # [B, count] the best positions.
return picked.scatter_(1, index, True) & eligible # Rows with too few eligible slots take fewer.
# %% 03. One sampling step with a planner
#
# The first four lines are the sampler of the implementation section: one forward pass, the
# zero-mask rule, a softmax, and one categorical draw per position. They are independent of the
# planner, so a trained checkpoint can be sampled this way with no retraining.
def planner_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
n_unmask: int, n_remask: int, kind: str = "confidence") -> torch.Tensor:
"""One reverse step that reveals the best ``n_unmask`` masks and remasks the worst tokens."""
logits = model(z, t) # [B, L, 5] the clean-token prediction at this noise level.
logits[..., MASK] = -1e9 # The zero-mask rule: no posterior mass belongs on m.
p = logits.softmax(-1) # [B, L, 5] the denoiser as probabilities.
x0 = torch.distributions.Categorical(probs=p).sample() # [B, L] one candidate per position.
score = planner_score(p[..., :DNA_SIZE], kind) # [B, L] how sure the denoiser is here.
masked = z == MASK # [B, L] the positions a reveal may act on.
pick = topk_where(score, masked, n_unmask) # [B, L] the best masked positions.
z = torch.where(pick, x0, z) # Unmask the best.
back = topk_where(-score, ~masked, n_remask) # [B, L] the worst visible positions.
return torch.where(back, torch.full_like(z, MASK), z) # Remask the worst, the corrector.
# %% 04. The corrector and its flux balance
#
# C_t = R_t + R_hat_t is in detailed balance with p_t, so running it at a fixed time holds the
# marginal in place. Under masking its forward part returns a visible token to the mask at rate
# beta(t) and its reverse part reveals a masked position at rate w(t).
def corrector_fluxes(t: float) -> tuple[float, float]:
"""alpha_t beta(t) and (1 - alpha_t) w(t), the two fluxes across the visible-masked boundary."""
alpha = 1.0 - t # The linear schedule.
return alpha * (1.0 / alpha), (1.0 - alpha) * (1.0 / t) # Both equal -alpha_dot, or one.
# %% 05. Bringing the planner into training
#
# The MDLM loss weights every masked position equally, right when the order of the reveals
# is uniform. Once a planner fixes the order, the weight lambda_i = M rho_i puts the gradient on
# the positions the sampler will reach first. Over the masked positions that weight averages to
# one.
def planner_weights(score: torch.Tensor, masked: torch.Tensor, strength: float) -> torch.Tensor:
"""lambda_i = 1 + strength (M rho_i - 1), with rho the planner's softmax over masked positions."""
# A large negative number in place of -inf, so a row with nothing masked still has a softmax.
ranked = score.masked_fill(~masked, -1e4) # [B, L] only masked positions compete.
rho = ranked.softmax(dim=1) # [B, L] the planner's distribution over which position to fill.
count = masked.sum(dim=1, keepdim=True).clamp(min=1) # [B, 1] M, the number of masked positions.
return 1.0 + strength * (count * rho - 1.0) # [B, L] a weight of one at strength zero.
def papl_loss(model: torch.nn.Module, x0: torch.Tensor, strength: float) -> torch.Tensor:
"""The MDLM loss with the planner weight inside the sum over masked positions."""
t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR) # [B] one noise level each.
keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None] # [B, L] the survival draw.
zt = torch.where(keep, x0, MASK) # [B, L] the corrupted sequence.
logits = model(zt, t) # [B, L, 5] the clean-token prediction.
logits[..., MASK] = -1e9 # The zero-mask rule.
ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none") # [B, L].
masked = zt == MASK # [B, L] where the loss lives.
score = planner_score(logits.softmax(-1)[..., :DNA_SIZE]).detach() # No gradient through these.
weights = planner_weights(score, masked, strength).clamp(min=0.0) # [B, L] lambda_i.
return ((1.0 / t)[:, None] * weights * ce * masked).sum(-1).mean() # Nats per sequence.
# %% 06. Sampling with and without the planner
#
# The number revealed per step follows the schedule: after a step down to s, the masked count is
# the one the survival probability sets, or s L on the linear schedule.
@torch.no_grad() # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 32,
n_remask: int = 0, kind: str = "confidence", planner: bool = True) -> torch.Tensor:
"""Generate from all masks. With planner false the reveals are drawn uniformly at random."""
device = next(model.parameters()).device
z = torch.full((batch, length), MASK, device=device) # [B, L] the all-mask start.
ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device) # The reverse grid.
for t, s in zip(ts[:-1], ts[1:]):
target = int(round(float(s) * length)) # The masked count the schedule sets at s.
current = int(z.eq(MASK).sum(dim=1).max()) # The largest masked count in the batch.
n_unmask = max(current - target, 0) + n_remask # Enough to reach the target after remasking.
if planner:
z = planner_step(model, z, t.expand(batch), n_unmask, n_remask, kind)
else: # The uniform order of the plain sampler, with the same number of reveals.
z = _random_step(model, z, t.expand(batch), n_unmask)
logits = model(z, ts[-1].expand(batch)) # The last step runs down to s = 0, where r = 1.
logits[..., MASK] = -1e9
return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)
def _random_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
n_unmask: int) -> torch.Tensor:
"""The same step with the masked positions ranked by a uniform draw in place of the score."""
logits = model(z, t) # [B, L, 5].
logits[..., MASK] = -1e9 # The zero-mask rule.
x0 = torch.distributions.Categorical(logits=logits).sample() # [B, L].
noise = torch.rand(z.shape, device=z.device) # [B, L] an order with no preference.
return torch.where(topk_where(noise, z == MASK, n_unmask), x0, z)
# %% 07. The report
def main() -> None:
parser = common_parser(__doc__.split("\n")[0])
parser.add_argument("--remask", type=int, default=1, help="positions the corrector sends back")
parser.add_argument("--planner-strength", type=float, default=1.0,
help="0 gives the MDLM loss, 1 gives the PAPL weights")
parser.add_argument("--score", choices=["confidence", "entropy"], default="confidence")
# Two samplers run on the same weights, so the default trains 1000 steps and measures 64
# sequences in each order, and that keeps the run inside about a minute.
parser.set_defaults(steps=1000, sample_steps=32, samples=64)
args = parser.parse_args()
set_seed(args.seed)
device = pick_device(args.device)
talk = (lambda *a: None) if args.quiet else print # --quiet keeps only the report block.
# Section 02, the worked selection on a six-position sequence.
scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]]) # [1, 6] one score per position.
masked = torch.tensor([[True, False, True, True, False, True]]) # Positions 1, 3, 4 and 6.
revealed = topk_where(scores, masked, 2).nonzero()[:, 1] + 1 # The two best masked positions.
returned = topk_where(-scores, ~masked, 1).nonzero()[:, 1] + 1 # The worst visible position.
talk(f"planner selection reveals {sorted(int(i) for i in revealed)} remasks "
f"{sorted(int(i) for i in returned)}")
# Section 04, the two fluxes across the boundary between visible and masked.
talk("corrector flux " + " ".join(
f"t={t}: beta flux {a:.1f}, w flux {b:.1f}" for t in (0.5, 0.25)
for a, b in [corrector_fluxes(t)]))
# Section 05, the worked PAPL weights on a state with four masked positions.
rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]]) # [1, 4] the planner's distribution.
cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]]) # [1, 4] nats at each masked position.
lam = 4.0 * rho # [1, 4] lambda_i = M rho_i, and it averages to one.
talk(f"PAPL weights {[round(float(v), 1) for v in lam[0]]} sum {float(lam.sum()):.1f} "
f"MDLM inner sum {float(cross_entropies.sum()):.2f} PAPL inner sum "
f"{float((lam * cross_entropies).sum()):.2f}")
# Train the denoiser with the planner weights inside the loss.
splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
train_tokens = splits["train_tokens"].to(device) # [N, L] clean sequences over ACGT.
valid_tokens = splits["valid_tokens"].to(device) # [M, L] held out for the loss below.
model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
talk(f"device={device} parameters={count_parameters(model):,} strength={args.planner_strength}")
def step_loss(step: int) -> torch.Tensor:
rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
return papl_loss(model, train_tokens[rows], args.planner_strength)
losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
with torch.no_grad():
weighted = sum(float(papl_loss(model, valid_tokens, args.planner_strength))
for _ in range(4)) / 4.0
plain = sum(float(papl_loss(model, valid_tokens, 0.0)) for _ in range(4)) / 4.0
# Sample the same weights twice, once in a uniform order and once in the planner's order.
runs = {}
for name, planner, remask in (("uniform order", False, 0),
("planner order", True, args.remask)):
torch.manual_seed(args.seed + 1) # The same draws for both samplers.
runs[name] = sample(model, args.samples, args.length, args.sample_steps,
remask, args.score, planner)
report = {
"revealed positions": sorted(int(i) for i in revealed),
"remasked position": sorted(int(i) for i in returned),
"flux at t = 0.5": corrector_fluxes(0.5)[0],
"flux at t = 0.25": corrector_fluxes(0.25)[1],
"PAPL weight at the planner's pick": float(lam[0, 1]),
"MDLM inner sum": float(cross_entropies.sum()),
"PAPL inner sum": float((lam * cross_entropies).sum()),
"training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
"validation PAPL loss": weighted,
"validation MDLM loss": plain,
"corrector remasks per step": args.remask,
}
for name, drawn in runs.items():
report[f"GC fraction, {name}"] = gc_fraction(drawn)
report[f"motif fraction, {name}"] = motif_fraction(drawn)
report[f"distinct, {name}"] = unique_fraction(drawn)
report["GC fraction in the data"] = gc_fraction(train_tokens)
report["motif fraction in the data"] = motif_fraction(train_tokens)
report["letter entropy, planner order"] = token_entropy(runs["planner order"], DNA_SIZE)
print_report("planner: the order of the reveals", report, decode_dna(runs["planner order"]))
save_run(args.out, model=model, config=vars(args), report=report,
losses=losses, samples=decode_dna(runs["planner order"]))
# %% 08. Run the script
if __name__ == "__main__": # Execute when launched directly.
main() # Work the selection, train with the planner weights, and sample both ways.