CIS6270 / lecture_5 /multi_objective.py
pranamanam's picture
Upload 87 files
0ba9d09 verified
Raw History Blame Contribute Delete
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()