Antigen-Specific Antibody Design via Direct Energy-based Preference Optimization

Thank you for your strong support! Please see below for our responses to the comments.

Q1: Lack of explanation of SE(3)-equivariant neural network. The author uses the diffusion model with such an equivariant neural network from Luo et al. Lacking such an explanation may hurt the understanding of the whole method.

A1: Thanks for your valuable suggestion. We will complete the explanation in the revision.

Q2: There are some writings inconsistent.

A2: Thanks for your careful inspection, we will fix it in the revision.

Q3: Why we need to use SE(3)-equivariant neural network?

A3: SE(3)-equivariance is a crucial property in protein design because the structure of proteins is independent to the observation view. Utilizing SE(3)-equivariance can ensure stable and predictable performance in the presence of nuisance transformations of the data input[1]. In addition, we chose to build our method on DiffAb, which uses equivariant neural networks. Other baselines, such as MEAN and dyMEAN, are also E(3)-equivariant. Therefore, we retained the equivariant neural networks to ensure a fair comparison. Nevertheless, we would like to emphasize that our proposed method is not specific to equivariant NNs and can be applied to other base models beyond equivariant NNs.

Reference:

[1] Fuchs, Fabian, Daniel Worrall, Volker Fischer, and Max Welling. "Se (3)-transformers: 3d roto-translation equivariant attention networks." Advances in neural information processing systems 33 (2020): 1970-1981.

Q4: What is the inference time for generating one sample or one batch?

A4: Our model could be regarded as an aligned model to the specific preference of a pre-trained model, so the inference time will not be changed. We record the detailed time cost in each antigen-antibody complex and show them below (batch_size=16, single A100 40G GPU). The inference time consists of two parts: model inference time and .pdb file reconstruction time. The inference time is related to data length. Most antigen-antibody complexes are truncated to 256 residues when fed into the model, and the corresponding inference time is around 16-17s. The reconstruction time is determined by the total number of residues in the complex. For the complex with 500 residues, the reconstruction time of one batch of samples is around 2s. We list the detailed time cost below (the last few rows are dropped due to the limitation of space)

Thank you for your helpful feedback! We address your questions as follows.

Q1: AAR is about 10% lower than dyMEAN. The authors did not provide a solid reason for what caused the result.

A1: As mentioned in Appendix A, AAR is easily hacked and conceals numerous issues. Certain intrinsic patterns within CDR-H3 sequences enable the model to achieve a seemingly satisfactory AAR by memorizing these patterns (also noted in the Limitations section of dyMEAN). However, this leads the model to produce highly similar sequences for all antigens, which is highly impractical. The possible reason is the severe lack of data, causing the model to merely learn the marginal distribution of residue types at each position. dyMEAN is a typical example of this scenario. To illustrate the sequences generated by dyMEAN with over 40% AAR, we reproduced dyMEAN, generated CDR-H3 for 55 antigens, and visualized the sequences. As shown in Fig. 2 in the uploaded pdf, through alignment (aligned and visualized by MAFFT[1]), it is evident that dyMEAN generates almost identical sequences, while AbDPO+ is considerably better. In this case, AAR is not a meaningful measure.

As for why AbDPO did not achieve a high AAR, we believe it is due to the different learning objectives between AbDPO and the baselines. The learning objective of baselines is to generate antibodies consistent with natural antibodies in terms of sequence and structure. Therefore, they perform better on metrics like AAR, which measures similarity to natural ones. However, AbDPO's learning objective is to better meet specific preferences, while not optimizing for consistency with natural ones. Consequently, its AAR is lower than dyMEAN. Nevertheless, sequence pLL leads to a significant improvement in AAR (31.25% to 36.27%) when used as an optimization objective.

Reference:

[1] Madeira, Fábio,et al. "The EMBL-EBI Job Dispatcher sequence analysis tools framework in 2024." Nucleic Acids Research.

Q2: The adaptation of RLHF into this approach should be further clarified.

A2: We can consider the antibody generation model obtained through conventional training methods as a pre-trained model. However, this pre-trained model does not meet our requirements, such as physical energy. This situation is similar to the challenges faced in current generative AI, where pre-trained models often require further alignment to match human preferences. Therefore, we can leverage methods like RLHF to align the antibody generation models with the desired properties, essentially optimizing the model. In AbDPO, we optimize the pre-trained model using such methods, allowing us to generate antibodies that satisfy multiple preferences.

Q3: Would it be helpful to directly train the diffusion model from scratch by optimizing the loss function using pyrosetta?

A3: Directly training the diffusion model from scratch is not practical.

The reasons are twofold:

1. Since the loss function using pyrosetta is not differentiable, RL algorithms can be applied, such as policy gradient methods. In this case, the antibody design problem is formulated as a high-dimensional non-linear decision making problem. The search space for this problem is large and involves both continuous and discrete variables. Directly training from scratch is hard due to its instability. Specifically, useless efforts might be made, especially at the beginning of the training, when the policy can only generate samples with unsatisfactory rewards. A feasible solution to address this challenge is to use a pre-trained policy that can significantly prune the search space and accelerate the convergence rate, which has been validated on many applications, such as [1] and [2]. And this solution is the pretraining-finetuning paradigm that we utilized in our work. The pre-trained diffusion model can be viewed as an expressive policy.

2. Another reason is that low pyrosetta energy is not our only purpose. We aim to generate valid antibodies with low energy. Some intricate properties cannot be explicitly formulated as rewards but can be learned from data by generative modeling. Our optimization objective is equivalent to maximizing rewards with regularization with reference model. The regularization term keeps knowledge learned from data, which cannot be achieved by training from scratch by optimizing the energy function.

References:

[1] Silver, David, et al. Mastering the game of go without human knowledge. Nature

[2] Haarnoja, Tuomas, et al. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. ICML.

Q4: Would it be helpful to test the method with other downstream tasks that may be closely related to the energy

A4: We test our method on the affinity optimization. Specifically, we test AbDPO, along with its counterpart, DiffAb by evaluating the best $\Delta\Delta G$ against the reference antibodies among 300 generated samples using [1] on 16 randomly selected antigens in the test set. Note that we do not use the metric as the signal in fine-tuning. The results are shown as follows:

Our results significantly outperform DiffAb, which shows the generalizability on other tasks. We believe that our method can perform even better if we consider the metric itself, ie., predicted $\Delta\Delta G$ against the reference antibodies, in our preference definition when fine-tuning.

Reference:

[1] Shan Sisi, et al. "Deep learning guided optimization of human antibody against SARS-CoV-2 variants with broad neutralization." Proceedings of the National Academy of Sciences

Thank you for your detailed feedback. Please see below for our responses to the comments.

Q1: The approach can design better binders, but comes at the expense of increasing the number of hydrophobic residues which is typically associated with non-specific binding.

A1: This is exactly why we included PHR as an optimization objective in AbDPO+. We will show the performance of each method when limiting the PHR in A6, and AbDPO still performs better.

Q2: It would be interesting to see the results relative to different physics-based simulators. I’m not sure how many different simulators were used to generate “fine-tuning” signal?

A2: All the three energy signals were calculated using Rosetta, which is capable of both energy calculation and side-chain packing, and is widely used by researchers. In early stages, we also tried using OpenMM for energy calculations and obtained consistent results. Additionally, we used OpenMM to calculate the potential energy of the antigen-antibody complex where CDR-H3 is generated by AbDPO fine-tuned according to Rosetta energy. For 6 complexes tested, AbDPO achieved lower energy in 5 compared to DiffAb. (the remaining 49 complexes raise errors in OpenMM workflow when adding Hydrogens, and are not influential in Rosetta workflow that does not require Hydrogens)

Q3: Would it be possible to include additional metrics such as lDDT, TM, and some metrics characterizing the fit on angles?

A3: For the antibodies generated by all methods, we measured both the TM-score[1] on the heavy chain (Hc_TM-score) and the torsion performance (Torsion-score) of residues on the CDR-H3. For Hc_TM-score, there is not much difference in methods other than HERN. The Torsion-score is derived from a two-dimensional KDE function, which is based on the joint distribution of Phi and Psi torsions observed in natural CDR-H3s. Diffusion-based methods, DiffAb and AbDPO, perform significantly better than other methods.

Reference:

[1] Zhang, Yang, and Jeffrey Skolnick. "TM-align: a protein structure alignment algorithm based on the TM-score." Nucleic acids research.

Q4: Could you also list how many mutations from the original CDR sequence are the improved binders (Table 1)?

A4: Yes, we calculate the minimum number of mutations of the generated samples that exhibit energy levels close to the natural samples (i.e., the "successful" samples in Table 1) and the corresponding CDR lengths.

Q5: Can you share the structural similarity scores between epitopes in the train and test folds?

A5: We follow the metric widely used in protein structure design tasks to calculate the structural similarity scores. The similarity between two protein sets is defined as the average TM-score of each sample in one set to its most similar protein in another set. The similarity between the train and test folds is 0.6965. Following MEAN, the epitope is defined as the 48 residues of the antigen closest to the paratope.

Additionally, we believe there is a misunderstanding here. The training data only appears during the pre-training phase, while in the optimization phase, we only used synthetic data.

Q6: If you limit the fraction of hydrophobic residues to the value for the original complex and accept only designs with lower score, how does Table 1 look?

A6: We perform an evaluation only on the samples that contain hydrophobic residues not exceeding the natural one, and the results are shown below. It can be seen that under this setting, the energy performance of almost all methods has deteriorated, but AbDPO and AbDPO+ still perform the best in terms of the two energies.

Q7: Can you list relative to how many test fold complexes each method achieves better binding free energy?

A7: Sure. For each complex, if a method achieves the best performance, it was considered a "win" for that method on that specific complex. We record the number of "win" complexes, N_win. We also observe that there are instances where none of the samples have a PHR lower than that of the natural sample. Therefore, we also list the number of complexes, denoted as N_phr, for which the method can generate samples with a PHR not exceeding that of the natural sample. The specific results are as follows:

Q8: How does Table 1 look like if you apply ITA to the vanilla baselines?

A8: We follow the ITA setting from the MEAN codebase and implement ITA for MEAN, dyMEAN, and DiffAb. We adapt $\text{CDR}E\_{\text{total}}$ and $\text{CDR-Ag}\Delta G$ as two targets to gather high-quality candidates for each antibody. We observe that directly generating high-quality candidates can negatively affect other preferences. This further demonstrates the effectiveness of AbDPO in optimizing multiple objectives simultaneously.

Thank you for your feedback. Please see below for our responses to the comments.

Q1: It is computationally expensive to guide diffusion models using Rosetta.

A1: Yes, the computational expense is indeed a drawback of Rosetta. We are aware of Rosetta's limitations and have discussed them in the appendix. However, despite these issues, the time cost of optimizing AbDPO for a single antigen is less than a day. Additionally, AbDPO is not specific to Rosetta. For example, other packing tools (like DiffPack) and energy calculation methods (such as OpenMM, which is significantly faster than Rosetta) can also be used. In fact, AbDPO supports optimization of any type of property, and we chose Rosetta because it can perform both packing and energy calculations and is widely used.

Q2: Rosetta side-chain packing is stochastic and non-deterministic. Therefore, if we relax the structure generated by diffusion model multiple times, the calculated Rosetta energy will be very different, and the standard deviation can be very high (sometimes it can be twice or three times higher than the mean). In other words, it is very tricky to construct a preference dataset because you need to compare the binding energy distribution between two CDR sequences, and their standard deviation is very high.

A2: We recognize that the same CDR-H3 (with the same sequence and structure) may result in different side-chain conformations after packing, and there will be energy differences between these side-chains. We observed this in our early exploration but still only performed the side-chain packing for each sample once. The reasons are as follows :

• The differences between different samples generated by existing methods (such as DiffAb) are generally much larger than the energy differences between different side chain conformations of the same sample. For natural antibodies, total energy is relatively low due to their reasonable structure, so different side chains may have a significant impact on total energy. However, for generated antibodies with several clashes, the influence of side-chain conformations is not significant as the total energy could be extremely high.
• Repeating the packing process on the same sample can indeed yield more accurate energy labels, but it will consume more time and we do not want the excellent performance of AbDPO to be based on a large amount of meticulously selected data. Although sampling only once may lead to incorrect judgments of win-lose relationships between two samples, this usually occurs when the two samples are quite similar. As long as the better sample has a higher probability of being sampled at lower energy, the model will generally be optimized toward generating antibodies with lower energy.
• DPO is derived based on the Bradley-Terry model. In the Bradley-Terry model, the labels of winner and loser are not deterministic but exist in probabilistic form. We perform packing for each sample only once, which is similar to sampling winner and loser once. This is also theoretically justifiable.

Furthermore, we use 1a14 as an example and sample 100 synthetic data from the 1a14 training dataset. For each synthetic sample, we repeat packing 128 times with different random seeds. Then we can compare the deviation between different synthetic data and different sidechains from the same synthetic data. The results are shown in Fig. 1 in the uploaded pdf, and the deviation between different synthetic data (the red line) is far greater than the deviation between different sidechains from the same synthetic data (the blue violin plot indicates the distribution of deviation within each synthetic data). This further supports our justification.

Q3: Can you report your model performance by running Rosetta relaxation multiple times with different random seeds?

A3: We repeat the optimization process 32 times with different random seeds for each AbDPOw/O generated sample. The average $\text{CDR} E_{\text{total}}$ and $\text{CDR-Ag} \Delta G$ of all generated antibodies are 65.19 and -2.67, which are slightly different from the reported values of 69.82 and -3.00. Therefore, the experiment results reported in our paper are reliable. The reason we did not report the averaged results over multiple runs is due to time cost considerations.

Q4: Can you report standard deviation of the reported binding energy for each method?

A4: Thanks for your reminder, we provide the standard deviation of two energies for each method below. It could be seen that significant deviations are commonly present in the results generated by each method, and the deviation of AbDPO is smaller than that of DiffAb.