Download lecture_5/multi_objective.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 14.1 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_5/multi_objective.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_5/multi_objective.py
-
curl -L -o multi_objective.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_5/multi_objective.py
14.1 kB
| """Several objectives at once: MOG-DFM's cone and pCoMole's rollout. | |
| A peptide is scored on binding, solubility and half-life together, and no dataset labels one | |
| sequence for all of them, so the single value function of Discrete Guidance is out of reach. | |
| MOG-DFM drops the ratio form and builds a tilt out of a rank inside each objective and a direction | |
| against a weight vector, with a cone that removes the candidates pointing the wrong way. pCoMole | |
| keeps the ratio and replaces the property likelihood by an exponentiated worst-objective utility, | |
| and its value at a candidate is the average terminal reward over rollouts of the base flow. | |
| 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). | |
| Chapter 5, Sections 5.30 to 5.32. The single-objective tilt is in ``lecture_5/guidance.py``. The | |
| rollout runs on the synthetic DNA of ``cis6270.data``, scored by GC fraction against AT fraction. | |
| Prints the Pareto front, the chapter's MOG-DFM step with its cone, the chapter's pCoMole rollout, | |
| and a rollout on DNA whose value reverses the base rates. | |
| python lecture_5/multi_objective.py # 600 rollouts per candidate on DNA | |
| python lecture_5/multi_objective.py --steps 4000 # a tighter estimate of the second value | |
| python lecture_5/multi_objective.py --steps 20 --quiet # twenty rollouts, for a first look | |
| """ | |
| # %% 01. Imports and configuration | |
| import math | |
| import torch | |
| from cis6270.data import DNA, DNA_SIZE, MOTIF, decode_dna, dna_splits | |
| from cis6270.runner import common_parser, print_report, save_run, set_seed | |
| MASK = DNA_SIZE # Index 4 is the mask symbol m, still held by a candidate off its own position. | |
| OMEGA = torch.tensor([0.5, 0.5], dtype=torch.double) # Equal weight on the two objectives of the DNA rollout. | |
| REFERENCE = torch.tensor([0.4, 0.3], dtype=torch.double) # The levels r_k the objectives sit above. | |
| BETA = 2.0 # The inverse temperature of the terminal reward. | |
| # %% 02. Several objectives at once | |
| # | |
| # With several objectives no single sequence is best. A candidate is on the Pareto front when no | |
| # other candidate matches it everywhere and beats it somewhere, and a weight vector picks one point | |
| # of that front by collapsing the scores to a single weighted number. | |
| def dominated(scores, row): | |
| """True when some other candidate is at least as good everywhere and better somewhere.""" | |
| others = torch.cat([scores[:row], scores[row + 1:]]) # Every candidate except this one. | |
| weak = (others >= scores[row]).all(dim=1) # At least as good on every objective. | |
| strict = (others > scores[row]).any(dim=1) # Strictly better on at least one. | |
| return bool((weak & strict).any()) | |
| def pareto_front(scores): | |
| """The indices of the candidates no other candidate dominates. scores: [M, N] -> list.""" | |
| return [row for row in range(len(scores)) if not dominated(scores, row)] | |
| # %% 03. MOG-DFM: rank, direction and a cone | |
| # | |
| # Without a joint likelihood we have no ratio to compute, so we score every candidate letter twice. | |
| # A rank inside each objective removes the units, the inner product with the weight vector measures | |
| # direction, and the two z-scored terms exponentiate into a tilt on the pretrained rate. | |
| def rank_scores(changes): | |
| """I_n(y^i, x), the rank of each candidate's change within each objective, over the alphabet size.""" | |
| order = changes.argsort(dim=0).argsort(dim=0) + 1 # Rank one for the smallest change. | |
| return order.to(changes.dtype) / changes.shape[0] # [K, N] in {1/K, ..., 1}. | |
| def z_score(values): | |
| """Norm[.], the z-score across the candidate letters at one position.""" | |
| spread = values.std(unbiased=False).clamp(min=1e-12) # The population deviation, held off zero. | |
| return (values - values.mean()) / spread # Zero everywhere when no candidate stands out. | |
| def combined_score(changes, importance, weights, lam=1.0): | |
| """Delta S = Norm[mean_n i_n I_n] + lambda Norm[Delta s . omega]. changes: [K, N] -> [K].""" | |
| ranks = (rank_scores(changes) * importance).mean(dim=1) # Importance-weighted average of ranks. | |
| direction = changes @ weights # [K] the directional score of each candidate. | |
| return z_score(ranks) + lam * z_score(direction), ranks, direction | |
| # %% 04. Accepting a move only inside a cone | |
| # | |
| # A z-scored score puts about half the candidates above zero, so the tilt alone promotes the least | |
| # bad letter at a position where every letter moves away from the requested trade-off. The angle | |
| # between a candidate's score change and omega removes those, and a controller on the logarithm of | |
| # the angle holds the surviving fraction near one minus the target rejection rate. | |
| def cone_angles(changes, weights): | |
| """alpha^i = arccos(Delta s . omega / (||Delta s|| ||omega||)), in degrees.""" | |
| cosine = (changes @ weights) / (changes.norm(dim=1) * weights.norm()) # [K]. | |
| return torch.rad2deg(torch.arccos(cosine.clamp(-1.0, 1.0))), cosine | |
| def update_cone(phi, rejected, target=0.2, gain=0.5, low=30.0, high=90.0): | |
| """Phi_{t+h} = clip(Phi_t exp(eta (r - tau)), Phi_min, Phi_max), proportional control on log Phi.""" | |
| return min(max(phi * math.exp(gain * (rejected - target)), low), high) # Wider when it rejects more. | |
| # %% 05. pCoMole: a reward in place of a likelihood | |
| # | |
| # If we score a sequence by its weakest objective, the sampler cannot trade one property away, and | |
| # a small multiple of the sum breaks ties. The exponentiated utility times a feasibility indicator | |
| # is the terminal reward, and the value of a candidate is the average reward over rollouts of the | |
| # base flow started from it. | |
| def utility(scores, weights, reference, rho=0.1): | |
| """U(x) = min_k omega_k (f_k(x) - r_k) + rho sum_k omega_k (f_k(x) - r_k).""" | |
| gaps = weights * (scores - reference) # [N] the weighted gap on each objective. | |
| return float(gaps.min() + rho * gaps.sum()) # The minimum, with the tie-breaking correction. | |
| def terminal_reward(scores, weights, reference, beta, feasible, rho=0.1): | |
| """G(x) = exp(beta U(x)) 1[x in F], zero for an infeasible sequence.""" | |
| return math.exp(beta * utility(scores, weights, reference, rho)) * float(feasible) | |
| def rollout_value(rollouts, feasible, weights, reference, beta, rho=0.1): | |
| """hat h_t(y_c), the average terminal reward over R continuations of the base flow.""" | |
| rewards = [terminal_reward(s, weights, reference, beta, f, rho) for s, f in zip(rollouts, feasible)] | |
| return sum(rewards) / len(rewards), rewards # A candidate with no feasible ending scores zero. | |
| # %% 06. A rollout on the synthetic DNA | |
| # | |
| # The base flow of a rollout is the mask-source sampler of this chapter with the letter frequency of | |
| # the data as its denoiser, so a continuation costs no network call. The GC fraction and the AT | |
| # fraction of the finished sequence pull against each other. The feasible set is the sequences that | |
| # carry the planted motif. | |
| def continuations(partial, frequency, rollouts, generator=None): | |
| """Fill every masked position of a candidate from the base flow. partial: [L] -> [R, L].""" | |
| grid = partial[None, :].repeat(rollouts, 1) # [R, L] one copy of the candidate per rollout. | |
| masked = grid == MASK # [R, L] the positions the flow still owes a letter. | |
| draws = torch.multinomial(frequency, int(masked.sum()), replacement=True, generator=generator) | |
| return grid.masked_scatter(masked, draws) # [R, L] the terminal sequences of the rollouts. | |
| def dna_objectives(sequences): | |
| """The GC fraction and the AT fraction of each sequence. The two sum to one. [R, L] -> [R, 2].""" | |
| gc = ((sequences == 1) | (sequences == 2)).to(torch.get_default_dtype()).mean(dim=1) # C and G. | |
| return torch.stack([gc, 1.0 - gc], dim=1) # [R, 2] raising one lowers the other. | |
| def carries_motif(sequences): | |
| """One per sequence that holds the planted motif, the hard constraint here. [R] ints.""" | |
| return [int(MOTIF in sequence) for sequence in decode_dna(sequences)] | |
| def dna_rollout_value(partial, frequency, rollouts, weights, reference, beta, generator=None): | |
| """hat h_t at one candidate, from ``rollouts`` continuations of the base flow. Returns a float.""" | |
| finished = continuations(partial, frequency, rollouts, generator) # [R, L]. | |
| scores = dna_objectives(finished) # [R, 2] the two objectives at the end of each rollout. | |
| value, _ = rollout_value(list(scores), carries_motif(finished), weights, reference, beta) | |
| return value # Zero when no continuation of this candidate is feasible. | |
| # %% 07. The report | |
| def main(): | |
| parser = common_parser(__doc__.split("\n")[0]) | |
| args = parser.parse_args() | |
| set_seed(args.seed) | |
| # Section 5.30: four candidates on two objectives, and the two weight vectors. | |
| candidates = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]], dtype=torch.double) | |
| front = pareto_front(candidates) | |
| report = { | |
| "pareto_front": front, # Candidates 1 and 2; 3 and 4 are dominated. | |
| "weighted_at_0.6_0.4": [round(float(s), 3) for s in candidates @ torch.tensor( | |
| [0.6, 0.4], dtype=torch.double)], # 0.60 for the first and 0.64 for the second. | |
| "weighted_at_0.8_0.2": [round(float(s), 3) for s in candidates @ torch.tensor( | |
| [0.8, 0.2], dtype=torch.double)], # 0.70 and 0.62, a reversal of the order. | |
| } | |
| # Section 5.31.2: one MOG-DFM step at a position holding A, with three candidate letters. | |
| changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]], dtype=torch.double) | |
| importance = torch.tensor([1.2, 0.8], dtype=torch.double) | |
| weights = torch.tensor([0.6, 0.4], dtype=torch.double) | |
| delta_s, ranks, direction = combined_score(changes, importance, weights) | |
| report["rank_average"] = [round(float(r), 3) for r in ranks] # (0.500, 0.700, 0.850, 0.450). | |
| report["directional_scores"] = [round(float(d), 3) for d in direction] # (0, 0.100, 0.090, -0.028). | |
| report["combined_score"] = [round(float(s), 3) for s in delta_s] # (-1.511, 1.541, 2.297, -2.327). | |
| posterior = torch.tensor([0.3, 0.3, 0.3, 0.1], dtype=torch.double) # The denoiser at t = 0.7. | |
| base = posterior[1:] / (1.0 - 0.7) # The pretrained rates out of A toward C, G and T. | |
| report["base_rates"] = [round(float(r), 4) for r in base] # (1.000, 1.000, 0.333). | |
| tilted = base * delta_s[1:].exp() | |
| report["guided_rates"] = [round(float(r), 3) for r in tilted] # (4.667, 9.948, 0.033). | |
| angles, cosines = cone_angles(changes[1:], weights) | |
| report["cone_cosines"] = [round(float(c), 4) for c in cosines] # (0.6727, 0.7894, -0.3032). | |
| report["cone_angles_deg"] = [round(float(a), 1) for a in angles] # (47.7, 37.9, 107.7). | |
| inside = angles <= 60.0 # The cone test at Phi = 60 degrees. | |
| survivors = tilted[inside] | |
| report["after_cone"] = [round(float(r), 3) for r in survivors / survivors.sum()] # (0.319, 0.681). | |
| rejected = float((~inside).to(torch.get_default_dtype()).mean()) # One candidate of three. | |
| report["next_cone_angle"] = round(update_cone(60.0, rejected), 1) # 64.1. | |
| # Section 5.32.2: three candidate edits, each rolled out four times. | |
| omega, reference = torch.tensor([0.5, 0.5], dtype=torch.double), torch.zeros(2, dtype=torch.double) | |
| edits = [ | |
| ([(0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5)], [1, 1, 0, 1], 0.6), | |
| ([(1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1)], [1, 1, 1, 1], 1.2), | |
| ([(0.4, 0.4)] * 4, [0, 0, 0, 0], 2.0), | |
| ] | |
| values = [rollout_value([torch.tensor(s, dtype=torch.double) for s in rollouts], | |
| feasible, omega, reference, 2.0)[0] for rollouts, feasible, _ in edits] | |
| base_rates = torch.tensor([rate for _, _, rate in edits], dtype=torch.double) | |
| products = torch.tensor(values) * base_rates # The base rate times the continuation value. | |
| report["rollout_values"] = [round(v, 3) for v in values] # (1.390, 1.181, 0). | |
| report["rollout_products"] = [round(float(p), 3) for p in products] # (0.834, 1.417, 0). | |
| report["rollout_selection"] = [round(float(p), 3) for p in products / products.sum()] | |
| report["rollout_base_selection"] = [round(float(p), 3) for p in base_rates / base_rates.sum()] | |
| # The same rollout on the shared DNA, from a partly revealed sequence with ACG already in place. | |
| splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed) | |
| counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double() | |
| frequency = counts / counts.sum() # [4] the denoiser of the base flow, GC-rich like the data. | |
| # The base rates below are kappa_dot/(1 - kappa) times that denoiser, read at t = 0.3. | |
| revealed = torch.full((args.length,), MASK, dtype=torch.long) | |
| revealed[2:5] = torch.tensor([0, 1, 2]) # Positions 3, 4 and 5 hold A, C and G. | |
| finish = {letter: revealed.clone() for letter in "TC"} | |
| for letter, candidate in finish.items(): | |
| candidate[5] = DNA.index(letter) # Position 6 takes T, completing ACGT, or C, and C does not. | |
| seen = {letter: dna_rollout_value(candidate, frequency, args.steps, OMEGA, REFERENCE, BETA) | |
| for letter, candidate in finish.items()} | |
| rates = torch.tensor([float(frequency[DNA.index(letter)]) for letter in finish]) / (1.0 - 0.3) | |
| guided = rates * torch.tensor([seen[letter] for letter in finish]) # The base rate times hat h. | |
| report["dna_rollout_value_T"] = round(seen["T"], 4) # Every continuation carries the motif. | |
| report["dna_rollout_value_C"] = round(seen["C"], 4) # Only the ones that plant it elsewhere. | |
| report["dna_base_rates_T_C"] = [round(float(r), 4) for r in rates] # C has the higher base rate. | |
| report["dna_base_selection"] = [round(float(r), 4) for r in rates / rates.sum()] | |
| report["dna_guided_selection"] = [round(float(r), 4) for r in guided / guided.sum()] | |
| print_report("several objectives at once", report) | |
| save_run(args.out, config=vars(args), report=report) | |
| if __name__ == "__main__": | |
| main() | |