Overview
I think that post-training is going to be really critical for bioML if we want to make sure that our models and systems can still be useful as the problems and targets we try to treat for become harder and harder. As an illustrative point, a pretrained antibody model knows what antibodies look like. If we ask it for a “good antibody”, then it’ll tell us what it confidently knows: antibodies that are within distribution. With post-training, we can readjust that pretrained distribution so that our models can better provide candidates for a “good antibody.”
Partly to learn, partly to actually try to contribute to getting there, I’m working on a more narrow approach to get a model to answer our questions more accurately. How can we steer a protein language model toward a binding objective, while making sure we keep its understanding of biology? In order to do that, I’m looking at how we can modify HCDR3 in the outputs of antibody language models — it’s the loop that carries most of the binding and relies much less on the conservation, so I think that’s where steering has the most expressivity in outcome. This looks like exploring different post-training methods (so far, mainly preference optimization techniques) and different mech-interp approaches to see how we can incorporate known biological heuristics to better steer our models (and be more data-efficient).
A couple findings I keep circling around are:
- If you push too hard on a model towards a function, it’ll reward-hack and give you all these sequences that score super well but fold super poorly. I think it’s interesting to explore which pretrained constraints survive this push and what ends up getting written over.
- Finding ways to optimize modelling and data usage have been really critical to getting this project pushed forward (especially since I, unfortunately, don’t have a thousand H100s to train on). It’s becoming a better consideration for me to figure out where to prioritize compute-efficiency and data-efficiency.
Right now, I’m working on adapting p-IgGen to HER2 and figuring out reproducible ways to make this work.
The current experiment: editing one loop of trastuzumab
Trastuzumab’s a great binder for HER2. The current experimental design is keeping the antibody fixed except for ten residues in
its HCDR3 binding loop and letting p-IgGen choose them. A
published library
measured 524,346 variants of that loop, grouped by binding as high, mid
or low. Since I’m constraining this experiment to one antigen, the model will only see (and have been pretrained on) the antibody sequences. Only in post-train does the model see binding affinities.
The first preference campaign seemed really promising until I looked more into the outputs. The benchmark, unfortunately, is pretty saturated, but after only 3 minutes of training, it ranked measured high binders a bit better than the SFT’ed model (0.9729 vs 0.9701 average precision). Worse than the mild gain, only 42% of 10,000 generated sequences were distinct, against about 98% for supervised training.
So far, to try to fix that, I’ve added a stop rule for runs that lose probability on sequences the parent model favored and then tried out IPO rather than DPO; my rationale being that targeting a finite gap prevents overtraining. I also added a “parent replay” constraint; where I sampled sequences from each starting model and trained the updated model to keep them probable (part of the goal here was to make sure we were remembering biology, not just memorizing a good CDR3 loop). IPO consistently beat matched supervised training by ~ 1.6 points in validation average precision (this metric was there to give more weight to the sequences that diverged from anything seen in the training dataset). This moderate replay also cut down the mode collapse roughly 5x while still retaining most of IPO’s ranking gain.
Something really interesting was that IPO improved this comparison that never got trained on. Because the preference pairs were only high and low variants, the model never saw the mid variant but, in testing, the model became better at accurately ordering high, mid, and low. IPO also increased measured dependence between residue positions in the loop (although stronger replay hyperparameters pulled it back); this could potentially be the model learning useful residue combinations but I want to test it a bit more, to make sure it’s not just memorizing a few patterns.
Residue combinations in the binding loop / Preliminary results
So far, the most interesting result has been that the model may be learning which residues work together. On 231 measured pairs that differed at heavy-chain position 98, IPO with replay puts the stronger-binding variant first 83.5% of the time. In a separate run that removed nearby training sequences (to make sure the model wasn’t just learning Hamming distance from training data), IPO with replay’s full score gained higher precision compared to a score that just came from separate “good residue” scores marginalized. My interpretation of this so far is that this comes from these complex conditionals.
A clue as to what these complex conditionals mean in a biological sense is threnonine-histidine-threonine (T-H-T) at positions H96-H98. T at H96 and H at H97 are on the lower side (represented in the low-binding group more often), but T-H-T more frequently fell into high-binding. In the literature, H98 can tolerate radical changes: D98W improved trastuzumab affinity.
I think it’d be really interesting to know whether post-training finds combinations the parent overlooks. T-H-T is a rare pattern in the measured library so this could be a good insight.
What I am testing now
Roughly 90% of the validation variants have a training neighbor that’s one mutation away. The current post-training run is checking whether the performance gain survives if we remove those “easy scoring” variants from training and, also, if those changed residue combinations actually help distinguish binding groups.