| """Lecture 4: the worked numbers of Chapter 4, asserted against the files that print them. |
| |
| Every assertion here is a number the chapter computes by hand. The tests import the lecture files |
| directly, so they check the same functions the scripts run, and none of them trains anything. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib.util |
| import math |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| LECTURE = Path(__file__).resolve().parents[1] / "lecture_4" |
|
|
|
|
| def load(name: str): |
| """Import one lecture file by path, since lecture_4 is a folder of scripts and not a package.""" |
| spec = importlib.util.spec_from_file_location(f"lecture_4_{name}", LECTURE / f"{name}.py") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| ctmc = load("ctmc") |
| mdlm = load("mdlm") |
| uniform = load("uniform") |
| d3pm = load("d3pm") |
| sedd = load("sedd") |
| block = load("block") |
| guidance = load("guidance") |
| planner = load("planner") |
| search = load("search") |
| simplex = load("simplex") |
|
|
|
|
| |
|
|
|
|
| def test_two_rules_and_the_dna_row(): |
| """A jumps to C at rate 2 and to G and T at rate 1, so its diagonal entry is -4.""" |
| assert ctmc.obeys_the_two_rules(ctmc.DNA_RATES) |
| assert float(ctmc.DNA_RATES[0, 0]) == pytest.approx(-4.0) |
| assert float(ctmc.exit_rates(ctmc.DNA_RATES)[0]) == pytest.approx(4.0) |
|
|
|
|
| def test_short_interval_row_and_the_step_bound(): |
| """P_{t,t+0.1}(A, .) = (0.6, 0.2, 0.1, 0.1), and the step must satisfy h <= 0.25.""" |
| row = ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.1) |
| assert [round(float(v), 4) for v in row] == [0.6, 0.2, 0.1, 0.1] |
| assert math.exp(-0.4) == pytest.approx(0.6703, abs=5e-5) |
| assert ctmc.largest_valid_step(ctmc.DNA_RATES) == pytest.approx(0.25) |
| assert float(ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.3)[0]) == pytest.approx(-0.2) |
|
|
|
|
| def test_waiting_time_from_the_chapters_uniform(): |
| """We invert S(s) = exp(-4s) at u = 0.42 and get a wait of 0.2169.""" |
| assert -math.log(0.42) / 4.0 == pytest.approx(0.2169, abs=5e-5) |
| assert math.log(2.0) / 4.0 == pytest.approx(0.1733, abs=5e-5) |
|
|
|
|
| def test_master_equation_on_four_dna_probabilities(): |
| """p_dot = (0.6, 0.2, -0.2, -0.6) and one Euler step gives (0.16, 0.22, 0.28, 0.34).""" |
| derivative = ctmc.master_equation(ctmc.START, ctmc.UNIT_RATES) |
| assert [round(float(v), 4) for v in derivative] == [0.6, 0.2, -0.2, -0.6] |
| arriving, leaving = ctmc.arrivals_and_departures(ctmc.START, ctmc.UNIT_RATES) |
| assert float(arriving[0]) == pytest.approx(0.9) |
| assert float(leaving[0]) == pytest.approx(0.3) |
| stepped = ctmc.START + 0.1 * derivative |
| assert [round(float(v), 4) for v in stepped] == [0.16, 0.22, 0.28, 0.34] |
| assert float(derivative.sum()) == pytest.approx(0.0, abs=1e-12) |
|
|
|
|
| def test_exact_evolution_matches_the_matrix_exponential(): |
| """1/4 + (p_0 - 1/4) exp(-4t) agrees with exp(tR) at t = 0.1 and at t = 1.""" |
| for t, expected in ((0.1, [0.1495, 0.2165, 0.2835, 0.3505]), |
| (1.0, [0.2473, 0.2491, 0.2509, 0.2527])): |
| exact = ctmc.exact_evolution(ctmc.START, t) |
| assert [round(float(v), 4) for v in exact] == expected |
| through_matrix = ctmc.START @ ctmc.transition_matrix(ctmc.UNIT_RATES, t) |
| assert float((exact - through_matrix).abs().max()) == pytest.approx(0.0, abs=1e-12) |
|
|
|
|
| def test_the_matrix_exponential_worked_for_dna(): |
| """At tau = log 2 the uniform kernel is 0.625 on the diagonal and 0.125 elsewhere.""" |
| pi = torch.full((4,), 0.25, dtype=torch.float64) |
| tau = math.log(2.0) |
| kernel = ctmc.transition_matrix(ctmc.r_pi(pi), tau) |
| assert float(kernel[0, 0]) == pytest.approx(0.625, abs=1e-12) |
| assert float(kernel[0, 1]) == pytest.approx(0.125, abs=1e-12) |
| assert float(kernel.sum(dim=1).min()) == pytest.approx(1.0) |
| closed = ctmc.pi_kernel(pi, math.exp(-tau)) |
| assert float((kernel - closed).abs().max()) == pytest.approx(0.0, abs=1e-12) |
| euler = {n: float(ctmc.euler_product(ctmc.r_pi(pi), tau, n)[0, 0]) for n in (1, 2, 10)} |
| assert [round(euler[n], 3) for n in (1, 2, 10)] == [0.480, 0.570, 0.616] |
|
|
|
|
| def test_masking_kernel_is_absorbing(): |
| """At alpha = 1/2 the masking kernel keeps each base with 0.5 and sends 0.5 to the mask.""" |
| mask_pi = torch.zeros(5, dtype=torch.float64) |
| mask_pi[4] = 1.0 |
| kernel = ctmc.pi_kernel(mask_pi, 0.5) |
| assert float(kernel[0, 0]) == pytest.approx(0.5) |
| assert float(kernel[0, 4]) == pytest.approx(0.5) |
| assert float(kernel[4, 4]) == pytest.approx(1.0) |
|
|
|
|
| def test_composition_eigenvalues_and_detailed_balance(): |
| """Intervals compose, R_pi has eigenvalues 0 and -1, and the cycle breaks detailed balance.""" |
| whole = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.7) |
| halves = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.3) @ ctmc.transition_matrix( |
| ctmc.UNIT_RATES, 0.4) |
| assert float((whole - halves).abs().max()) == pytest.approx(0.0, abs=1e-12) |
| pi = torch.full((4,), 0.25, dtype=torch.float64) |
| eigenvalues = sorted(round(float(v.real), 6) for v in torch.linalg.eigvals(ctmc.r_pi(pi))) |
| assert eigenvalues == [-1.0, -1.0, -1.0, 0.0] |
| assert float((pi @ ctmc.r_pi(pi)).abs().max()) == pytest.approx(0.0, abs=1e-12) |
| assert ctmc.detailed_balance_gap(ctmc.r_pi(pi), pi) == pytest.approx(0.0, abs=1e-12) |
| cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64) |
| assert float((cycle_pi @ ctmc.CYCLE).abs().max()) == pytest.approx(0.0, abs=1e-12) |
| assert ctmc.detailed_balance_gap(ctmc.CYCLE, cycle_pi) == pytest.approx(1.0 / 3.0) |
|
|
|
|
| def test_time_reversal_retraces_the_marginals(): |
| """With p_t = (0.5, 0.25, 0.125, 0.125) the reverse rate out of C toward A is 2.""" |
| marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64) |
| backward = ctmc.reverse_rates(ctmc.UNIT_RATES, marginal) |
| assert float(backward[1, 0]) == pytest.approx(2.0) |
| assert float(backward[1, 2]) == pytest.approx(0.5) |
| assert float(ctmc.exit_rates(backward)[1]) == pytest.approx(3.0) |
| net = marginal @ backward + marginal @ ctmc.UNIT_RATES |
| assert float(net.abs().max()) == pytest.approx(0.0, abs=1e-12) |
|
|
|
|
| |
|
|
|
|
| def test_reverse_step_worked(): |
| """At alpha_s = 0.7 and alpha_t = 0.4 the reveal probability is 0.5.""" |
| assert mdlm.reveal_probability(0.7, 0.4) == pytest.approx(0.5) |
| law = mdlm.reverse_step_distribution(0.7, 0.4, [0.6, 0.4, 0.0, 0.0]) |
| assert [round(v, 4) for v in law] == [0.3, 0.2, 0.0, 0.0, 0.5] |
| assert sum(law) == pytest.approx(1.0) |
|
|
|
|
| def test_one_term_of_the_loss_worked(): |
| """w(0.5) = 2 and the masked cross-entropy of 1.04982 give a weighted term of 2.09964.""" |
| assert mdlm.reveal_rate(0.5) == pytest.approx(2.0) |
| masked = -math.log(0.7) - math.log(0.5) |
| assert masked == pytest.approx(1.04982, abs=5e-6) |
| assert 2.0 * masked == pytest.approx(2.09964, abs=5e-6) |
|
|
|
|
| def test_the_rate_error_collapses_to_the_cross_entropy(): |
| """At a masked position the rate errors sum to w(t) times the negative log probability.""" |
| rates = [2.0 * d for d in (0.2, 0.6, 0.1, 0.1)] |
| assert [round(v, 4) for v in rates] == [0.4, 1.2, 0.2, 0.2] |
| total = mdlm.rate_error(2.0, rates[1]) + sum( |
| mdlm.rate_error(0.0, r) for i, r in enumerate(rates) if i != 1) |
| assert total == pytest.approx(2.0 * -math.log(0.6), abs=1e-9) |
| assert total == pytest.approx(1.0217, abs=5e-5) |
| assert 2.0 * -math.log(0.4) == pytest.approx(1.8326, abs=5e-5) |
|
|
|
|
| def test_corruption_and_sampling_shapes(): |
| """The corruption keeps or masks, and the sampler returns clean DNA with no mask left.""" |
| torch.manual_seed(0) |
| x0 = torch.randint(0, 4, (8, 6)) |
| zt = mdlm.corrupt(x0, torch.full((8, 1), 0.5)) |
| assert bool(((zt == x0) | (zt == mdlm.MASK)).all()) |
| kept = mdlm.corrupt(x0, torch.zeros(8, 1)).eq(x0).to(torch.get_default_dtype()) |
| assert float(kept.mean()) == 1.0 |
| assert int(mdlm.corrupt(x0, torch.ones(8, 1)).eq(mdlm.MASK).all()) |
| model = _tiny_denoiser(5, 6) |
| drawn = mdlm.sample(model, 4, 6, steps=4) |
| assert drawn.shape == (4, 6) and int(drawn.max()) < 4 |
|
|
|
|
| def _tiny_denoiser(vocab: int, length: int, num_classes: int = 0): |
| """A very small SequenceDenoiser, so the sampler tests cost nothing.""" |
| from cis6270.nets import SequenceDenoiser |
|
|
| torch.manual_seed(0) |
| return SequenceDenoiser(vocab, length, width=16, depth=1, heads=2, num_classes=num_classes) |
|
|
|
|
| |
|
|
|
|
| def test_the_uniform_kernel_and_the_clean_factor(): |
| """At alpha_t = 1/2 the kernel is 0.625 at the clean base and 0.125 elsewhere.""" |
| assert float(uniform.kernel(0.5)[0, 0]) == pytest.approx(0.625) |
| assert float(uniform.kernel(0.5)[0, 1]) == pytest.approx(0.125) |
| assert uniform.clean_factor(0.5) == pytest.approx(5.0) |
|
|
|
|
| def test_the_worked_posterior_average_is_two(): |
| """The chapter's posterior average of the conditional reverse rates on one base comes out 2.""" |
| prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64) |
| posterior = uniform.posterior_over_clean(0.5, 1, prior) |
| assert [round(float(v), 4) for v in posterior] == [0.375, 0.625, 0.0, 0.0] |
| assert float(uniform.conditional_rates(0.5, current=1, clean=0)[0]) == pytest.approx(5.0) |
| assert float(uniform.conditional_rates(0.5, current=1, clean=1)[0]) == pytest.approx(0.2) |
| assert uniform.marginal_reverse_rate(0.5, 1, 0, prior) == pytest.approx(2.0) |
| assert uniform.marginal_reverse_rate(0.5, 1, 2, prior) == pytest.approx(0.5) |
| marginal = prior @ uniform.kernel(0.5) |
| assert [round(float(v), 4) for v in marginal] == [0.5, 0.25, 0.125, 0.125] |
| assert float(marginal[0] / marginal[1]) == pytest.approx(2.0) |
|
|
|
|
| def test_the_rate_error_on_one_base(): |
| """A denoiser of (0.5, 0.3, 0.1, 0.1) gives rates 2.76 and 1.16 and a rate error of 0.754.""" |
| denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64) |
| current = torch.tensor([[1]]) |
| alpha = torch.tensor([[0.5]], dtype=torch.float64) |
| rates = uniform.model_rates(denoiser, uniform.ratio_table(current, alpha))[0, 0] |
| rates = rates / (uniform.BETA / 4) |
| assert [round(float(v), 4) for v in rates] == [2.76, 1.0, 1.16, 1.16] |
| target = uniform.conditional_rates(0.5, current=1, clean=0) |
| error = float(uniform.rate_error(target[[0, 2, 3]], rates[[0, 2, 3]]).sum()) |
| assert error == pytest.approx(0.754, abs=5e-4) |
| assert float(uniform.rate_error(torch.tensor(5.0), torch.tensor(2.76))) == pytest.approx( |
| 0.731, abs=5e-4) |
|
|
|
|
| def test_the_rate_error_is_minimized_at_the_posterior_average(): |
| """0.375 l(5, B) + 0.625 l(0.2, B) is smallest at B = 2, where it equals 1.430.""" |
| def mixture(b: float) -> float: |
| return float(0.375 * uniform.rate_error(torch.tensor(5.0), torch.tensor(b)) |
| + 0.625 * uniform.rate_error(torch.tensor(0.2), torch.tensor(b))) |
|
|
| assert mixture(2.0) == pytest.approx(1.430, abs=5e-4) |
| assert mixture(5.0) == pytest.approx(2.598, abs=5e-4) |
| assert mixture(0.2) == pytest.approx(4.235, abs=5e-4) |
| assert mixture(2.0) < min(mixture(1.9), mixture(2.1)) |
|
|
|
|
| |
|
|
|
|
| def test_cumulative_transition_matrix(): |
| """With beta_1 = 0.2 and beta_2 = 0.3 the product is 0.67 on the diagonal and 0.11 off it.""" |
| product = d3pm.one_step(0.2) @ d3pm.one_step(0.3) |
| assert float(product[0, 0]) == pytest.approx(0.67, abs=5e-3) |
| assert float(product[0, 1]) == pytest.approx(0.11, abs=5e-3) |
| assert float((product - d3pm.cumulative(0.56)).abs().max()) == pytest.approx(0.0, abs=1e-6) |
| assert (1.0 - 1.0 / 1000) ** 1000 == pytest.approx(0.3677, abs=5e-5) |
|
|
|
|
| def test_reversing_one_finite_step(): |
| """We observe C under a clean A and get (0.5795, 0.3523, 0.0341, 0.0341) at the earlier step.""" |
| first, second = d3pm.one_step(0.2), d3pm.one_step(0.3) |
| clean_a = torch.tensor([[[1.0, 0.0, 0.0, 0.0]]]) |
| posterior = d3pm.exact_posterior(torch.tensor([[1]]), clean_a, second[None], first[None])[0, 0] |
| assert [round(float(v), 4) for v in posterior] == [0.5795, 0.3523, 0.0341, 0.0341] |
|
|
|
|
| def test_rate_matrices_chosen_for_the_data(): |
| """The Gaussian row is 0.4088 to a neighbour, and exp(0.3 R) has first row (0.6747, ...).""" |
| row = d3pm.gaussian_row() |
| assert float(row[1]) == pytest.approx(0.4088, abs=5e-5) |
| assert float(row[0]) == pytest.approx(0.0912, abs=5e-5) |
| adjacency = torch.tensor([[0.0, 1.0, 0.5, 0.0], [1.0, 0.0, 1.0, 0.0], |
| [0.5, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]) |
| R = d3pm.graph_rates(adjacency) |
| assert [round(float(v), 4) for v in R[0]] == [-1.5, 1.0, 0.5, 0.0] |
| kernel = torch.matrix_exp(0.3 * R) |
| assert [round(float(v), 4) for v in kernel[0]] == [0.6747, 0.1964, 0.1122, 0.0167] |
|
|
|
|
| |
|
|
|
|
| def test_score_entropy_at_three_predictions(): |
| """With the target ratio 2, a prediction of 1.8 costs 0.0107 nats and 1.0 costs 0.3863.""" |
| target = torch.tensor(2.0) |
| assert float(sedd.score_entropy(target, torch.tensor(1.8))) == pytest.approx(0.0107, abs=5e-5) |
| assert float(sedd.score_entropy(target, torch.tensor(1.0))) == pytest.approx(0.3863, abs=5e-5) |
| assert float(sedd.score_entropy(target, torch.tensor(3.0))) == pytest.approx(0.1891, abs=5e-5) |
| assert float(sedd.score_entropy(target, target)) == pytest.approx(0.0, abs=1e-6) |
|
|
|
|
| def test_the_target_ratio_from_the_denoiser_route(): |
| """0.375(5) + 0.625(0.2) is 2, and the kernel ratio reproduces it at a clean A.""" |
| assert 0.375 * 5.0 + 0.625 * 0.2 == pytest.approx(2.0) |
| x0 = torch.tensor([[0]]) |
| z = torch.tensor([[1]]) |
| ratio = sedd.kernel_ratio(x0, z, torch.tensor([[0.5]])) |
| assert float(ratio[0, 0, 0]) == pytest.approx(5.0) |
| assert float(ratio[0, 0, 2]) == pytest.approx(1.0) |
|
|
|
|
| |
|
|
|
|
| def test_the_worked_block_bound(): |
| """ACGTAC in blocks of two at t = 0.5 gives 1.8326, 0.7134 and 0.2107, summing to 2.7567.""" |
| weight = 2.0 |
| contributions = [weight * (-math.log(0.5) - math.log(0.8)), |
| weight * -math.log(0.7), weight * -math.log(0.9)] |
| assert [round(v, 4) for v in contributions] == [1.8326, 0.7133, 0.2107] |
| assert sum(contributions) == pytest.approx(2.7567, abs=5e-5) |
|
|
|
|
| def test_the_one_position_identity(): |
| """At blocks of one position the time integral of w(t)(1 - alpha_t) is exactly one.""" |
| assert block.one_position_integral() == pytest.approx(1.0, abs=1e-9) |
| generator = torch.Generator().manual_seed(0) |
| estimate = block.one_position_estimate(0.7, 40000, generator) |
| assert estimate == pytest.approx(-math.log(0.7), abs=0.05) |
|
|
|
|
| def test_block_input_keeps_the_prefix_and_hides_the_future(): |
| """The prefix stays clean, the current block is corrupted, and the suffix is all mask.""" |
| x0 = torch.arange(8).remainder(4)[None, :].repeat(2, 1) |
| z = block.block_input(x0, torch.tensor([1, 1]), 2, torch.full((2, 1), 1.0)) |
| assert bool((z[:, :2] == x0[:, :2]).all()) |
| assert bool((z[:, 2:4] == block.MASK).all()) |
| assert bool((z[:, 4:] == block.MASK).all()) |
|
|
|
|
| |
|
|
|
|
| def test_h_transform_keeps_a_rate_matrix(): |
| """When we reweight by h(y)/h(x) the row sums stay at zero, and a constant h changes nothing.""" |
| R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0)) |
| transformed = guidance.h_transform(R, torch.tensor([2.0, 1.0, 0.5, 0.5])) |
| assert float(transformed.sum(dim=1).abs().max()) == pytest.approx(0.0, abs=1e-6) |
| constant = guidance.h_transform(R, torch.full((4,), 3.0)) |
| assert float((constant - R).abs().max()) == pytest.approx(0.0, abs=1e-6) |
|
|
|
|
| def test_the_reverse_process_as_an_h_transform(): |
| """h_t = p_t/pi is (2, 1, 0.5, 0.5), and the rates out of C come back as 2, 0.5 and 0.5.""" |
| pi = torch.full((4,), 0.25) |
| values = guidance.value_from_marginal(torch.tensor([0.5, 0.25, 0.125, 0.125]), pi) |
| assert [round(float(v), 4) for v in values] == [2.0, 1.0, 0.5, 0.5] |
| R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0)) |
| reverse = guidance.h_transform(R, values) |
| assert float(reverse[1, 0]) == pytest.approx(2.0, abs=1e-5) |
| assert float(reverse[1, 2]) == pytest.approx(0.5, abs=1e-5) |
| clean_a = guidance.value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi) |
| assert float(clean_a[0] / clean_a[1]) == pytest.approx(5.0) |
| assert float(clean_a[1] / clean_a[0]) == pytest.approx(0.2) |
|
|
|
|
| def test_the_worked_guided_jump(): |
| """The classifier turns the rates 2, 0.5, 0.5 into 1, 1.5, 0.5, and gamma = 2 sharpens them.""" |
| unguided = torch.tensor([2.0, 0.5, 0.5]) |
| ratio = torch.tensor([0.1, 0.6, 0.2]) / 0.2 |
| guided = unguided * ratio |
| assert [round(float(v), 4) for v in guided] == [1.0, 1.5, 0.5] |
| assert float(guided.sum()) == pytest.approx(3.0) |
| shares = guided / guided.sum() |
| assert [float(v) for v in shares] == pytest.approx([1 / 3, 1 / 2, 1 / 6], abs=1e-6) |
| sharper = unguided * ratio ** 2 |
| assert [round(float(v), 4) for v in sharper] == [0.5, 4.5, 0.5] |
| assert float(sharper.sum()) == pytest.approx(5.5) |
| assert float((sharper / sharper.sum())[1]) == pytest.approx(0.8182, abs=5e-5) |
| assert 1.0 / float(sharper.sum()) == pytest.approx(0.1818, abs=5e-5) |
|
|
|
|
| def test_exact_and_gradient_ratios_agree_at_first_order(): |
| """Both estimates of p(c | z^{i <- a})/p(c | z) have the right shape and a positive value.""" |
| torch.manual_seed(0) |
| classifier = guidance.NoisyClassifier(6, width=16) |
| z = torch.tensor([[0, guidance.MASK, 2, guidance.MASK, 1, 3]]) |
| t = torch.tensor([0.5]) |
| with torch.no_grad(): |
| exact = guidance.exact_ratio(classifier, z, t) |
| approximate = guidance.gradient_ratio(classifier, z, t) |
| assert exact.shape == approximate.shape == (1, 6, 4) |
| assert bool((exact > 0).all()) and bool((approximate > 0).all()) |
| |
| assert float(exact[0, 0, 0]) == pytest.approx(1.0, abs=1e-5) |
| assert float(approximate[0, 0, 0]) == pytest.approx(1.0, abs=1e-5) |
|
|
|
|
| |
|
|
|
|
| def test_the_worked_planner_selection(): |
| """Scores (0.42, 0.91, 0.55, 0.78) at the masked positions reveal 3 and 6, and remask 5.""" |
| scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]]) |
| masked = torch.tensor([[True, False, True, True, False, True]]) |
| revealed = sorted(int(i) + 1 for i in planner.topk_where(scores, masked, 2).nonzero()[:, 1]) |
| returned = sorted(int(i) + 1 for i in planner.topk_where(-scores, ~masked, 1).nonzero()[:, 1]) |
| assert revealed == [3, 6] |
| assert returned == [5] |
| assert int(planner.topk_where(scores, masked, 0).sum()) == 0 |
|
|
|
|
| def test_the_corrector_fluxes_match(): |
| """alpha_t beta(t) and (1 - alpha_t) w(t) both equal -alpha_dot, or one here.""" |
| for t in (0.5, 0.25, 0.75): |
| forward, reverse = planner.corrector_fluxes(t) |
| assert forward == pytest.approx(1.0) |
| assert reverse == pytest.approx(1.0) |
|
|
|
|
| def test_the_worked_papl_weights(): |
| """A planner distribution of (0.10, 0.50, 0.15, 0.25) turns 2.9 nats into 2.12.""" |
| rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]]) |
| cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]]) |
| lam = 4.0 * rho |
| assert [round(float(v), 4) for v in lam[0]] == [0.4, 2.0, 0.6, 1.0] |
| assert float(lam.sum()) == pytest.approx(4.0) |
| assert float(cross_entropies.sum()) == pytest.approx(2.9) |
| assert float((lam * cross_entropies).sum()) == pytest.approx(2.12) |
|
|
|
|
| def test_planner_weights_average_to_one_and_reduce_to_mdlm(): |
| """lambda_i averages to one over the masked positions, and strength zero gives every weight one.""" |
| score = torch.tensor([[0.9, 0.1, 0.5, 0.2]]) |
| masked = torch.tensor([[True, True, False, True]]) |
| weights = planner.planner_weights(score, masked, 1.0) |
| assert float((weights * masked).sum()) == pytest.approx(float(masked.sum()), abs=1e-5) |
| plain = planner.planner_weights(score, masked, 0.0) |
| assert float((plain - 1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) |
|
|
|
|
| |
|
|
|
|
| def test_the_worked_dirichlet_parameters(): |
| """At alpha_t = 1/2 the parameters are (5, 1, 1, 1) and at alpha_s = 0.8 they are (17, 1, 1, 1).""" |
| assert simplex.concentration(0.5) == pytest.approx(8.0) |
| assert simplex.concentration(0.8) == pytest.approx(20.0) |
| mean_t = simplex.belief_mean(0.5, 0) |
| assert [round(float(v), 4) for v in mean_t] == [0.625, 0.125, 0.125, 0.125] |
| assert [round(float(v), 4) for v in simplex.belief_mean(0.8, 0)] == [0.85, 0.05, 0.05, 0.05] |
| at_t = simplex.dirichlet_parameters(0.5, 0) |
| at_s = simplex.dirichlet_parameters(0.8, 0) |
| assert [round(float(v), 4) for v in at_t] == [5.0, 1.0, 1.0, 1.0] |
| assert [round(float(v), 4) for v in at_s] == [17.0, 1.0, 1.0, 1.0] |
| total = at_t.sum() |
| deviation = (at_t * (total - at_t) / (total ** 2 * (total + 1))).sqrt() |
| assert float(deviation[0]) == pytest.approx(0.161, abs=5e-4) |
|
|
|
|
| def test_the_increment_sits_on_the_clean_coordinate(): |
| """The increment is (12, 0, 0, 0), so the innovation is the one-hot vector of the clean base.""" |
| step = simplex.increment(0.5, 0.8, 0) |
| assert [round(float(v), 4) for v in step] == [12.0, 0.0, 0.0, 0.0] |
| first, second = simplex.thinning_law(0.5, 0.8) |
| assert (round(first, 4), round(second, 4)) == (8.0, 12.0) |
| weight = first / (first + second) |
| assert weight == pytest.approx(0.4) |
| mixed = weight * simplex.belief_mean(0.5, 0) + (1.0 - weight) * torch.eye(4)[0] |
| assert [round(float(v), 4) for v in mixed] == [0.85, 0.05, 0.05, 0.05] |
|
|
|
|
| def test_the_thinning_representation_reproduces_the_dirichlet(): |
| """W V + (1 - W) onehot(A) has the mean and the spread of Dir(17, 1, 1, 1).""" |
| torch.manual_seed(0) |
| at_t, at_s = simplex.dirichlet_parameters(0.5, 0), simplex.dirichlet_parameters(0.8, 0) |
| first, second = simplex.thinning_law(0.5, 0.8) |
| draws = 200_000 |
| thinned = torch.distributions.Dirichlet(at_t).sample((draws,)) |
| weight = torch.distributions.Beta(torch.tensor(first), torch.tensor(second)).sample((draws,)) |
| drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(4)[0] |
| total = at_s.sum() |
| exact = (at_s * (total - at_s) / (total ** 2 * (total + 1))).sqrt() |
| assert [round(float(v), 4) for v in exact] == [0.0779, 0.0476, 0.0476, 0.0476] |
| assert float(drawn.mean(0)[0]) == pytest.approx(0.85, abs=2e-3) |
| assert float(drawn.std(0)[0]) == pytest.approx(float(exact[0]), abs=2e-3) |
|
|
|
|
| def test_the_parameter_field_and_the_reverse_step(): |
| """The batched parameters match the per-position ones, and one step stays on the simplex.""" |
| x0 = torch.tensor([[0, 1], [2, 3]]) |
| field = simplex.parameter_field(x0, torch.full((2, 1), 0.5)) |
| assert [round(float(v), 4) for v in field[0, 0]] == [5.0, 1.0, 1.0, 1.0] |
| assert [round(float(v), 4) for v in field[1, 1]] == [1.0, 1.0, 1.0, 5.0] |
| torch.manual_seed(0) |
| belief = simplex.corrupt(x0, torch.full((2, 1), 0.5)) |
| assert belief.shape == (2, 2, 4) |
| assert float(belief.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) |
| moved = simplex.reverse_step(belief, x0, 0.5, 0.8) |
| assert float(moved.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) |
| assert bool((moved >= 0).all()) |
|
|
|
|
| def test_a_one_hot_belief_recovers_the_token_denoiser(): |
| """The architecture changes only at the input, so one-hot beliefs give identical logits.""" |
| torch.manual_seed(0) |
| model = simplex.BeliefDenoiser(6, width=16) |
| tokens = torch.tensor([[0, 1, 2, 3, 0, 1]]) |
| t = torch.tensor([0.4]) |
| one_hot = torch.nn.functional.one_hot(tokens, 4).to(torch.get_default_dtype()) |
| with torch.no_grad(): |
| through_belief = model(one_hot, t) |
| through_tokens = model.inner(tokens, t) |
| assert float((through_belief - through_tokens).abs().max()) == pytest.approx(0.0, abs=1e-6) |
| drawn = simplex.sample(model, 3, 6, steps=4) |
| assert drawn.shape == (3, 6) and int(drawn.max()) < 4 |
|
|
|
|
| |
|
|
|
|
| def test_the_worked_value_estimate(): |
| """Rollouts averaging 0.35 and 0.75 reweight a rate of 1.2 to 2.571.""" |
| parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean()) |
| child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean()) |
| assert parent == pytest.approx(0.35) |
| assert child == pytest.approx(0.75) |
| assert child / parent == pytest.approx(2.143, abs=5e-4) |
| assert 1.2 * child / parent == pytest.approx(2.571, abs=5e-4) |
| assert 0.2 / parent == pytest.approx(0.571, abs=5e-4) |
|
|
|
|
| def test_the_worked_pareto_archive(): |
| """P3 at (0.4, 0.2) is dominated by P1 at (0.8, 0.3), and P1 and P2 both stay.""" |
| archive = search.ParetoArchive() |
| for name, scores in (((1,), (0.8, 0.3)), ((2,), (0.5, 0.7)), ((3,), (0.4, 0.2))): |
| archive.add(name, scores) |
| assert sorted(q for q, _ in archive.entries) == [(1,), (2,)] |
| archive.add((4,), (0.9, 0.9)) |
| assert [q for q, _ in archive.entries] == [(4,)] |
|
|
|
|
| def test_the_two_objectives(): |
| """A GC-free sequence scores zero on the first axis, and the motif scores one on the second.""" |
| tokens = torch.tensor([[0, 1, 2, 3, 0, 0], [0, 0, 3, 3, 0, 3]]) |
| scores = search.objectives(tokens) |
| assert float(scores[0, 0]) == pytest.approx(2 / 6) |
| assert float(scores[0, 1]) == pytest.approx(1.0) |
| assert float(scores[1, 0]) == pytest.approx(0.0) |
| assert float(scores[1, 1]) < 1.0 |
|
|
