Design of Ligand-Binding Proteins with Atomic Flow Matching

We really appreciate your valueable feedbacks on our paper. Thank you for affirming the technical contribution and the overall experiment results. We address the issues you discussed below, and hope the explanation may lead to a more positive judgement of our paper.

Q1. A unified distance encoding.

Thanks for pointing out a possible enhancement of the distance encoding in AtomFlow. During our training, we find that the model is quite capable of dealing with the uneven split of binning. On the contrary, directly provide the original distance resulted in unsatisfactory training trajectory. We hypothesize that a more universal distance encoding should encode embed the raw distance into a high-demensional vector, and this is a great direction for future exploration.

Q2. Training details.

We've added more training details to the appendix. Regarding your question, ligand tokens are not given during the first training stage. The first stage trains an unconditional protein generation model, while the second stage turns it to a conditional protein binder + ligand conformer generation model.

Q3. On the training objective.

We've added the following explanation to the appendix. The FAPE loss is defined as an average of all pairs of tokens in the original paper, so the calculation process first yield a FAPE matrix and then produce the average value of the matrix. The protein-protein, protein-ligand and ligand-ligand loss calculates the average value of the sub-matrixs defined as (row: protein, col: protein), (row: protein, col: ligand), and (row: ligand, col:ligand).

Q4. On the ablation.

First, we sincerely apologize for a typo in the original text. The training steps for the two stages should be 400k and 300k, rather than 40k and 30k as originally stated. This error has now been corrected.

Regarding the ablation, since training a protein design model is significantly time-consuming, the design choices of our training strategy is largely determined by grid searching possible design space and save the training trajectory of the first 30~50k steps. We compare the training trajectories and select the best configuration that meets the following criteria:

• a) The final distogram loss should close the minimum we get among the configurations (around 2.0).

• b) The L_CFM-FAPE should not decline too fast at the first 10k steps. The first 10k steps is for the transformer stack to learn a relatively steady output, indicated by the decline of the distogram loss. A decline of L_CFM-FAPE at this stage will resulted in an undesired local minimum. Then L_CFM-FAPE should decline fast right after the distogram loss turns to decline much smoother. We select the configuration with the lowest L_CFM-FAPE at the end of training.

We decide the end of each training stage when the training converges, with the following criteria:

• a) the decline rate of every single loss is small.

• b) the structural violence of sampled structures (counts of CA atom violation) converges.

An initial study on directly train the second stage shows unsatisfactory training trajectory. Since the ligand conformer is way easier to generate compared to protein folds, the FAPE loss declines too fast even before the distogram loss, resulted in unstable TransformerStack output, and leading to a diverge of the model after ~30k steps. The resulted model with minimum loss is able to predict the ligand structure, with random protein residue position, which is unusable.

Due to resource limit, we could not provide a comprehensive ablation study since only the configuration we selected has been trained to the end. We hope the details regarding our training process provide enough information to justify our decisions.

We also want to share a new experimental result with you. During the review process, several AlphaFold 3 replications, including the original models, were released. Leveraging these AlphaFold 3-like model predictions, we developed what is likely a more reasonable in silico metric for evaluating binding affinity. On this metric, AtomFlow outperforms RFDiffusionAA. The detailed results have been included in the appendix of our paper, and a more comprehensive analysis will be provided in the final version. We hope this introduces an additional evaluation approach that can benefit the community.

We hope our explanation has addressed your questions and concerns thoroughly. We sincerely appreciate your valuable feedback and are always happy to clarify further or discuss ways to improve our work. If you find our responses satisfactory, we would be grateful for your consideration of a higher score. Thank you!

Thank you for your insightful feedback. We recognize your concerns regarding the novelty and experimental sufficiency of our work. In this rebuttal, we address these issues comprehensively by:

• Demonstrating AtomFlow’s methodological contributions and its distinctions from prior works.
• Presenting additional experimental comparisons and explaining the exclusion of certain baselines.
• Discussing the applicability and limitations of AlphaFold 3 in our setting.

On Methodological Novelty

We appreciate your comments regarding the clarity of AtomFlow’s novelty. We acknowledge the relevance of the works you mentioned and have expanded our discussion to emphasize AtomFlow’s unique contributions. Below, we explicitly compare our methodology with existing literature to clarify its novelty.

• A new paradigm for ligand-binding protein design

  AtomFlow addresses a key challenge in real-world flexible design of ligand-binding proteins: the lack of predefined templates or fixed binding poses. Existing methods often assume a known ligand conformation [2] or binding template [3,4], which is rarely available in de novo design. This gap limits their applicability to novel targets [1]. AtomFlow overcomes this by generating both protein and ligand structures without requiring prior knowledge of binding poses, which provides a useful tool for biological discovery.

• A novel architecture for direct design of interaction

  A critical aspect of ligand binder generation lies in the fusion of two different types of biomolecules. Protein design models commonly use rigid frames for translation and rotation [5,6], whereas ligand generation methods rely on torsion angles derived from chemical graphs [7]. These distinction pose challenges for capturing spatial interactions between them. AtomFlow addresses this with a unified architecture that directly enables the modelling of biological interaction between small molecule atoms and protein residues, treating both of them as a series of movable biotokens, and defininig a probability flow on representative atomic coordinates. This unified approach also supports the generation of diverse biomolecule types, though this work primarily focuses on proteins and ligands. AF3 has tackled similar challenges, but it is (1) primarily centered on predicting interactions rather than designing interactions, and (2) it was not publicly available at the time our work was completed. This was also mentioned in your review, and we will discuss it in more detail later.

• A probability flow on conformation, not representation

  AtomFlow defines probability flow directly on molecular conformations (which could be represented by many different coordinate representations under global translation and rotation) rather than specific coordinates. This approach treats all SE(3)-transformed variants of a conformation as equivalent, ensuring that the training objective is invariant to global translation and rotation. In contrast, mainstream protein design models [5,6,7,8] rely on coordinate-based representations and train with mean squared error (MSE) loss. This penalizes structurally similar conformations under different SE(3) transformations, making the generation process less robust. By avoiding these limitations, AtomFlow offers a more natural framework for biomolecular design.

• Augmented model structure for unified features

  We repurposed key building blocks from AlphaFold 2 to construct a prediction network capable of generating structures from noisy inputs. While leveraging core components such as IPA, we enhanced the architecture to support arbitrary biomolecules by introducing atom frames and a unified biotoken embedder. This design allows for flexible conditional inputs at both sequence and token pair levels, enabling the model to accommodate diverse control requirements without further modifications. These improvements not only extend the applicability of our framework but also lay the groundwork for future extensions in biomolecular design.

We've revised Sections 1 of our paper to explicitly highlight our methodological contribution.

On AlphaFold 3

Regarding AlphaFold 3, while its diffusion-based approach represents an advance in biomolecule prediction, its heavy reliance on predefined sequences limits its applicability to de novo protein-ligand design. In our experiments, neither AlphaFold 3 nor its replications (e.g., Protenix [9], Chai [10]) could reliably generate ligand-binding proteins from scratch. This highlights AtomFlow’s tailored capabilities for this task.

• While AlphaFold 3 does not directly apply to our setting, we found your suggestion of adapting a structure prediction model for design tasks to be highly insightful. Inspired by this, we conducted a preliminary experiment using two AlphaFold 3 replications, Protenix and Chai, which were forcibly adapted to accept a series of [UNK] residues along with a ligand (FAD) as input. Our findings are as follows: approximately 20% of proteins generated by Protenix achieved scRMSD < 2, with reasonable ligand conformations but no visible interactions between ligand and protein. Chai performed slightly better, with around 50% of its generated proteins achieving scRMSD < 2, though none of the ligands exhibited correct conformations. Notably, the original AlphaFold 3 model is hard-coded to require a valid residue sequence, rendering direct comparisons infeasible. These results suggest that while AlphaFold 3 shows potential, it currently cannot generalize to binder design without substantial retraining and engineering, which falls outside the scope of this work.

• At the time of writing our draft, no publicly available AlphaFold 3 model was available for local experimentation. Inspired by your suggestion, we revisited AlphaFold 3 and leveraged recent replications (e.g., Chai) to conduct a new experiment. Specifically, we developed an alternative in silico metric for binding affinity based on minimum interchain pAE predictions, inspired by AlphaProteo [11]. Using this metric, AtomFlow demonstrated superior results compared to RFDiffusionAA. This new experiment has been added to the appendix, and we hope that introducing this metric to the community will further support our contribution.

• Additionally, we hypothesize that AlphaFold 3 could serve as a foundation for protein hallucination approaches, similar to BindCraft [12], which utilizes AlphaFold 2. Although this direction lies beyond the scope of our current work, it represents a promising avenue for future exploration.

We hope our explanation has addressed your questions and concerns thoroughly. We sincerely appreciate your valuable feedback and are always happy to clarify further or discuss ways to improve our work. If you find our responses satisfactory, we would be grateful for your consideration of a higher score. Thank you!

[1] Bick, M. J., Greisen, P. J., Morey, K. J., Antunes, M. S., La, D., Sankaran, B., ... & Baker, D. (2017). Computational design of environmental sensors for the potent opioid fentanyl. Elife, 6, e28909.

[2] Krishna, R., Wang, J., Ahern, W., Sturmfels, P., Venkatesh, P., Kalvet, I., ... & Baker, D. (2024). Generalized biomolecular modeling and design with RoseTTAFold All-Atom. Science, 384(6693), eadl2528.

[3] Zhang, Z., Shen, W. X., Liu, Q., & Zitnik, M. (2024). Efficient generation of protein pockets with PocketGen. Nature Machine Intelligence, 1-14.

[4] Stärk, H., Jing, B., Barzilay, R., & Jaakkola, T. (2023). Harmonic self-conditioned flow matching for multi-ligand docking and binding site design. arXiv preprint arXiv:2310.05764.

[5] Yim, J., Trippe, B. L., De Bortoli, V., Mathieu, E., Doucet, A., Barzilay, R., & Jaakkola, T. (2023). SE (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277.

[6] Yim, J., Campbell, A., Foong, A. Y., Gastegger, M., Jiménez-Luna, J., Lewis, S., ... & Noé, F. (2023). Fast protein backbone generation with SE (3) flow matching. arXiv preprint arXiv:2310.05297.

[7] Lin, Y., Lee, M., Zhang, Z., & AlQuraishi, M. (2024). Out of Many, One: Designing and Scaffolding Proteins at the Scale of the Structural Universe with Genie 2. arXiv preprint arXiv:2405.15489.

[8] Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., ... & Jumper, J. M. (2024). Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, 1-3.

[9] https://github.com/bytedance/Protenix

[10] Chai Discovery, Boitreaud, J., Dent, J., McPartlon, M., Meier, J., Reis, V., ... & Wu, K. (2024). Chai-1: Decoding the molecular interactions of life. bioRxiv, 2024-10.

[11] Zambaldi, V., La, D., Chu, A. E., Patani, H., Danson, A. E., Kwan, T. O., ... & Wang, J. (2024). De novo design of high-affinity protein binders with AlphaProteo. arXiv preprint arXiv:2409.08022.

[12] Pacesa, M., Nickel, L., Schmidt, J., Pyatova, E., Schellhaas, C., Kissling, L., ... & Correia, B. E. (2024). BindCraft: one-shot design of functional protein binders. bioRxiv, 2024-09.

On Experiment Sufficiency

We appreciate your suggestion regarding additional baselines such as PocketGen [3] and FlowSite [4]. While these models address ligand-protein interactions, their focus is limited to refining pocket residues within a predefined radius. They lack the capacity to design full protein folds, making direct comparison with AtomFlow infeasible. Following your suggestion, we further conducted an unfair experiment with PocketGen by providing a template binder to it, as detailed below. Despite this, the results demonstrate that AtomFlow consistently outperforms PocketGen in terms of fold quality across all radii.

For this experiment, we used the natural binders of four ligands—FAD (7bkc), SAM (7c7m), IAI (5sdv), and OQO (7v11)—as input. To evaluate the design capability of PocketGen (PG) under different constraints, we progressively increased the design radius (minimum distance to ligand) from 3.5 to 9.5. The masked target area expanded with the radius, requiring the model to redesign increasingly larger regions of the protein. When the radius exceeded the protein's dimensions (radius ≥ the protein size), all residues were masked, simulating our full-design setting. The table below presents the min/median scRMSD values for designs generated by PocketGen at each radius. For reference, scRMSD < 2 is generally considered a successful design. Notably, PocketGen’s performance deteriorated significantly as the radius increased, reflecting its reliance on template residues. (At radius=8 for OQO, PocketGen generated designs with several residues misaligned with the ligand, leading to abnormally high scRMSD values.) Additionally, PocketGen does not support radius settings beyond 10, preventing direct simulation of AtomFlow’s fully template-free design scenario. The results of AtomFlow is derived from our main experiment. |ligand|AtomFlow (r=inf)|PG (r=3.5)|PG (r=5)|PG (r=6.5)|PG (r=8)|PG (r=9.5)| |-|-|-|-|-|-|-| |FAD|0.79/3.74|7.10/7.38|6.75/7.81|7.29/8.35|20.92/24.23|23.12/25.23| |SAM|0.83/2.01|2.12/2.62|2.77/2.99|2.94/4.03|12.39/14.49|13.79/14.74| |IAI|0.56/1.82|0.71/0.85|0.95/1.02|2.04/2.28|3.59/5.53|9.02/11.71| |OQO|0.59/1.63|1.20/1.26|1.70/1.79|2.40/2.45|11.59/11.94|2.13/2.41|

Notably, PocketGen’s performance deteriorates significantly as the radius increases, highlighting its reliance on template information. In contrast, AtomFlow maintains high performance without such dependencies. Based on similar settings, we expect FlowSite to exhibit comparable performance patterns to PocketGen. We added explicit description on the exclusion of these methods in our paper.

We thank a lot for the detailed response of the authors. The extra expriments solve a lots of my concerns. Besides, the detailed explanation of the methods compared with other competitors is convincing. However, I still think that the contribution of method novelty is limited, since the use of probabilistic models has already been thoroughly explored in the relevant field, and the application of flow matching in molecular and protein design has become quite common.

Therefore, after reviewing your response, I reread the article and reassessed its quality. I have decided to raise my initial rating from 3 to 5.

Thanks for your review on our work. We are really happy to have a feedback from your hands-on experience of AtomFlow, and we are glad to see that AtomFlow successfully designs protein binders for your selected ligands. We address each of your concerns or the issues encountered below.

1. The infinite loop.

We'll remove the infinite loop in the final release.

2. The diversity and novelty of designs.

a. Diversity. We replicated your encountered situation that AtomFlow consistently generates one same fold for a given ligand in our released version, but the version used in evaluation seems not showing similar phenomenon. This should be a bug and will be fixed in the final release. We've generated 10*5 binders each for the L1000 set in CASP16 with lengths in [100, 150, 200, 250, 300] and uploaded them to https://anonymous.4open.science/r/AtomFlow-suppl-data-7384/L1000_results.zip .

b. Novelty. We admit that the overall designs generated by AtomFlow tends to be conservative. After intensive search, the most novel fold found by AtomFlow in our experiment has pdbTM=0.347. We believe there's space for us to explore to increase the novelty of designs, and thank you for pointing out the novelty issue. We're trying to improve the design novelty of our model, perhaps by adding more training data and designing a better training strategy.

3. The affinity results.

a. Energy function as a proxy. Thanks for pointing out that Vina is not a best proxy of affinity. It seems that you somewhat misunderstood the motivation of our affinity experiment. Both AtomFlow and RFDiffAA are not designed to predict the binding affinity, and the results we report is an in silico calculation on the structures they generate. The experiment is to confirm that AtomFlow is really generating binders for the given ligand, rather than generating a random fold without considering the target. We know that there does not exist an in silico metric that perfectly aligns with real binding affinity, but a metric is needed and we follow the practice of exising literature [1,2] to adopt an energy calculated by docking method as a proxy. We also thank you for suggesting adding the ground truth value of Vina energy to the paper, and we'll update the figures in our paper. The ground truth energy for the four ligands are: IAI (-9.864), FAD (-13.459), SAM (-7.115), OQO (-9.594). AtomFlow successfully designs binders with lower Vina energy than ground truth for IAI, SAM and OQO. No methods (including RFDiffAA) generate binders with lower Vina energy than ground truth for FAD.

b. Data splitting (Your Question 3). We properly splitted the PDBBind dataset following the official split, following DiffDock [3]. FAD and SAM are in the training set, while IAI and OQO are not. We adopt this split since the main purpose of AtomFlow is to generate binders for ligands no matter seen or not. We believe splitting data in a way that hides several folds from training and see whether the model can generate an unseen fold is a very interesting research direction. However, this is beyond the scope of our research and should better start from addressing the novelty issue on unconditional protein design.

c. AlphaFold 3-like model as an affinity proxy. Inspired by your comments on Vina energy, we feel it necessary to introduce an additional in silico proxy for binding affinity since the paper will set up a standard for future research to follow. We noticed that AlphaProteo [4] adopts several metrics produced by AlphaFold 3 as in silico filters. At the time we finished our draft, there's no publicly available AlphaFold 3 for us to run locally. Recently, several AlphaFold 3 replications and the original AlphaFold 3 are released. Though we're not abel to run the AlphaFold 3 model locally due to a licensing issue, we developed an alternative in silico metric based on Chai-1 [5], calculating the minimum value across all interchain terms in the PAE matrix (min_ipAE, lower is better). Note that this metric is proved to be a good indicator for protein-protein binder design, but not validated on small molecule-protein binders. We've added the results of AtomFlow-N and RFDiffAA to the appendix. The results shows that the AtomFlow-generated binders have lower min_ipAE than the ones from RFDiffAA, and both models have the ability to generate binders with similar min_ipAE as natural ones. We're working to develop better metrics as in silico proxy for ligand-binding protein design based on wet lab verification.

Thank you for your detailed feedback on our submission. We appreciate your insights, particularly your concerns regarding novelty and the sufficiency of experiments, which may have influenced your current assessment of our work. In this rebuttal, we aim to address these points comprehensively and provide further evidence to clarify the contribution of our work.

1. The contribution of our work

a. Significance of the Problem

AtomFlow seeks to bypass the limitation of requiring a fixed binding pose of ligand in RFDiff and RFDiffAA, since the bound structure of ligand is often unknown in our designing practice. A bound conformer usually diverges greatly from the ideal structure, which is shown in the introduction of our paper, and it's difficult to determine which pose will be optimal before design since it's determined by the binder [1]. In real practice, a known bound conformer is usually unavailable for a novel target [2], which greatly limits the adoption of RFDiffAA in the workflow.

b. Novelty of our Methodology

We proposed a novel general framework for an important biological problem, which is built on top of AlphaFold 2 but with non-trivial augmentation and adaptation. AtomFlow address the problem of co-modelling protein and small molecule ligands, by introducing a more generalized biotoken representation and a probability flow model on the representative atoms. We acknowledge that the final form of our probability flow is similar to AlphaFlow. However, the derivation, noise form, and final loss is different from theirs. Furthermore, AlphaFlow is designed to generate multiple protein ensembles for a FIXED sequence, which means that its generation process is largely restricted by the given protein sequence, and the flow matching process is for increasing the diversity of AlphaFold prediction, rather than modelling the distribution of all possible folds. We've tested the original AlphaFlow code on general protein design and find that it could not properly generate a fold without specifying a proper sequence.

Though not highlighted in the paper, we're actually the first to prove that it's possible to train a model to generate arbitrary protein fold with an SE(3)-invariant loss function, rather an MSE loss used in FrameDiff, Genie and AlphaFold 3 [3,4,5,6]. An SE(3)-invariant loss gives the same supervising signal for a predicted structure under arbitrary SE(3) transformation, while an MSE loss requires the model to generate a structure with a specific global translation and rotation, falsely punishing similar structure under different SE(3) transformation, which we believe is suboptimal. AlphaFlow partially proves this, but some issues related to their derivation leads to unsatisfactory results when training a design model without sequence. Our derivation has a much better theorectical ground and leads to satisfactory model performance, proving this is a working route.

Furthermore, we enhanced our evaluation according to your suggestion by adding an alternative in silico metric for binding affinity based on minimum interchain pAE predictions to the appendix, inspired by AlphaProteo [7]. During the review process, several AlphaFold 3 replications, including the original models, were released. Though we're not abel to run the AlphaFold 3 model locally due to a licensing issue, we developed an alternative in silico metric based on Chai-1 [8], calculating the minimum value across all interchain terms in the PAE matrix (min_ipAE, lower is better). Note that this metric is proved to be a good indicator for protein-protein binder design, but not validated on small molecule-protein binders. We've added the results of AtomFlow-N and RFDiffAA to the appendix. The results shows that the AtomFlow-generated binders have lower min_ipAE than the ones from RFDiffAA, and both models have the ability to generate binders with similar min_ipAE as natural ones. We're working to develop better metrics as in silico proxy for ligand-binding protein design based on wet lab verification.

Thanks for the detailed response and updated experiments. My primary concern is that, with only 4 ligands and a minor tweak in problem formulation from RFDiff-AA, the paper does not sufficiently demonstrate a value-add on top of RFDiff-AA. It should be noted that 4 experimental validations are very different from 4 in silico successes, as experimental success is what actually matters and in silico metrics are just a noisy approximation for that. This is not to say I am not sympathetic to the authors' claims; quite to the contrary, I believe they may very well have a model on their hands that is better than RFDiff-AA. However, with the current limited evaluations, statistical significance is hard to judge, and publication may be premature.

Thank you for your thoughtful feedback! We truly appreciate your recognition of our model's contributions and are dedicated to thoroughly addressing your concerns in the following sections.

1. Evaluation on Additional Dataset

We emphasize that the initial draft of our paper already included an evaluation on an additional dataset of 20 ligands in the appendix, with atom counts ranging from 21 to 104. During the rebuttal stage, we expanded our evaluation metrics to include iPAE for binding affinity, geometric bond distribution, Ramachandran plot, and chemical validity metrics including SA, QED, and LogP. These additional evaluations further strengthened our results. The extended results corroborate the trends in the main text (as highlighted by Reviewer IAfd) and further validate the robustness of AtomFlow’s outputs.

Note: Extending experiments is a resource-intensive process, as a comprehensive generation for different lengths is required for each ligand, making the process significantly more demanding than evaluating unconditional design models. We are working towards incorporating even more comprehensive supplementary evaluations in the final version.

We acknowledge that our paper and previous rebuttal may not have sufficiently emphasized the additional evaluation in the appendix, and we hope that our response addresses this gap clearly.

2. Future Work on Laboratory Validation

We agree that laboratory success is the ultimate benchmark for a design model. While this is beyond the scope of the current study, it is a key focus for our future work.

3. On RFDiffusionAA & Our Contribution to the ML Community

We recognize the exceptional utility of RFDiffusionAA in protein design, and the primary goal of this work is not to replace RFDiffusionAA.

From a machine learning perspective, AtomFlow introduces a user-friendly tool with a solid theoretical foundation in the SE(3)-invariant quotient space of conformer structures for designing binders to entirely new targets with unknown ligand conformers (as supported by Reviewer Z34b). AtomFlow incorporates a unified, novel biotoken-based atomic flow matching framework for designing flexible interactions. Additionally, we are the first to demonstrate the effectiveness of SE(3)-invariant flow in de novo design.

Inspired by recent advances in protein-protein binder design (e.g., AlphaProteo), we introduced the iPAE prediction of AlphaFold3-like models as part of our binding affinity evaluation. During the rebuttal stage, we significantly improved our supplementary experiments, providing insights for future machine learning researchers and further enhancing our contributions. These aspects mark AtomFlow as a complementary and valuable tool for the community.

We hope this response clarifies our contributions and addresses your concerns comprehensively. Thank you once again for your insights.