Upload 87 files

0ba9d09 verified about 1 month ago

11.1 kB

"""Block diffusion on DNA, with the block size as the dial between diffusion and autoregression.
The chain rule over blocks is exact, and each factor is a distribution over one block of tokens
conditioned on a clean prefix. Masked diffusion models that factor directly. At one block the
bound is the MDLM loss, at blocks of a single position it is the autoregressive negative
log-likelihood with no slack, and every size in between interpolates.
Chapter 4, "Blocks between diffusion and autoregression" and "The block size as a dial between two
models". 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 bound of 2.7567 nats, the one-position identity,
and the sample statistics.
    python lecture_4/block.py                 # train and generate with blocks of four
    python lecture_4/block.py --block 1       # the autoregressive end of the dial
    python lecture_4/block.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The worked bound is a sum of logarithms.

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. One block, with a clean prefix and a masked suffix
#
# The rate matrix acts only on the current block. Everything before it is clean and conditions the
# prediction, and everything after it is hidden from the denoiser.
def block_input(x0: torch.Tensor, block: torch.Tensor, size: int, t: torch.Tensor) -> torch.Tensor:
    """The network input for one block: clean prefix, corrupted block, masked suffix. [B, L]."""
    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L] the column index.
    start = (block * size)[:, None]  # [B, 1] where the current block begins.
    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)  # [B, L] the survival draw.
    corrupted = torch.where(keep, x0, MASK)  # [B, L] masking applied inside the block.
    return torch.where(position < start, x0, torch.where(inside, corrupted, MASK))  # [B, L].


def block_loss(model: torch.nn.Module, x0: torch.Tensor, size: int) -> torch.Tensor:
    """One Monte Carlo sample of the block bound, with one block and one time per sequence."""
    blocks = math.ceil(x0.shape[1] / size)  # B, the number of blocks the sequence splits into.
    block = torch.randint(0, blocks, (x0.shape[0],), device=x0.device)  # [B] which factor to sample.
    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
    z = block_input(x0, block, size, t[:, None])  # [B, L] the denoiser's input.
    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L].
    start = (block * size)[:, None]  # [B, 1].
    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
    logits = model(z, t)  # [B, L, 5] the clean-token prediction at every position.
    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
    scored = ce * ((z == MASK) & inside)  # [B, L] only the masked positions of the current block.
    # We draw one block uniformly to estimate a sum over blocks, so the term carries that count.
    return (blocks * (1.0 / t)[:, None] * scored).sum(-1).mean()  # Nats per sequence.


# %% 03. The block size as a dial
#
# At a block of one position the time integral collapses. The position is masked with probability
# 1 - alpha_t and weighted by w(t) = -alpha_dot/(1 - alpha_t), so the integrand is -alpha_dot times
# the cross-entropy, and the remaining integral is alpha_0 - alpha_1 = 1.
def one_position_integral(grid: int = 2048) -> float:
    """The integral of w(t)(1 - alpha_t) over [0, 1], made exactly one by the collapse."""
    t = (torch.arange(grid).double() + 0.5) / grid  # [grid] midpoints of a uniform partition.
    return float(((1.0 / t) * t).mean())  # w(t)(1 - alpha_t) = 1 at every t on the linear schedule.


def one_position_estimate(probability: float, draws: int, generator) -> float:
    """The same integral by Monte Carlo, the way the training loss evaluates it."""
    t = torch.rand(draws, generator=generator).clamp(min=TIME_FLOOR)  # [N] uniform times.
    masked = torch.rand(draws, generator=generator) < t  # [N] the single position is masked or not.
    return float(((1.0 / t) * masked * -math.log(probability)).mean())  # Should land on -log d.


# %% 04. Generating block by block
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, size: int,
           steps: int = 32) -> torch.Tensor:
    """Fill one block at a time by masked diffusion, with the finished blocks as a clean prefix."""
    device = next(model.parameters()).device
    z = torch.full((batch, length), MASK, device=device)  # [B, L] every position starts masked.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid per block.
    position = torch.arange(length, device=device)[None, :]  # [1, L] the column index.
    for start in range(0, length, size):  # Blocks are completed in order, left to right.
        inside = (position >= start) & (position < start + size)  # [1, L] the current block.
        for t, s in zip(ts[:-1], ts[1:]):
            logits = model(z, t.expand(batch))  # [B, L, 5] the prefix conditions every prediction.
            logits[..., MASK] = -1e9  # The zero-mask rule, at sampling time.
            x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L] candidate tokens.
            unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
            z = torch.where((z == MASK) & unmask & inside, x0, z)  # Reveal inside this block only.
        logits = model(z, ts[-1].expand(batch))  # The last step of the block runs down to s = 0.
        logits[..., MASK] = -1e9
        final = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
        z = torch.where((z == MASK) & inside, final, z)  # The block is clean before the next starts.
    return z  # [B, L] clean DNA.


# %% 05. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--block", type=int, default=4, help="L', the block length")
    # Generation runs the whole reverse grid once per block, so --sample-steps stays at 16 while
    # training rises to 1000 steps and the report measures 64 sequences, inside about a minute.
    parser.set_defaults(steps=1000, sample_steps=16, 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, one draw of the bound on ACGTAC with blocks of two at t = 0.5.
    weight = 1.0 / 0.5  # w(0.5) = 2 on the linear schedule.
    contributions = [  # Block 1 masked twice, block 2 once, block 3 once.
        weight * (-math.log(0.5) - math.log(0.8)),
        weight * -math.log(0.7),
        weight * -math.log(0.9),
    ]
    talk(f"worked block bound  {[round(v, 4) for v in contributions]}  total "
         f"{sum(contributions):.4f} nats on ACGTAC")

    # Section 03, the one-position identity that makes blocks of length one autoregressive.
    generator = torch.Generator().manual_seed(args.seed)  # A fixed stream for the estimate.
    exact = one_position_integral()  # The integral of w(t)(1 - alpha_t), equal to one.
    estimate = one_position_estimate(0.7, 20000, generator)  # The same number by Monte Carlo.
    talk(f"one-position dial   integral {exact:.4f}  Monte Carlo bound {estimate:.4f}  against "
         f"-log 0.7 = {-math.log(0.7):.4f}")

    # Train the denoiser on the shared synthetic DNA at the chosen block size.
    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 bound below.
    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
    blocks = math.ceil(args.length / args.block)  # How many factors the chain rule has.
    talk(f"device={device} parameters={count_parameters(model):,} "
         f"blocks={blocks} of length {args.block}")

    def step_loss(step: int) -> torch.Tensor:
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return block_loss(model, train_tokens[rows], args.block)  # One clean batch into one scalar.

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    with torch.no_grad():
        bound = sum(float(block_loss(model, valid_tokens, args.block)) for _ in range(8)) / 8.0

    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.block, args.sample_steps)
    report = {
        "worked bound on ACGTAC": sum(contributions),
        "first block of the worked bound": contributions[0],
        "integral of w(t)(1 - alpha_t)": exact,
        "one-position Monte Carlo bound": estimate,
        "blocks": blocks,
        "block length": args.block,
        "training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation bound, nats/seq": bound,
        "bound, nats per token": bound / args.length,
        "GC fraction": gc_fraction(samples),
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction": motif_fraction(samples),
        "motif fraction in the data": motif_fraction(train_tokens),
        "distinct samples": unique_fraction(samples),
        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
    }
    print_report(f"block: diffusion in blocks of {args.block}", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


# %% 06. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the block bound, train at one block size, and generate block by block.

	"""Block diffusion on DNA, with the block size as the dial between diffusion and autoregression.

	The chain rule over blocks is exact, and each factor is a distribution over one block of tokens
	conditioned on a clean prefix. Masked diffusion models that factor directly. At one block the
	bound is the MDLM loss, at blocks of a single position it is the autoregressive negative
	log-likelihood with no slack, and every size in between interpolates.

	Chapter 4, "Blocks between diffusion and autoregression" and "The block size as a dial between two
	models". 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 bound of 2.7567 nats, the one-position identity,
	and the sample statistics.

	python lecture_4/block.py # train and generate with blocks of four
	python lecture_4/block.py --block 1 # the autoregressive end of the dial
	python lecture_4/block.py --steps 20 # a few seconds, for a first look
	"""

	# %% 01. Imports and configuration
	import math # The worked bound is a sum of logarithms.

	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. One block, with a clean prefix and a masked suffix
	#
	# The rate matrix acts only on the current block. Everything before it is clean and conditions the
	# prediction, and everything after it is hidden from the denoiser.
	def block_input(x0: torch.Tensor, block: torch.Tensor, size: int, t: torch.Tensor) -> torch.Tensor:
	"""The network input for one block: clean prefix, corrupted block, masked suffix. [B, L]."""
	position = torch.arange(x0.shape[1], device=x0.device)[None, :] # [1, L] the column index.
	start = (block * size)[:, None] # [B, 1] where the current block begins.
	inside = (position >= start) & (position < start + size) # [B, L] the current block.
	keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t) # [B, L] the survival draw.
	corrupted = torch.where(keep, x0, MASK) # [B, L] masking applied inside the block.
	return torch.where(position < start, x0, torch.where(inside, corrupted, MASK)) # [B, L].


	def block_loss(model: torch.nn.Module, x0: torch.Tensor, size: int) -> torch.Tensor:
	"""One Monte Carlo sample of the block bound, with one block and one time per sequence."""
	blocks = math.ceil(x0.shape[1] / size) # B, the number of blocks the sequence splits into.
	block = torch.randint(0, blocks, (x0.shape[0],), device=x0.device) # [B] which factor to sample.
	t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR) # [B] one noise level each.
	z = block_input(x0, block, size, t[:, None]) # [B, L] the denoiser's input.
	position = torch.arange(x0.shape[1], device=x0.device)[None, :] # [1, L].
	start = (block * size)[:, None] # [B, 1].
	inside = (position >= start) & (position < start + size) # [B, L] the current block.
	logits = model(z, t) # [B, L, 5] the clean-token prediction at every position.
	logits[..., MASK] = -1e9 # The zero-mask rule: no posterior mass belongs on m.
	ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none") # [B, L].
	scored = ce * ((z == MASK) & inside) # [B, L] only the masked positions of the current block.
	# We draw one block uniformly to estimate a sum over blocks, so the term carries that count.
	return (blocks * (1.0 / t)[:, None] * scored).sum(-1).mean() # Nats per sequence.


	# %% 03. The block size as a dial
	#
	# At a block of one position the time integral collapses. The position is masked with probability
	# 1 - alpha_t and weighted by w(t) = -alpha_dot/(1 - alpha_t), so the integrand is -alpha_dot times
	# the cross-entropy, and the remaining integral is alpha_0 - alpha_1 = 1.
	def one_position_integral(grid: int = 2048) -> float:
	"""The integral of w(t)(1 - alpha_t) over [0, 1], made exactly one by the collapse."""
	t = (torch.arange(grid).double() + 0.5) / grid # [grid] midpoints of a uniform partition.
	return float(((1.0 / t) * t).mean()) # w(t)(1 - alpha_t) = 1 at every t on the linear schedule.


	def one_position_estimate(probability: float, draws: int, generator) -> float:
	"""The same integral by Monte Carlo, the way the training loss evaluates it."""
	t = torch.rand(draws, generator=generator).clamp(min=TIME_FLOOR) # [N] uniform times.
	masked = torch.rand(draws, generator=generator) < t # [N] the single position is masked or not.
	return float(((1.0 / t) * masked * -math.log(probability)).mean()) # Should land on -log d.


	# %% 04. Generating block by block
	@torch.no_grad() # Generation needs no gradient tracking.
	def sample(model: torch.nn.Module, batch: int, length: int, size: int,
	steps: int = 32) -> torch.Tensor:
	"""Fill one block at a time by masked diffusion, with the finished blocks as a clean prefix."""
	device = next(model.parameters()).device
	z = torch.full((batch, length), MASK, device=device) # [B, L] every position starts masked.
	ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device) # The reverse grid per block.
	position = torch.arange(length, device=device)[None, :] # [1, L] the column index.
	for start in range(0, length, size): # Blocks are completed in order, left to right.
	inside = (position >= start) & (position < start + size) # [1, L] the current block.
	for t, s in zip(ts[:-1], ts[1:]):
	logits = model(z, t.expand(batch)) # [B, L, 5] the prefix conditions every prediction.
	logits[..., MASK] = -1e9 # The zero-mask rule, at sampling time.
	x0 = torch.distributions.Categorical(logits=logits).sample() # [B, L] candidate tokens.
	unmask = torch.rand(z.shape, device=device) < (t - s) / t # [B, L] the reveal draw.
	z = torch.where((z == MASK) & unmask & inside, x0, z) # Reveal inside this block only.
	logits = model(z, ts[-1].expand(batch)) # The last step of the block runs down to s = 0.
	logits[..., MASK] = -1e9
	final = torch.distributions.Categorical(logits=logits).sample() # [B, L].
	z = torch.where((z == MASK) & inside, final, z) # The block is clean before the next starts.
	return z # [B, L] clean DNA.


	# %% 05. The report
	def main() -> None:
	parser = common_parser(__doc__.split("\n")[0])
	parser.add_argument("--block", type=int, default=4, help="L', the block length")
	# Generation runs the whole reverse grid once per block, so --sample-steps stays at 16 while
	# training rises to 1000 steps and the report measures 64 sequences, inside about a minute.
	parser.set_defaults(steps=1000, sample_steps=16, 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, one draw of the bound on ACGTAC with blocks of two at t = 0.5.
	weight = 1.0 / 0.5 # w(0.5) = 2 on the linear schedule.
	contributions = [ # Block 1 masked twice, block 2 once, block 3 once.
	weight * (-math.log(0.5) - math.log(0.8)),
	weight * -math.log(0.7),
	weight * -math.log(0.9),
	]
	talk(f"worked block bound {[round(v, 4) for v in contributions]} total "
	f"{sum(contributions):.4f} nats on ACGTAC")

	# Section 03, the one-position identity that makes blocks of length one autoregressive.
	generator = torch.Generator().manual_seed(args.seed) # A fixed stream for the estimate.
	exact = one_position_integral() # The integral of w(t)(1 - alpha_t), equal to one.
	estimate = one_position_estimate(0.7, 20000, generator) # The same number by Monte Carlo.
	talk(f"one-position dial integral {exact:.4f} Monte Carlo bound {estimate:.4f} against "
	f"-log 0.7 = {-math.log(0.7):.4f}")

	# Train the denoiser on the shared synthetic DNA at the chosen block size.
	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 bound below.
	model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
	blocks = math.ceil(args.length / args.block) # How many factors the chain rule has.
	talk(f"device={device} parameters={count_parameters(model):,} "
	f"blocks={blocks} of length {args.block}")

	def step_loss(step: int) -> torch.Tensor:
	rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
	return block_loss(model, train_tokens[rows], args.block) # One clean batch into one scalar.

	losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
	with torch.no_grad():
	bound = sum(float(block_loss(model, valid_tokens, args.block)) for _ in range(8)) / 8.0

	torch.manual_seed(args.seed + 1) # Fix the reveal draws, so the samples repeat.
	samples = sample(model, args.samples, args.length, args.block, args.sample_steps)
	report = {
	"worked bound on ACGTAC": sum(contributions),
	"first block of the worked bound": contributions[0],
	"integral of w(t)(1 - alpha_t)": exact,
	"one-position Monte Carlo bound": estimate,
	"blocks": blocks,
	"block length": args.block,
	"training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
	"validation bound, nats/seq": bound,
	"bound, nats per token": bound / args.length,
	"GC fraction": gc_fraction(samples),
	"GC fraction in the data": gc_fraction(train_tokens),
	"motif fraction": motif_fraction(samples),
	"motif fraction in the data": motif_fraction(train_tokens),
	"distinct samples": unique_fraction(samples),
	"letter entropy, nats": token_entropy(samples, DNA_SIZE),
	}
	print_report(f"block: diffusion in blocks of {args.block}", report, decode_dna(samples))
	save_run(args.out, model=model, config=vars(args), report=report,
	losses=losses, samples=decode_dna(samples))


	# %% 06. Run the script
	if __name__ == "__main__": # Execute when launched directly.
	main() # Work the block bound, train at one block size, and generate block by block.