Blame for lecture_5/dirichlet.py

"""Dirichlet Flow Matching: a token relaxed into a point of the simplex.

A straight line toward a corner shrinks its support as t grows, and past t = 1/k at most k - 1
corners remain possible, so the denoiser target becomes discontinuous in time and in space. If we
replace the line by the Dirichlet family Dir(1 + t e_i), every corner stays possible at every finite
time. The velocity of that path points at the data corner and changes only its speed. To solve for
that speed we match the probability flux across the cut x_i = c.

Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y). On this path the label x_1
still means the data sample while its subscript stops meaning time one, since time runs to infinity
and inference stops at t = 8.

Chapter 5, Sections 5.19.3 and 5.20.1 to 5.20.5. Trained on the synthetic DNA of ``cis6270.data``.
Prints the worked speeds, the continuity check, the posterior average, and the measurements of a
generated batch.

    python lecture_5/dirichlet.py                     # train, then 100 Euler steps to t = 8
    python lecture_5/dirichlet.py --steps 20 --quiet  # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.special import betainc, betaln

from cis6270.data import DNA_SIZE, decode_dna, dna_splits
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters
from cis6270.runner import common_parser, print_report, save_run, set_seed, train

T_MAX = 8.0  # The inference horizon, where the data coordinate has mean (t+1)/(t+K) = 0.75 for DNA.
EDGE = 1e-6  # The speed divides by (1 - x_i) and by x_i, so a coordinate is held off both ends.


# %% 02. The Dirichlet family as a conditional path
#
# With the concentration 1 + t in the data coordinate and 1 everywhere else, every exponent except
# the i-th is zero and the density collapses to a single factor. At t = 0 the path is the uniform
# distribution on the simplex, and for every finite t it is positive wherever x_i is, so no corner is
# ever ruled out.


def dirichlet_density(x, i, t, K=DNA_SIZE):
    """p_t(x | x_1 = e_i) = Gamma(t + K)/Gamma(t + 1) x_i^t."""
    prefactor = math.exp(math.lgamma(t + K) - math.lgamma(t + 1.0))  # Gamma(4) = 6 at t = 0.
    return prefactor * float(x[i]) ** t  # Constant over the simplex at t = 0.


def sample_dirichlet(targets, t, K=DNA_SIZE):
    """Draw x ~ Dir(1 + t e_i) at every position. targets: [B, L], t: [B] -> [B, L, K]."""
    alpha = torch.ones(*targets.shape, K) + t[:, None, None] * F.one_hot(targets, K)  # [B, L, K].
    gamma = torch._standard_gamma(alpha)  # One Gamma draw per letter, with unit rate.
    return gamma / gamma.sum(-1, keepdim=True)  # Once normalized, the Gamma draws are a Dirichlet.


def straight_line_corners(x, t):
    """How many corners the straight-line path still allows at x, the count of x_i >= t."""
    return int((x >= t).sum())  # Proposition 5.5 bounds this by k - 1 once t > 1/k.


# %% 03. The velocity by matching flux across a cut
#
# Cut the simplex at x_i = c. The mass below the cut is then the Beta(t+1, K-1) distribution
# function at c, and its rate of change is the derivative of that function in its first parameter.
# We set the loss rate equal to minus the outward flux and solve for the speed. That leaves one
# scalar per coordinate, and at general K it has no closed form.


def d_betainc_da(x, a, b, eps=1e-6):
    """tilde I_x(a, b), the derivative of the regularized incomplete beta function in a."""
    return (betainc(a + eps, b, x) - betainc(a - eps, b, x)) / (2.0 * eps)  # A central difference.


def speed(x_i, t, K=DNA_SIZE):
    """C(x_i, t) = -tilde I_{x_i}(t+1, K-1) B(t+1, K-1) / ((1 - x_i)^{K-1} x_i^t).

    The data coordinate may be a scalar or a tensor of any shape, since every factor broadcasts,
    and a whole batch of relaxed tokens costs the same two calls to the incomplete beta function
    that one token costs. The central difference subtracts two numbers that agree to about six
    decimals, and at single precision nothing of the difference survives. We therefore compute this
    one quantity in double whatever the ambient dtype, and the caller casts the result back.
    """
    a, b = t + 1.0, float(K - 1)  # The two Beta parameters of the data coordinate.
    beta = math.exp(betaln(a, b))  # The complete beta function, 1/12 at a = 2 and b = 3.
    held = torch.as_tensor(x_i, dtype=torch.double).clamp(EDGE, 1.0 - EDGE)  # Off both endpoints.
    derivative = torch.as_tensor(d_betainc_da(held.numpy(), a, b))  # Negative, so the speed is positive.
    return -derivative * beta / ((1.0 - held) ** (K - 1) * held ** t)  # Positive.


def conditional_velocity(x, i, t):
    """u_t(x | x_1 = e_i) = C(x_i, t)(e_i - x), tangent to the simplex. x: [K] -> [K]."""
    direction = F.one_hot(torch.tensor(i), len(x)).to(x) - x  # Coordinates summing to zero.
    return float(speed(float(x[i]), t, len(x))) * direction  # The speed scales the direction.


# %% 04. Two letters, in closed form
#
# With K = 2 the second Beta parameter is one, the incomplete beta function is x^a, and the speed
# comes out in elementary functions. If we differentiate the density in time and the flux in space,
# the two cancel term by term, the continuity equation at K = 2.


def speed_two_letters(x, t):
    """C(x_1, t) = -x_1 log x_1 / ((t + 1)(1 - x_1)), the K = 2 case of the general formula."""
    return -x * math.log(x) / ((t + 1.0) * (1.0 - x))  # Positive on the whole open interval.


def continuity_terms(x, i, t, h=1e-2):
    """partial_t p and div(p u) in the chart (x_1, ..., x_{K-1}), with sum zero on this path."""
    x = x.to(torch.double)  # Double precision, because the stencil sits on a numerical derivative.
    K = len(x)

    def density(point):
        return dirichlet_density(point, i, t, K)  # The Dirichlet density at a simplex point.

    def chart(point, coordinate):
        return density(point) * float(conditional_velocity(point, i, t)[coordinate])  # p u_j.

    time = density(x) * (_digamma(t + K) - _digamma(t + 1.0) + math.log(float(x[i])))  # partial_t p.
    divergence = 0.0
    for j in range(K - 1):  # The chart drops the last coordinate, fixed by the sum constraint.
        step = torch.zeros(K, dtype=x.dtype)
        step[j], step[K - 1] = 1.0, -1.0  # A direction inside the simplex, keeping the sum at one.
        divergence += (  # A five-point stencil, since the speed itself is a numerical derivative.
            -chart(x + 2 * h * step, j) + 8 * chart(x + h * step, j)
            - 8 * chart(x - h * step, j) + chart(x - 2 * h * step, j)
        ) / (12.0 * h)
    return time, divergence  # Equal and opposite to the precision of the difference.


def _digamma(z):
    """The digamma function, used only to differentiate the Gamma prefactor in time."""
    return float(torch.digamma(torch.tensor(z, dtype=torch.double)))


# %% 05. A worked speed on a DNA base
#
# A coordinate already holding a large share has a larger speed, so the flow concentrates gradually.


def speed_vector(x, t):
    """C(x_k, t) at every letter of one relaxed base. x: [K] -> [K]."""
    return speed(x, t, x.shape[-1]).to(x)  # One vectorized evaluation, cast back to x's dtype.


# %% 06. The posterior average and the loss
#
# The data corner is unknown at generation time, so the four conditional velocities combine with the
# denoiser as weights. Each one is tangent to the simplex and a weighted average of tangent vectors
# is tangent, so the learned field needs no constraint. Cross-entropy on the data token trains the
# denoiser, and the conditional velocity appears only inside the sampler.


def marginal_velocity(x, p_hat, t):
    """v_t(x) = sum_k u_t(x | e_k) p_{1|t}(e_k | x). x: [K], p_hat: [K] -> [K]."""
    columns = torch.stack([conditional_velocity(x, k, t) for k in range(len(x))], dim=1)  # [K, K].
    return columns @ p_hat  # The posterior average of the four conditional velocities.


def batched_velocity(x, p_hat, t):
    """The same average at every position of a batch. x, p_hat: [B, L, K] -> [B, L, K]."""
    speeds = speed(x, t, x.shape[-1]).to(x)  # [B, L, K] one speed per coordinate.
    weighted = speeds * p_hat  # [B, L, K] the coefficient of each conditional velocity.
    return weighted - weighted.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.


def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
    model = SequenceDenoiser(K, length, width=width, depth=depth)
    model.token_embedding = nn.Linear(K, width, bias=False)  # It now takes a point of the simplex.
    return model


# %% 07. Training on DNA, and sampling by Euler steps
#
# A training step draws t from an exponential with rate one, places a point on the conditional path,
# and scores the true letter. To generate we start from the uniform Dirichlet and integrate the
# marginal velocity to t = 8, projecting back onto the simplex after each step.


@torch.no_grad()
def generate(model, batch, length, steps, K=DNA_SIZE):
    """Integrate the marginal velocity from the uniform simplex to t = T_MAX. Returns [B, L]."""
    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) is the uniform source.
    x = gamma / gamma.sum(-1, keepdim=True)  # [B, L, K].
    h = T_MAX / steps  # The time increment, 0.08 over 100 steps.
    for k in range(steps):
        t = k * h
        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
        x = x + h * batched_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
        x = x.clamp(min=1e-8)  # A large step can push a coordinate below zero.
        x = x / x.sum(-1, keepdim=True)  # Back onto the simplex.
    return x.argmax(-1)  # Read the token off the largest coordinate.


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1000, width=64, samples=64, sample_steps=100)
    args = parser.parse_args()
    set_seed(args.seed)

    report = {
        "density_at_t_0": round(dirichlet_density(torch.full((4,), 0.25), 0, 0.0), 4),  # 6.
        "corners_left_at_0.3": straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3),
        "corners_left_at_0.6": straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6),
        "speed_K2_closed": f"{speed_two_letters(0.3, 1.0):.7f}",  # 0.2579942.
        "speed_K2_general": f"{float(speed(0.3, 1.0, K=2)):.7f}",  # The same to six decimals.
    }
    point = torch.tensor([0.5, 0.3, 0.2], dtype=torch.double)  # A base of three letters, at t = 0.8.
    time_term, space_term = continuity_terms(point, 0, 0.8)
    report["continuity_time_term"] = round(time_term, 4)  # 0.6355.
    report["continuity_space_term"] = round(space_term, 4)  # -0.6355.
    report["continuity_residual"] = f"{time_term + space_term:.2e}"  # Their sum, near zero.

    # Section 5.20.4: K = 4, t = 1 and the relaxed base (0.40, 0.30, 0.20, 0.10).
    x = torch.tensor([0.40, 0.30, 0.20, 0.10], dtype=torch.double)
    report["speeds"] = [round(float(c), 3) for c in speed_vector(x, 1.0)]  # 0.232, 0.199, 0.159, 0.104.
    toward_a = conditional_velocity(x, 0, 1.0)  # [4] the velocity toward the A corner.
    report["velocity_toward_A"] = [round(float(v), 3) for v in toward_a]
    report["velocity_sums_to"] = f"{float(toward_a.sum()):.1e}"  # Zero, so the step stays on it.

    # Section 5.20.5: the same base with the denoiser (0.60, 0.25, 0.10, 0.05).
    p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05], dtype=torch.double)
    v = marginal_velocity(x, p_hat, 1.0)
    report["marginal_velocity"] = [round(float(c), 4) for c in v]  # (0.0553, -0.0133, -0.0262, -0.0158).
    report["euler_step_h_0.1"] = [round(float(c), 4) for c in x + 0.1 * v]  # (0.4055, ...).
    report["step_sums_to"] = round(float((x + 0.1 * v).sum()), 9)  # One, still on the simplex.

    # Training on the shared DNA, then generation.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    data = splits["train_tokens"]  # [N, L].
    model = relaxed_denoiser(args.length, args.width)

    def loss_fn(step):
        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
        t = torch.distributions.Exponential(1.0).sample((args.batch_size,))  # t ~ Exp(1).
        xt = sample_dirichlet(x1, t)  # [B, L, K] a point on the conditional path.
        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # Score the true letter.

    if not args.quiet:
        print(f"parameters={count_parameters(model):,} horizon={T_MAX}")
    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
    tokens = generate(model, args.samples, args.length, args.sample_steps)

    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
    print_report("Dirichlet flow matching", report, decode_dna(tokens))
    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))


if __name__ == "__main__":
    main()