src/diffusion_lm/diffusion.py

"""Absorbing-mask forward corruption, objective, and reverse samplers."""

from __future__ import annotations

from dataclasses import dataclass
from collections.abc import Iterator
from typing import Callable, Literal, Protocol

import torch
import torch.nn.functional as F
from torch import Tensor


class Denoiser(Protocol):
    config: object

    def __call__(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        output_positions: Tensor | None = None,
        attn_mask: Tensor | None = None,
    ) -> Tensor: ...


@dataclass(frozen=True)
class CorruptionBatch:
    noisy_tokens: Tensor
    mask: Tensor
    mask_probability: Tensor
    valid_mask: Tensor


@dataclass(frozen=True)
class LossOutput:
    loss: Tensor
    masked_accuracy: Tensor
    masked_tokens: int


@dataclass(frozen=True)
class UnmaskStep:
    """One observable state of the reverse diffusion process."""

    step: int
    total_steps: int
    tokens: Tensor
    masked_remaining: int


def sample_mask_probabilities(
    batch_size: int,
    *,
    device: torch.device | str,
    eps: float = 1e-3,
    low_discrepancy: bool = True,
    generator: torch.Generator | None = None,
) -> Tensor:
    """Sample linear noise levels in ``[eps, 1]``.

    A random cyclic shift of an evenly spaced grid preserves uniform marginals
    while covering the complete noise range in every reasonably sized batch.
    """

    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    if not 0.0 < eps < 1.0:
        raise ValueError("eps must be in (0, 1)")

    if low_discrepancy:
        offset = torch.rand((), device=device, generator=generator)
        unit = (offset + torch.arange(batch_size, device=device) / batch_size) % 1.0
    else:
        unit = torch.rand(batch_size, device=device, generator=generator)
    return eps + (1.0 - eps) * unit


def corrupt_tokens(
    clean_tokens: Tensor,
    mask_token_id: int,
    *,
    valid_mask: Tensor | None = None,
    mask_probability: Tensor | None = None,
    eps: float = 1e-3,
    low_discrepancy: bool = True,
    generator: torch.Generator | None = None,
) -> CorruptionBatch:
    """Apply the absorbing forward process at one random time per sequence."""

    if clean_tokens.ndim != 2:
        raise ValueError("clean_tokens must have shape [batch, sequence]")
    batch_size, _ = clean_tokens.shape
    if valid_mask is None:
        valid_mask = torch.ones_like(clean_tokens, dtype=torch.bool)
    elif valid_mask.shape != clean_tokens.shape:
        raise ValueError("valid_mask must match clean_tokens")
    else:
        valid_mask = valid_mask.bool()

    if mask_probability is None:
        mask_probability = sample_mask_probabilities(
            batch_size,
            device=clean_tokens.device,
            eps=eps,
            low_discrepancy=low_discrepancy,
            generator=generator,
        )
    else:
        mask_probability = torch.as_tensor(
            mask_probability, device=clean_tokens.device, dtype=torch.float32
        )
        if mask_probability.ndim == 0:
            mask_probability = mask_probability.repeat(batch_size)
        if mask_probability.shape != (batch_size,):
            raise ValueError("mask_probability must be scalar or have shape [batch]")
        if bool(((mask_probability <= 0) | (mask_probability > 1)).any()):
            raise ValueError("mask probabilities must be in (0, 1]")

    random_values = torch.rand(clean_tokens.shape, device=clean_tokens.device, generator=generator)
    mask = (random_values < mask_probability[:, None]) & valid_mask
    noisy_tokens = torch.where(mask, mask_token_id, clean_tokens)
    return CorruptionBatch(noisy_tokens, mask, mask_probability, valid_mask)


def diffusion_cross_entropy(
    logits: Tensor,
    clean_tokens: Tensor,
    corruption: CorruptionBatch,
) -> LossOutput:
    """Compute the continuous-time masked-diffusion likelihood bound.

    ``logits`` may contain all positions as ``[B, L, V]`` or only the masked
    positions as ``[N_masked, V]``. The latter is substantially more memory
    efficient for small models with non-trivial vocabularies.
    """

    if clean_tokens.shape != corruption.noisy_tokens.shape:
        raise ValueError("clean_tokens must match the corruption batch")
    targets = clean_tokens[corruption.mask]
    if logits.ndim == 3:
        if logits.shape[:2] != clean_tokens.shape:
            raise ValueError("full logits must have shape [batch, sequence, vocab]")
        selected_logits = logits[corruption.mask]
    elif logits.ndim == 2:
        selected_logits = logits
    else:
        raise ValueError("logits must have shape [B, L, V] or [N_masked, V]")
    if selected_logits.shape[0] != targets.numel():
        raise ValueError("selected logits count does not match the number of masked tokens")

    masked_tokens = int(targets.numel())
    if masked_tokens == 0:
        zero = logits.sum() * 0.0
        return LossOutput(zero, zero.detach(), 0)

    per_token = F.cross_entropy(selected_logits.float(), targets, reduction="none")
    probabilities = corruption.mask_probability[:, None].expand_as(clean_tokens)
    weights = probabilities[corruption.mask].reciprocal()
    normalizer = corruption.valid_mask.sum().clamp_min(1)
    loss = (per_token * weights).sum() / normalizer
    accuracy = (selected_logits.argmax(dim=-1) == targets).float().mean()
    return LossOutput(loss, accuracy, masked_tokens)


def _sample_categorical(
    logits: Tensor,
    temperature: float,
    generator: torch.Generator | None,
) -> tuple[Tensor, Tensor]:
    """Sample with fp64 Gumbel noise and return token ids plus model confidence."""

    if temperature < 0:
        raise ValueError("temperature must be non-negative")
    log_probs = F.log_softmax(logits.float(), dim=-1)
    if temperature == 0:
        tokens = logits.argmax(dim=-1)
    else:
        # MPS has no float64 kernels. Preserve fp64 categorical sampling by
        # moving only the sampling calculation to CPU on Apple Silicon.
        sampling_device = torch.device("cpu") if logits.device.type == "mps" else logits.device
        if logits.device.type == "mps":
            logits64 = logits.float().cpu().double() / temperature
        else:
            logits64 = logits.double() / temperature
        sampling_generator = generator
        if generator is not None and generator.device != sampling_device:
            sampling_generator = None
        uniform = torch.rand(
            logits64.shape,
            device=sampling_device,
            dtype=torch.float64,
            generator=sampling_generator,
        ).clamp_(1e-12, 1.0 - 1e-12)
        gumbel = -torch.log(-torch.log(uniform))
        tokens = (logits64 + gumbel).argmax(dim=-1).to(logits.device)
    confidence = log_probs.gather(-1, tokens[:, None]).squeeze(-1).exp()
    return tokens, confidence


def iterative_unmask_steps(
    model: Denoiser,
    input_ids: Tensor,
    mask_token_id: int,
    *,
    steps: int = 64,
    temperature: float = 1.0,
    strategy: Literal["ancestral", "confidence", "left_to_right"] = "ancestral",
    blocked_token_ids: tuple[int, ...] = (),
    attn_mask: Tensor | None = None,
    generator: torch.Generator | None = None,
    logits_fn: Callable[[Tensor, Tensor], Tensor] | None = None,
) -> Iterator[UnmaskStep]:
    """Yield each state while filling masks and clamping visible prompt tokens.

    ``ancestral`` implements the absorbing reverse transition from mask rate
    ``t`` to ``s``. ``confidence`` reveals an equal-sized highest-confidence
    group on each pass; it is faster-looking and often useful, but is a heuristic.
    ``left_to_right`` reveals equal-sized position-ordered groups, which keeps
    arithmetic left operands visible before their results are committed.

    ``logits_fn(tokens, masked)`` overrides how predictions are obtained, so a caller
    holding a key/value cache can score only the masked window instead of the whole
    sequence. The revealing schedule is unchanged either way.
    """

    if input_ids.ndim != 2:
        raise ValueError("input_ids must have shape [batch, sequence]")
    if steps <= 0:
        raise ValueError("steps must be positive")
    if strategy not in {"ancestral", "confidence", "left_to_right"}:
        raise ValueError("strategy must be ancestral, confidence, or left_to_right")

    tokens = input_ids.clone()
    batch_size, _ = tokens.shape
    yield UnmaskStep(0, steps, tokens.detach(), int(tokens.eq(mask_token_id).sum()))

    for step in range(steps):
        masked = tokens.eq(mask_token_id)
        if not bool(masked.any()):
            break

        # This function is itself a generator, so a decorator would leave the
        # inference context before iteration begins. Scope it around each pass.
        with torch.inference_mode():
            if logits_fn is not None:
                logits = logits_fn(tokens, masked)
            else:
                # Kept as a conditional kwarg so mask-free denoiser doubles stay valid.
                extra = {} if attn_mask is None else {"attn_mask": attn_mask}
                logits = model(tokens, output_positions=masked, **extra)
            if blocked_token_ids:
                logits = logits.clone()
                for token_id in blocked_token_ids:
                    logits[:, token_id] = torch.finfo(logits.dtype).min
            predictions, confidence = _sample_categorical(logits, temperature, generator)

            proposed = tokens.clone()
            proposed[masked] = predictions
            reveal = torch.zeros_like(masked)
            steps_left = steps - step

            if strategy == "ancestral":
                # Linear t grid: P(unmask from t to s | still masked) = 1 - s/t.
                reveal_probability = 1.0 / steps_left
                reveal = (
                    torch.rand(tokens.shape, device=tokens.device, generator=generator)
                    < reveal_probability
                ) & masked
            elif strategy == "left_to_right":
                for row in range(batch_size):
                    masked_positions = masked[row].nonzero(as_tuple=True)[0]
                    remaining = int(masked_positions.numel())
                    count = (remaining + steps_left - 1) // steps_left
                    if count:
                        reveal[row, masked_positions[:count]] = True
            else:
                confidence_grid = torch.full(
                    tokens.shape,
                    -torch.inf,
                    device=tokens.device,
                    dtype=confidence.dtype,
                )
                confidence_grid[masked] = confidence
                for row in range(batch_size):
                    remaining = int(masked[row].sum())
                    count = (remaining + steps_left - 1) // steps_left
                    if count:
                        positions = confidence_grid[row].topk(count).indices
                        reveal[row, positions] = True

            tokens = torch.where(reveal, proposed, tokens)

        yield UnmaskStep(
            step + 1,
            steps,
            tokens.detach(),
            int(tokens.eq(mask_token_id).sum()),
        )

    if bool(tokens.eq(mask_token_id).any()):
        raise RuntimeError(
            "sampler finished with masked positions; this indicates an internal error"
        )


@torch.no_grad()
def iterative_unmask(
    model: Denoiser,
    input_ids: Tensor,
    mask_token_id: int,
    *,
    steps: int = 64,
    temperature: float = 1.0,
    strategy: Literal["ancestral", "confidence"] = "ancestral",
    blocked_token_ids: tuple[int, ...] = (),
    attn_mask: Tensor | None = None,
    generator: torch.Generator | None = None,
    logits_fn: Callable[[Tensor, Tensor], Tensor] | None = None,
) -> Tensor:
    """Return the final state from :func:`iterative_unmask_steps`."""

    final_state: UnmaskStep | None = None
    for state in iterative_unmask_steps(
        model,
        input_ids,
        mask_token_id,
        steps=steps,
        temperature=temperature,
        strategy=strategy,
        blocked_token_ids=blocked_token_ids,
        attn_mask=attn_mask,
        generator=generator,
        logits_fn=logits_fn,
    ):
        final_state = state
    if final_state is None:  # Defensive: the iterator always yields its initial state.
        raise RuntimeError("sampler produced no state")
    return final_state.tokens