Design of Ligand-Binding Proteins with Atomic Flow Matching

## Abstract

Abstract Designing novel proteins that bind to small molecules is a long-standing challenge in computational biology, with applications in developing catalysts, biosensors, and more.
Current computational methods rely on the assumption that the binding pose of the target molecule is known, which is not always feasible, as conformations of novel targets are often unknown and tend to change upon binding.
In this work, we formulate proteins and molecules as unified biotokens, and present AtomFlow , a novel deep generative model under the flow-matching framework for the design of ligand-binding proteins from the 2D target molecular graph alone.
Operating on representative atoms of biotokens, AtomFlow captures the flexibility of ligands and generates ligand conformations and protein backbone structures iteratively.
We consider the multi-scale nature of biotokens and demonstrate that AtomFlow can be effectively trained on a subset of structures from the Protein Data Bank, by matching flow vector field using an SE(3) equivariant structure prediction network.
Experimental results show that our method can generate high fidelity ligand-binding proteins and achieve performance comparable to the state-of-the-art model RFDiffusionAA , while not requiring bound ligand structures.
As a general framework, AtomFlow holds the potential to be applied to various biomolecule generation tasks in the future.

## 1 Introduction

Proteins are indispensable macromolecules that drive the essential processes of living organisms.
A crucial mechanism by which they accomplish this is through binding with small molecules (Schreier et al., ).
Continuous progress has been made to design ligand-binding proteins with various biological functions, such as catalysts and biosensors (Bennett et al., ).
However, the problem remains challenging due to the complex interactions between proteins and molecules, as well as the inherent flexibility of ligands.
The most well-established approaches depend on shape complementarity to dock molecules onto native protein scaffold structures (Bick et al., ; Polizzi & DeGrado, ), which are computationally expensive.

Recently, RFDiffusionAA (Krishna et al., ), a de novo protein design method based on the all-atom structure prediction model RoseTTAFoldAA (Krishna et al., ), has shown remarkable performance in designing novel ligand-binding proteins for small molecules.
This method explicitly captures the interactions between proteins and molecules, achieving superior performance compared to its predecessor RFDiffusion (Watson et al., ), which can only model interactions between amino acid residues.
Despite their great potential for ligand-binding protein design, current approaches assume that the bound conformation of the target molecule is known and rigid.
However, the binding pose of the target molecule is not always available, especially for molecules that do not bind to any known natural proteins (Bick et al., ).
While it is possible to mitigate this limitation by sampling a diverse set of conformers and subsequently filtering them using expert knowledge (Krishna et al., ), this approach demands potentially prohibitive computational resources.
Additionally, the constraint of ligand rigidity is suboptimal, as ligands often undergo significant conformation changes upon binding with proteins (Mobley & Dill, ).
We illustrate this phenomenon in Figure.1.
Some pioneering efforts have been made to account for ligand flexibility (Zhang et al., ; Stark et al., ), however, these methods can only design the portions of proteins that directly interact with the ligands and require the rest part of the proteins as input.

Figure: Figure 1: The conformer of OQO deforms upon binding to coagulation factor XIa. Green: ideal conformer. Orange: bound conformer.
Refer to caption: x1.png

To address the aforementioned issues, we present Atomic Flow-matching (AtomFlow),
a novel deep generative model grounded in the
flow-matching framework (Lipman et al., ) for the design of ligand-binding proteins from 2D molecular graphs alone.
AtomFlow considers ligand flexibility and can iteratively generate ligand conformations and bound protein backbone structures. Instead of relying on a fixed ligand conformer, AtomFlow learns to update the ligand structure along with the structure of the protein binder.
Inspired by recent advances in all-atom structure modeling (Krishna et al., ; Abramson et al., ), we conceptualize proteins and molecules as biotokens with representative atoms, which are associated with various type-specific attributes and can be modeled by a single, unified network.
This approach maximizes the information aggregation between different molecular types, while encouraging the model to focus on the key interaction patterns.
Following the rectified flow approach (Liu et al., ) for generative modelling, we define a flow on the representative atoms as a linear interpolation between the bound protein-ligand complex structures and noisy structures.
We demonstrate that, with minor approximations, the vector field of the defined flow can be effectively learned using an SE(3)-equivariant structure prediction module and a variant of Frame Aligned Point Error (FAPE) loss (Jumper et al., ) that compensates for the multi-scale nature of their geometric features(^1^11The size of a protein is often much larger than that of a molecule. The size disparity should be considered when designing flow-matching models for stable training and inference.).
After training, protein-ligand complex structures can be sampled from the approximated vector field, which iteratively transforms and refines noisy structures based on 2D molecular graphs.
The idea of regressing the vector field using a structure prediction module is also explored in a concurrent work (Jing et al., ), but their focus is on protein structure prediction.
Notably, as a general generative model operating on biotokens, AtomFlow is versatile for different molecular types and has the potential to be applied to various biomolecule generation tasks.

We follow the *in silico* evaluation pipeline of the state-of-the-art method RFDiffusionAA, evaluating AtomFlow on several key metrics including self-consistency, binding affinity, diversity and novelty. AtomFlow matches the overall performance of RFDiffusionAA and demonstrates advantages in various situations. An ablation study further highlights that when the bound structure is unknown, AtomFlow successfully designs protein binders with high binding affinity, whereas RFDiffusionAA can be constrained by its dependence on a fixed, suboptimal ligand structure.

## 2 Related Work

Ligand-binding Protein Design.
Traditional approaches to ligand-binding protein design mainly rely on docking molecules onto large sets of shape-complementary protein pockets (Polizzi & DeGrado, ; Lu et al., ). While the screening process can be accelerated with deep learning models (An et al., ), conventional methods are computationally expensive and often depend on domain experts (Bick et al., ).
Recent advances in deep generative models have paved the way for data-driven approaches, and a variety of models have been proposed to design proteins conditioned on binding targets (Shi et al., ; Kong et al., ; Watson et al., ; Zhang et al., ).
Focusing on molecule binder design, RFDiffusion (Watson et al., ) generates novel proteins from scratch, using a heuristic attractive-repulsive potential to measure shape complementarity.
The follow-up work RFDiffusionAA (Krishna et al., ) improves the performance by explicitly modeling the interactions between proteins and molecules with an all-atom formulation.
These approaches assume binding poses of ligands are known and impose rigidity constraints on ligand structures.
Another line of research focuses on designing binding pockets for small molecules (Stark et al., ; Zhang et al., ). While taking ligand flexibility into consideration, they can only design the portions of proteins that interact with the ligands and require the rest part of the proteins as input.
Our model also accounts for the ligand flexibility, but is able to design full ligand-binding proteins from 2D molecular graph alone.

Protein Generative Model and Structure Prediction.
Recently, various deep generative models for protein generation have emerged (Ingraham et al., ; Lin & AlQuraishi, ; Yim et al., ; ; Wu et al., ; Watson et al., ; Krishna et al., ).
For example, Genie (Lin & AlQuraishi, ) introduces a diffusion process defined on C$\alpha$ coordinates of proteins and allows for the incorporation of motif structures as conditions.
FrameDiff (Yim et al., ) takes a step further by generating novel protein backbone structures using an SE(3) diffusion process applied to residue frames.
Its successor, FrameFlow (Yim et al., ), accelerates the generation process by leveraging the flow-matching framework.
However, these approaches are tailored for single-chain protein generation and fall short in modeling multiple biomolecules.
In contrast, we treat multiple biomolecules, e.g., proteins and molecules, as biotokens and define a novel flow-matching model on their representative atoms.
This allows us to design ligand-binding proteins based solely on molecular graphs, effectively capturing the flexibility of biomolecules and the intricate interactions between them.
Our work is also related to approaches that perform protein structure predictions within the all-atom framework, such as RoseTTAFoldAA (Krishna et al., ) and AlphaFold 3 (Abramson et al., ).
These methods tokenize various types of biomolecules into unified tokens, aiming to develop a universal structure prediction model for all molecular types presented in the Protein Data Bank.
Our AtomFlow adopts the same practice, and we believe this formulation can maximize the information flow between proteins and molecules (Bryant et al., ), while our structural modelling on the representative atoms encourages the model to focus on the key patterns of biointeractions.

## 3 Preliminaries

### 3.1 Notations and Problem Formulation

Notations. In this work, a protein-ligand complex is represented as a series of $N$ biotokens
$\{a_{i}\mid a_{i}=(s_{i},x_{i}),i=1,2,\dots,N\}$, where each token $a_{i}$ corresponds
to either a protein residue or a ligand atom, $s_{i}$ denotes the token type, and $x_{i}\in\mathbb{R}^{3}$
denotes the token position, i.e. the coordinate of its representative atom. Let $\mathcal{S}_{\text{protein}}$ and $\mathcal{S}_{\text{atom}}$ be the set of amino acid types and chemical elements, respectively. For protein residues, $s_{i}\in\mathcal{S}_{\text{protein}}$, with $x_{i}$ being the position of the C-$\alpha$ carbon.
For ligand atoms, $s_{i}\in\mathcal{S}_{\text{atom}}$,
with $x_{i}$ being the atomic position. We define the protein token set as
$\mathcal{P}=\{a_{i}\mid s_{i}\in\mathcal{S}_{\text{protein}}\}$, with $N_{p}=|\mathcal{P}|$
being the number of protein residues, and the ligand token set as
$\mathcal{M}=\{a_{i}\mid s_{i}\in\mathcal{S}_{\text{atom}}\}$, with $N_{m}=|\mathcal{M}|$
representing the number of ligand atoms. In our settings, $N=N_{p}+N_{m}$. The biotokens are attributed with token-level features $f^{\text{token}}\in\mathbb{R}^{N\times c_{t}}$ and pair-level features $f^{\text{pair}}\in\mathbb{R}^{N\times N\times c_{p}}$, where $c_{t}$ and $c_{p}$ denote the feature dimensions.

Problem Formulation. Given a ligand molecule represented as a chemical graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ and a residue count $N_{p}$ for the protein binder to be designed, we aim to generate a protein-ligand complex, where a conformer of $\mathcal{G}$ is docked to a protein binder with $N_{p}$ residues. Specifically, by describing the target protein-ligand complex as a series of biotokens, we generate the token positions $\{x_{i}\}$, with $\mathbf{x}_{m}=\{x_{i}\mid a_{i}\in\mathcal{M}\}$ being a valid conformer for $\mathcal{G}$, and $\mathbf{x}_{p}=\{x_{i}\mid a_{i}\in\mathcal{P}\}$ being a protein binder with high binding affinity to $\mathbf{x}_{m}$. Following previous works (Krishna et al., ; Yim et al., ), we additionally generate the token frames $\{T_{i}=(r_{i},t_{i})\mid a_{i}\in\mathcal{P}\}$ for protein tokens as described in Appendix A.1, which can be used to recover full backbone coordinates of residues. The design of residue types $\{s_{i}\mid a_{i}\in\mathcal{P}\}$ is delegated to an existing reverse folding model (Dauparas et al., ).

### 3.2 Flow Matching

Building upon the significant success of diffusion models in various generative tasks, flow matching models (Albergo & Vanden-Eijnden, ; Liu et al., ) allows for faster and more reliable sampling from a distribution learnt from data. The generative process of flow matching models is usually defined by a probability path $p_{t}({x}),t\in[0,1]$ that gradually transforms from a known noisy distribution $p_{0}({x})=q(x)$, such as $\mathcal{N}({x}|0,I)$ for ${x}\in\mathbb{R}$, to an approximate data distribution $p_{1}\approx p_{\mathrm{data}}(x)$. A vector field $u_{t}({x})$, which leads to an ODE $\frac{\mathrm{d}\phi_{t}(\mathbf{x})}{\mathrm{d}t}=u_{t}(\phi_{t}(\mathbf{x}))$, is used to generate the probability path via the push-forward equation,

$$ $p_{t}=[\phi_{t}]_{*}p_{0}=p_{0}(\phi_{t}^{-1}(x))\mathrm{det}\left[\frac{ \partial\phi_{t}^{-1}}{\partial x}(x)\right],$ (1) $$

which could be approximated with a trainable network $\hat{v}_{t}(x;\theta)$.

Due to the complexity of defining an appropriate $p_{t}$ and $u_{t}$ , we could alternatively define a conditional probability path $p_{t}(x|x_{1})$, which is usually derived through a conditional vector field $u_{t}(x|x_{1})$ for each data point $x_{1}$ (Lipman et al., ). The conditional vector field is then approximated with a trainable network $\hat{v}_{t}(x;\theta)$. Lipman et al. () has proved that the conditional flow matching loss,

$$ $\mathcal{L}_{\mathrm{CFM}}(\theta)=\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}( x|x_{1})}\|\hat{v}_{t}(x;\theta)-u_{t}(x|x_{1})\|,$ (2) $$

has identical gradients w.r.t. $\theta$ with $\mathcal{L}_{\mathrm{FM}}=\mathbb{E}_{t,p_{\mathrm{data}}(x)}||\hat{v}_{t}(x;
\theta)-u_{t}(x)||$, which means the model can generate a marginal vector field by simply learning from the $x_{1}$-conditioned vector fields, without access to $p_{t}(x)$ and $u_{t}(x)$. After training, a neural ODE is obtained, ready for sampling from $p_{0}$ to $p_{t}$ by an ODE solver (Jardine, ) .

## 4 Method

AtomFlow adopts a unified biotoken representation to generate the protein binder and ligand structure by learning the joint distribution of the token positions conditioned on a ligand chemical graph $\mathcal{G}$, $p(\{x_{i}\}|\mathcal{G})$, from known structures of proteins and protein-ligand complexes. To achieve this, we define a rectified flow on the space of all token positions $\mathbf{x}\in\mathbb{R}^{N\times 3}$, and the corresponding vector field is approximated with an SE(3)-equivariant structure prediction module. The structure predicted at the last generation step is adopted as the final result. In this section, we introduce the flow matching model in Section 4.1, the biotoken feature representation in Section 4.2, the structure prediction module in Section 4.3, and the training and inference procedures in Section 4.4. The overview of our method is illustrated in Figure 2.

Figure: Figure 2: The inference process of AtomFlow. We represent the protein-ligand complex as a series of biotokens and embed their token and pair level features. Starting from a noisy sample, the flow matching procedure gradually generates the designed structure $x_{1}$ with a structure prediction network.
Refer to caption: x2.png

### 4.1 Flow Matching for Protein-Ligand Complex Generation

We jointly design the complex structure $\mathbf{x}=\mathbf{x}_{m}\cup\mathbf{x}_{p}$, which lies in the space of $\mathbb{R}^{N\times 3}$, with a flow matching model. Considering that different structures obtained under arbitrary SE(3) transformations correspond to the same complex, we treat each structure as an element in the quotient space $\mathcal{Q}:\mathbb{R}^{N\times 3}/\text{SE}(3)$, where two structures are identical if they could be perfectly aligned with an SE(3) transformation (Jing et al., ). This quotient space is proved to be a Riemannian manifold when defined with suitable care (Diepeveen et al., ).

Following Riemannian Flow Matching (Chen & Lipman, ), we define a rectified flow on this manifold with a premetric $d:\mathcal{Q}\times\mathcal{Q}\rightarrow\mathbb{R}$. We denote $\text{align}_{x}(y)$ for $x,y\in\mathbb{R}^{N\times 3}$ as aligning structure $y$ to $x$ to minimize RMSD, then the premetric $d(x,y)$ could be defined as the minimum point-wise root mean square deviation (RMSD) among all pairs of possible structures in the original space $\mathbb{R}^{N\times 3}$ for two elements in the quotient space

$$ $d(x,y)=\min_{\tau\in\mathrm{SE}(3)}\text{RMSD}\,(\tau(y),x)=\text{RMSD}\,( \text{align}_{x}(y)-x)$ (3) $$

###### Proposition 1 .

The premetric in equation 3 is a qualified premetric on $Q$.

With such premetric at hand, we could obtain a well-defined conditional vector field that decreases the premetric linearly from the prior distribution to the data distribution

$$ $u_{t}(x|x_{1})=\dfrac{1}{1-t}\left(\text{align}_{x}(x_{1})-x\right).$ (4) $$

We leave the proof of Proposition 1 and the derivation of equation 4 to Appendix A.3. Since the vector field is defined as a function of $\mathbf{x}_{1}$, we could learn the vector field with a structure prediction model $\hat{\mathbf{x}}_{1}(\mathbf{x},t;\theta)$. By substituting equations 4 into equation 2, we obtain the training loss

$$ $\mathcal{L}_{\mathrm{CFM}}(\theta)=\mathbb{E}_{t,p_{\text{data}}(\mathbf{x}_{1 }),p_{t}(\mathbf{x}|\mathbf{x}_{1})}\left\|\dfrac{1}{1-t}(\text{align}_{ \mathbf{x}}(\hat{\mathbf{x}}_{1}(x,t;\theta))-\text{align}_{\mathbf{x}}( \mathbf{x}_{1}))\right\|,$ (5) $$

This loss calculates the $(1-t)$-normalized distance between predicted $\hat{\mathbf{x}}_{1}$ and $\mathbf{x}_{1}$ in the data distribution aligned to the noisy structure of current step, which is SE(3)-equivariant to both the predicted and ground truth structure. The structure module is designed to predict the token frames (Section 3), while the token positions are extracted from them during the generation process. The last prediction output is adopted as the final result.

Defining a unified flow matching procedure on the joint distribution enables the model to directly learn the structure characteristics that leads to a tightly binded complex, as well as the conformation deformation of both the proteins and the ligands, which is essential to designing a satisfactory ligand-binding protein.

###### Proposition 1 .

### 4.2 Representation of Conditional Features

The generation process of AtomFlow is conditioned on the ligand chemical graph $\mathcal{G}$ and a designated protein length $N_{p}$. We model such conditions as an additional condition to the vector field $u$. As a result, the inputs of the prediction network $\hat{\mathbf{x}}_{1}$ is augmented to accept conditional features. With the biotoken representation, we embed all such features as $f^{\text{token}}$ and $f^{\text{pair}}$ as illustrated in Figure 2A.

For a ligand chemical graph $\mathcal{G}$, we embed the chemical element, as well as other known chemical properties as $f^{\text{token}}$ of ligand tokens. The chemical bonds $\mathcal{E}$ are embeded in $f^{\text{pair}}$ as a multi-dimensional adjacency tensor, each dimension representing a bond type. We also embed the relative residue position (Shaw et al., ) as a pair feature, while the residue tokens may also be attributed with other known conditions. We concatenate the protein and ligand features to form a unified feature tensor, eliminating the need to distinguish different types of token when processing the features.

### 4.3 Structure Prediction Network

The structure prediction network $\hat{\mathbf{x}}_{1}(\mathbf{x},t;\theta)$ (^2^22 Though $\hat{\mathbf{x}}_{1}$ is a function of $\mathbf{x},t,f^{\text{token}},f^{feat}$, we omitted certain parameters to simplify the text.) predicts the token frames $\{T_{i}\}$, which can be used to extract token positions $\mathbf{x}_{1}$, given a series of noisy positions $\mathbf{x}$ at timestamp $t$. It encodes $\mathbf{x}$, along with $f^{\text{token}}$ and $f^{\text{pair}}$, with an SE(3) invariant encoding module, processing the representation with a transformer stack, and generates the predicted structure with a structure module based on invariant-point attention (IPA) (Jumper et al., ), as illustrated in Figure 2B. The network jointly processes two kinds of biotokens, protein residues and ligand atoms, with different spatial scales, and handles such difference with special care.

Distance Map. The input coordinates $\mathbf{x}$ are encoded by projecting the one-hot binned distance map between input coordinates for each token pair to the feature space

$$ $t_{i,j}=\text{Linear}(\text{BinRepr}(\|\mathbf{x}^{(i)}-\mathbf{x}^{(j)}\|)),$ (6) $$

where the bins are not divided equally considering different precision requirement between residues and atoms. This representation is SE(3) invariant, since the internal distance does not change under rigid transformation. (^3^33 To accommodate the precision differences between ligands and proteins, the bin intervals are dense between 1Å (approximate length of a chemical bond) and 3.25Å (approximate distance between adjacent amino acids) and sparser beyond 3.25Å.)

Feature Embedder. The feature embedder generates a single representation $s\in{\mathbb{R}}^{N\times c_{s}}$ and pair representation $z\in{\mathbb{R}}^{N\times N\times c_{z}}$ from distance map $h$, noise level $t$, $f^{\text{token}}$ and $f^{\text{pair}}$ for the following steps. The noise level is encoded with Gaussian Fourier embedding (Song et al., ). The local features are concatenated and projected to single representation $s$ and pair representation $z$, $s_{i}=\text{Linear}(f^{\text{local}}_{i})$.
The pair features and input encoding are projected and added to $z$

$$ $z_{i,j}=\text{Linear}(f^{\text{local}}_{i})+\text{Linear}(f^{\text{local}}_{j} )+f^{\text{pair}}_{i,j}+t_{i,j}.$ (7) $$

As described in Section 4.2, different token types can be treated the same and processed uniformly.

Structure Module. The structure module generates a predicted complex structure, represented as a series of token frames $T^{N}$. For ligand atoms, the rotation of the predicted frame is always identity rotation, while the translation equals to its position. It first processes $z$ through a deep transformer stack (Appendix A.4) to obtain a denoised pair representation $z^{\prime}$, and converts $s$ and $z^{\prime}$ to $T^{N}$ through a series of shared-weight IPA block

$$ $T_{1\cdots N}=\text{IPAStack}(s_{1\cdots N},\text{TransformerStack}(z_{1\cdots N ,1\cdots N})).$ (8) $$

The IPA stack outputs a sequence of frames for each token, while the rotations for atom tokens are dropped and replaced with the atom frame demonstrated in Section 4.2. The final output represents the full complex structure $\hat{\mathbf{x}}$, while token positions $\hat{x_{1}}$ is calculated as previously described. The Transformer stack on the unified token sequence allows us to smoothly model the interactions between different types of biological entities in a joint feature space, while the IPA blocks are proved to be efficient when the final structure is properly embedded in the transformer output (Jumper et al., ).

Auxiliary Head. We add an auxiliary head to predict the pairwise binned distance from the denoised pair representation $z^{\prime}$, $h_{i}=\text{softmax}(\text{Linear}(z^{\prime}_{i}))$, which directly supervise the input of structure module and has been proved to be helpful during training (Jumper et al., ). The bins are also unevenly divided to accommodate the multi-scale characteristics of the predicted complex.

### 4.4 Training and Inference

We train the network $\hat{\mathbf{x}}_{1}$ by sampling data points and timestamps, calculating the noisy input, and supervising the predicted results. At inference time, we transforms the token positions sampled from the prior distribution through the predicted vector field with an ODE solver, and outputs the structure we obtained at the final step.

Loss. We supervise the predicted complex structure $T$ with a metric that measures the structural difference between the observed structure and the predicted structure. Preliminary experiments show that the $\mathcal{L}_{\text{CFM}}$ in equation 5 leads to a fluctuating training trajectory since the aligning object $\mathbf{x}$ varies upon training. With approximation (Appendix A.4), we replace the loss function to a variant of the widely-adopted FAPE function (Jumper et al., ),

$$ $\mathcal{L}_{\text{CFM-FAPE}}(\theta)=\mathbb{E}_{t,p_{\text{data}}(\mathbf{x} _{1}),p_{t}(\mathbf{x}|\mathbf{x}_{1})}\left[\dfrac{1}{1-t}\mathrm{FAPE}(\hat{ \mathbf{x}}_{1}(\mathbf{x},t;\theta),\mathbf{x}_{1})\right].$ (9) $$

We show that this substitution does not change the training objective in the appendix. Since the normalization factor $Z$ in the FAPE loss is related to the numerical range of distance, we divide the FAPE loss into protein-protein interaction, protein-ligand interaction, ligand-ligand interaction, and assign different $Z$s for the three parts. For the auxiliary head, we adopt the cross-entropy loss averaged over all token pairs for the predicted distance. The final training loss

$$ $\mathcal{L}=\alpha_{1}\mathcal{L}_{\text{CFM-FAPE-pp}}+\alpha_{2}\mathcal{L}_{ \text{CFM-FAPE-pl}}+\alpha_{3}\mathcal{L}_{\text{CFM-FAPE-ll}}+\alpha_{4} \mathcal{L}_{\text{aux}}.$ (10) $$

Further details are elaborated in Appendix A.4.

Training. We sample the timestamp $t$ from the logit normal distribution, assigning more weight on intermediate steps, which helps the model to achieve better performance on hard timestamps (Esser et al., ; Karras et al., ). The prior distribution $q(x)$ is selected as $\mathcal{N}(0,\sigma_{\text{data}})$, where $\sigma_{\text{data}}=10$. The input $\mathbf{x}$ is given by interpolating the data point and a sample from the prior distribution. The training procedure is shown in Algorithm 1.

Figure: Algorithm 1 Training

Inference. A scheduler of noise levels $\{t_{i}\}_{i=0}^{m},t_{0}=0,t_{m}=1$ is used to determine the noise level $t_{i}$ of each sampling step $x_{t_{i}}$. Starting from a noisy sample $x_{t_{i}}=x_{0}$ as the initial model input, the structure prediction network predicts the vector field, which gives $x_{t_{i+1}}$ with the Euler’s Method, i.e.

$$ $\mathbf{x}_{t_{i+1}}=\mathbf{x}_{t_{i}}+\dfrac{t_{i+1}-t_{i}}{1-t_{i}}\bigg{(} \text{align}_{\mathbf{x}_{t_{i}}}\Big{(}\text{Extract}\big{(}\hat{\mathbf{x}}_ {1}(\mathbf{x}_{t_{i}},t_{i};\theta)\big{)}\Big{)}-\mathbf{x}_{t_{i}}\bigg{)},$ (11) $$

where the Extract function extracts the token positions from the predicted token frames. The model output at the last step is adopted as the final result. The inference procedure is shown in Algorithm 2.

## 5 Experiments

Following previous protein design models (Yim et al., ; Lin & AlQuraishi, ; Watson et al., ) and binder design models (Krishna et al., ), we evaluate AtomFlow through in silico experiments on key metrics of our generated binder including self-consistency, binding affinity, diversity and novelty.

### 5.1 Experiment Setup

Training Data. We train the denoising model on two datasets: PDBBind (Liu et al., ), a protein-ligand conformer dataset derived from the Protetin Data Bank (PDB) (Berman et al., ), and SCOPe (Chandonia et al., ), a structure categorical dataset for protein. The model is firsted trained on solely generating the protein structure for 40k steps, and then finetuned on co-generating both of the protein and ligand structure for 30k steps.

Baseline and Model Variant. We compare AtomFlow with the state-of-the-art binder generation method RFDiffusionAA (Krishna et al., ), which is extensive trained on almost all known data. Since RFDiffusionAA requires a fixed ligand structure at the binding state as input, we extend our method to work under its setting. For AtomFlow, besides the original setting (AtomFlow-N), we also train a version of our model with the pairwise distance matrix of the bound structure as an auxiliary hint input (AtomFlow-H). This version still needs to generate the ligand structure itself, rather than rely on a fixed structure, as other specifications is not modified.

Evaluation Set. We mainly evaluate all methods on a selected ligand set (evaluation set) from RFDiffusionAA (FAD, SAM, IAI, OQO). The evaluation set comprises ligands from inside and outside the training set, with both long and short lengths. We conduct evaluation on an extended ligand set (extended set, see Appendix A.5) to further demonstrate the performance of AtomFlow.

### 5.2 Self-consistency and Conformer Legitimacy

In this section, we evaluate the legitimacy of the generated protein structure by self-consistency RMSD and the predicted ligand structure at binding state by detecting structural violence in the conformer. Legitimacy is crucial in binder design, given that the model output is not guaranteed to be valid, while a design with higher legitimacy is more possible to fold as expected.

Figure: Figure 3: Designed structures for different ligands at different lengths. We align the ESMFold predicted structure to the designed structure, and report the scRMSD metric. Green: designed protein; Orange: designed ligand conformer; Grey: ESMFold predicted protein.
Refer to caption: x3.png

Protein Structure. For protein structures, self-consistency RMSD is widely adopted as a metric to evaluate their legitimacy (Lin & AlQuraishi, ; Watson et al., ), which compares the generated structure and the folding of its sequence predicted by an accurate model. We adopt LigandMPNN (Dauparas et al., ) to predict possible sequences from the generated structures. We first generate 8 sequences for all designed structure with LigandMPNN, then predict the corresponding protein structure with ESMFold (Lin et al., ), and the metric for each generated structure is calculated as the minimum rooted mean squared distance between the designed structure and predicted structure (scRMSD). For each ligand in the evaluation set, we generates 10 structures for lengths in [100, 150, 200, 250, 300]. The results are shown in Table 1. We illustrate several generated samples in Figure 3, and the cumulative distribution of scRMSD among them in Figure 4A and 4D. The results on the extended set are shown in Appendix A.5.

**Table 1: Proportion of samples with scRMSD $<2$ on the evaluation set (higher is better).**
| Method | Overall | SAM | FAD | IAI | OQO |
| --- | --- | --- | --- | --- | --- |
| AtomFlow-H | 0.57 | 0.60 | 0.36 | 0.58 | 0.74 |
| AtomFlow-N | 0.50 | 0.50 | 0.38 | 0.58 | 0.54 |
| RFDiffusionAA | 0.52 | 0.60 | 0.58 | 0.48 | 0.42 |
| RFDiffusion | 0.33 | 0.04 | 0.50 | 0.44 | 0.32 |

AtomFlow and RFDiffusionAA outperforms RFDiffusion on all ligands in the evaluation set, while both AtomFlow-H and AtomFlow-N reach comparable results to RFDiffusionAA, and exhibits advantages over RFDiffusionAA on several cases. The restricted performance of RFDiffusion is as expected since its binding potential for guiding the protein-ligand interaction may lead to structural destruction. Both AtomFlow and RFDiffusionAA models the interaction directly, thus not requiring a strong potential to interfere the generative process, and lead to better generation results. Notably, without relying on structural guidance from the input ligand conformer, AtomFlow-N achieves close performance to AtomFlow-H, thereby successfully augmenting the setting to flexible design.

### 5.3 Binding Affinity

Figure: Figure 4: A: Self-consistency RMSD distribution curve demonstrating the ratio among all designed samples for the evaluation set with $\text{scRMSD}\leq x$ (higher is better). AtomFlow outperforms RFDiffusion with a curve similar to RFDiffusionAA. AtomFlow-H generates achieves the best result among the methods. The ratio of samples with $\text{scRMSD}<2$ is highlighted. B: PoseBusters score distribution of AtomFlow generated samples on the extended set. Most ligand conformers generated by AtomFlow-N only fails $\leq 1$ metric of its evaluations. C: Vina score distribution over all designs on the evaluation set (lower is better). AtomFlow achieves comparable performance to RFDiffusionAA, outperforming RFDiffusion. D: scRMSD curve and Vina energy distribution over designs for each ligand in the evaluation set. AtomFlow outperforms RFDiffusion on all cases and metrics. AtomFlow and RFDiffusionAA each exhibit advantages on different ligands, with comparable overall results.
Refer to caption: x4.png

In this section, we evaluate the binding affinity of the designed protein binder by calculating an energy function for the atom-level interaction between the protein and the ligand. Binding affinity is the key metric to reveal whether the designed binders are able to bind the target molecule. Though the real binding affinity could only be determined through experiments in the wet lab, an energy function is usually adopted as an in silico alternative (Zhang et al., ). We calculate the AutoDock Vina Score (Eberhardt et al., ) for all 8 sequences packed by the Rosetta packer (Leaver-Fay et al., ), and the reported energy for a structure is the minimum score among all packed proteins. We calculate the energy for all generated structures for the selected ligand set in Section 5.2 and compare AtomFlow with RFDiffusionAA and RFDiffusion. The result is illustrated in Figure 4C and 4D.

We find that the binding affinity of RFDiffusion is quite poor since it does not model the protein-ligand interaction directly. AtomFlow has reached comparable binding affinity to RFDiffusionAA, though marginally lower on several cases. We attribute this to the exhaustive training process of RFDiffusionAA on all known data in PDB, while AtomFlow could be further trained and this will be investigated further in our future work. The minimum binding energy generated by AtomFlow-H is slightly higher than that of AtomFlow-N, possibly because the provided conformer hint hinders the model from exploring additional binding states.

Figure: Figure 5: AtomFlow-N designs binders with lower vina energy distribution than RFDiffusionAA on 2GJ without the bound structure. Illustration of one sample for each method with PLIP demonstrates that the AtomFlow-N designed binder has more chemical interaction with the ligand.
Refer to caption: x5.png

We further compare AtomFlow with RFDiffusionAA on a realistic setting where the bound conformer is unknown. We set the target ligand as luminespib (PDB id: 2GJ), an Hsp90 inhibitor (Piotrowska et al., ). A designed protein binder for luminespib may act as a protein drug carrier to enhance drug efficacy. Luminespib is a molecule ligand with 33 heavy atoms, so that the conformer is quite flexible when docked to different receptors. We design 10 binders for luminespib using AtomFlow and RFDiffusionAA. The ideal conformer from PDB is provided to RFDiffusionAA, while no conformer is provided to AtomFlow. The binding energy of the designed structures and one designed sample with PLIP (Adasme et al., ) to demonstrate the protein-ligand interaction are illustrated in Figure 5. It is shown that AtomFlow generates more binders with higher binding affinity than RFDiffusionAA, and significantly outperforms RFDiffusionAA on the lowest energy among all generated structures. This demonstrate that a proper bound structure is crucial to the performance of RFDiffusionAA, while AtomFlow does not rely on such structure and generates proper conformers by co-modelling the structure space of proteins and ligands.

### 5.4 Diversity and Novelty

In this section, we report the diversity and novelty of AtomFlow, following common practice in literature (Krishna et al., ; Yim et al., ). Diversity refers to the structural divergence of the designed binders for a certain ligand, while novelty refers to how close a designed protein is to the known proteins. For diversity, we generate 100 structures with 200 residues for each ligand, and then use MaxCluster (Herbert, ) to calculate pairwise structural distance of the outputs and report the number of clusters using different thresholds of maximum distance within cluster. For novelty, we generate 4 structures with residue count in $[100,101,\cdots,300]$ for each ligand, and then calculate the highest TM-score (Zhang, ) between a designed structure and any similar structure searched by FoldSeek (Kempen et al., ) (pdbTM), as well as the protein scRMSD. The search range of pdbTM is all known protein structures in PDB.

Figure: Figure 6: A: Cluster count based on different thresholds for maximum difference within cluster for each ligand in the evaluation set. AtomFlow generates diverse binder folds for all ligands, not restricted to the existing binder structure. B: Scatter plot of designability (scRMSD) vs. novelty (pdbTM) for ligands in the evaluation set. AtomFlow successfully designs self-consistent structures with high pdbTM, demonstrating high novelty.
Refer to caption: x6.png

Figure 6A shows that the structures generated by AtomFlow is quite diverse for all four ligands, and the diversity varies among different ligands. Though existing protein-ligand complexes only provides limited folds for possible binders, by adding protein-only data to the training set, our model successfully learns from the protein structure distribution to generate more possible folds, instead of replicating known patterns. The scatter plot of scRMSD vs. pdbTM shown in Figure 6B reveals that AtomFlow has the ability to generate structures that are quite different from existing proteins with acceptable designability. Note that most designable structures are still similar to known ones, which is as expected since most protein folds are already discovered, while novel folds are quite sparse and hard to derive.

## 6 Conclusion and Future Work

In this work, we proposed AtomFlow, a de novo protein binder design method for small molecule ligands considering the flexibility of ligand structure. Unlike previous works, AtomFlow no longer relies on a given bound ligand conformer as input. We represent the protein-ligand complex as unified biotokens, learning the structure distribution of both the proteins and the ligands simultaneously from the data with an SE(3)-equivariant flow matching model on the representative atoms. During evaluation, AtomFlow shows comparable design quality to the state-of-the-art model RFDiffusionAA, which requires the ligand conformer to be fixed before design. Further evaluation exhibits the advantage of AtomFlow at the circumstance when the ligand conformer is not known. A direct future work is to support more precise control of the generated structures, and we’re working to migrate AtomFlow to all kinds of biomolecules, including DNA, RNA, and metal ions.

## Appendix A Appendix

### A.1 Protein Frames

Proteins are composed of amino acid chains linked by peptide bonds, forming a backbone with protruding side chains. Each amino acid’s position and orientation is described by a local coordinate system, or protein frame, centered on three key backbone atoms: the alpha carbon (C$\alpha$), the carbonyl carbon (C), and the amide nitrogen (N). These atoms act as reference points for establishing the frame. The alpha carbon (C$\alpha$) typically acts as the origin. The vector from C$\alpha$ to the amide nitrogen (N) is normalized to define one axis of the frame. A second axis is defined by the normalized vector from C$\alpha$ to the carbonyl carbon (C). The third axis is formed by the cross product of these two vectors, creating an orthogonal, right-handed coordinate system. The residue frame is typically represented as an SE(3) transformation $T=(R,t)$, which maps a vector from this local system to the global coordinate system. In this transformation, $t$ corresponds to the position of C$\alpha$ in the global system, and $R$ represents the rotation needed to align the residue’s structure within the global context.

Figure: Figure 7: A protein frame illustration. The C$\alpha$, C, N atoms form a panel, which is the xy panel. The x axis is defined as the orientation from C$\alpha$ to N, while the y axis is on the panel and perpendicular to the x axis. The z axis is perpendicular to the xy panel.
Refer to caption: x7.png

### A.2 Details on Biotokens

#### Token Features

For ligand atom tokens, the token-level feature set includes: chirality, degree, formal charge, implicit valence, number of H atoms, number of radical electrons, orbital hybridisation, aromaticity, ring size. The pair-level feature is provided as one-hot embedding of the bond type. For residue tokens, no token-level feature is known, while the pair-level feature only contains the binned distance of residue index between residues. All features are encoded as a one-hot vector and concatenated.

#### Token Frames

The final loss we adopted $\mathcal{L}_{\text{CFM-FAPE}}$ requires aligning the predicted structure to the local frame of every token. The frames of protein residues can be naturally defined as in Section 3. However, the frames of ligand atoms could not be choosed directly. Since a frame could be calculated from the coordinate of 3 atoms, we need to choose an atom triplet for every atom token.

We first obtain an canonical rank of every atom that does not depend on the input order (Schneider et al., ). The atoms are then renamed to its rank. For atoms $x$ with a degree greater than or equal to 2, we select the lexicographically smallest triplet $(u,x,v)$ to define the frame, where $u$ and $v$ are neighbors of $x$ . For atoms with a degree of 1, $u$ is the only neighbor of $x$ , and $v$ is chosen as one of $u$ ’s neighbors. This method ensures that each atom’s frame is defined in a consistent manner, irrespective of its position in the input sequence, thereby facilitating the model to learn a consistent structural target.

#### Extending Token Types and Features

Though AtomFlow only considers the interaction between protein and molecule ligands, the unified biotoken has the potential to extend to all biological entities, including DNA, RNA, etc, by defining the token position, token frame, local and pair features, and the representation of the internal structure. For example, an RNA can be represented as a sequence of nucleotides, with the token position defined as its mass center, and the token frame calculated from an atom triplet, such as C2-N1-C6.

The token features can also be extended to support more types of known information. For example, the local features could also contain an embedding to indicate preferred secondary-structure, or whether a ligand atom is required to be closer to the designed protein; the pair features could also contain the motif information with a distance map.

### A.3 Details on the Flow Matching Process

For all types of tokens, we only consider their token positions to simplify the flow matching process. Thus, the positions of all tokens lies in the euclidean space $\mathbb{R}^{N\times 3}$. Since a complex could be arbitrarily moved or rotated in the coordinate space without changing its structure, we need an algorithm that treats different position series as the same if they could be aligned with an SE(3) translation. Thus, every data point we consider now lies in the quotient space $\mathbb{R}^{N\times 3}/\text{SE(3)}$. This quotient space is proved to be a Riemannian manifold (Diepeveen et al., ).

For a Riemannian manifold, the flow matching process could be defined using a premetric (Chen & Lipman, ). A premetric $d:\mathcal{M}\times\mathcal{M}\rightarrow\mathbb{R}$ should satisfy: 1. $d(x,y)\geq 0\text{ for all }x,y\in\mathcal{M}$; 2. $d(x,y)=0\text{ iff }x=y$; 3. $\nabla d(x,y)\neq 0\text{ iff }x\neq y$.

We define our premetric as the minimum point-wise rooted sum of squared distance (RMSD) among all pairs of possible structures in the original space $\mathbb{R}^{N\times 3}$ for two elements in the quotient space $d(x,y)=\|\text{align}_{x}(y)-x\|$, which satisfies all three conditions (Proposition 1).

###### Proof.

Since the premetric is defined as a norm, it satisfies condition 1 by nature. When $x=y$, the best alignment that aligns $y$ to $x$ could derive the exact same position as $x$, yielding an zero norm. When $x\neq y$, when $y$ is aligned to $x$, there’s still structural difference between the structures, thus the premetric is not zero. For condition 3, by defining $y^{\prime}=\text{align}_{x}(y)$, we have

$$ $\nabla d(x,y)=\nabla\sqrt{\sum_{i=1}^{n}(y_{i}^{\prime}-x_{i})^{2}}=\dfrac{y^{ \prime}-x}{||y^{\prime}-x||}=\dfrac{\text{align}_{x}(y)-x}{||\text{align}_{x}( y)-x||}\geq 0.$ (12) $$

Thus $d(x,y)$ satisfies all the conditions as a qualified premetric.
∎

With such premetric, and a monotonically decreasing differentiable scheduler $\kappa(t)=1-t$, we could obtain a well-defined conditional vector field that linearly interpolates between the noisy and real data
(Chen & Lipman, )

$$ $u_{t}(x|x_{1})=\dfrac{\mathrm{d}\log\kappa(t)}{\mathrm{d}t}d(x,x_{1})\dfrac{ \nabla d(x,x_{1})}{\|\nabla d(x,x_{1})\|^{2}}=\dfrac{1}{1-t}(\text{align}_{x}( x_{1})-x).$ (13) $$

The vector field in equation 13 is calculated by substituting equation 12 into the left side.
This vector field provides the direction for moving straight towards $x_{1}$, and generates a probability flow that interpolate linearly between noisy sample $x_{0}$ and data sample $x_{1}$.

Since the vector field is defined as a function of $x_{1}$, we could learn the vector field with a structure prediction model $\hat{x_{1}}(x,t;\theta)$. By substituting equation 4 into equation 2, we obtain the training loss

$$ $\mathcal{L}_{\mathrm{CFM}}(\theta)=\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}( x|x_{1})}\left\|\dfrac{1}{1-t}(\text{align}_{x}(\hat{x_{1}}(x,t;\theta))-\text {align}_{x}(x_{1}))\right\|.$ (14) $$

###### Proof.

### A.4 Details on the Prediction Network

#### Structure Module Specifications

The main components of the structure module is derived from Alphafold 2 (Jumper et al., ), while our implementation builds on top of the widely acknowledged reimplementation OpenFold (Ahdritz et al., ). The TransformerStack consists of 14 layers of simplified Evoformer block, and the IPAStack consists of 4 layers of Invariant Point Attention (IPA) blocks. The MSA operations in the Evoformer block is simplified by replacing the operations on the MSA feature matrix to the single representation $s_{i}$. The weights of the IPA blocks are shared, and the structural loss is calculated on the outputs of each block and averaged.

#### Training Details

During training, we equally sample data from the SCOPe dataset (v2.08) and the PDBBind dataset (2020). We simply drop the data with more than 512 tokens, and we don’t crop the filtered complexes since the cutoff is large enough and only filters out a relative small portion of data. We train our model on 10 NIVIDA RTX 4090 acceleration card, with a batch size set to 10, which means the batch size on each device is set to 1. We use the Adam Optimizer (Kingma, ) with a weight-decaying learning rate scheduler, starting from $10^{-3}$ and decays the learnign rate by $0.95$ every 50k steps. We seperate the training process into two stages: 1) initial training, $\alpha_{1}=0.5,\alpha_{4}=0.3,\alpha_{2}=\alpha_{3}=0$; 2) finetuning, $\alpha_{1}=\alpha_{2}=\alpha_{3}=0.5,\alpha_{4}=0.3$.

#### Loss Function

$\mathcal{L}_{\text{CFM}}$ calculates an aligned RMSD by aligning $\mathbf{x_{1}}$ and $\hat{\mathbf{x_{1}}}$ to $\mathbf{x}$, while the FAPE loss calculates an averaged RMSD by aligning $\mathbf{\hat{x_{1}}}$ to each residue frame of $\mathbf{x_{1}}$, which could be extended to the token frame (Appendix A.2). Let $\text{align}_{x,i}(y)$ denote aligning $y$ to the $i$-th token frame of $x$, we have

$$ $\displaystyle\mathcal{L}_{\text{CFM}}$ $\displaystyle=\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}(x|x_{1})}\left\| \dfrac{1}{1-t}(\text{align}_{x}(\hat{x_{1}}(x,t;\theta))-\text{align}_{x}(x_{1 }))\right\|$ $\displaystyle\approx\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}(x|x_{1})}\left \|\dfrac{1}{1-t}\cdot\dfrac{1}{N}\sum_{i=1}^{N}\left(\text{align}_{x,i}(\hat{x _{1}}(x,t;\theta))-\text{align}_{x,i}(x_{1}))\right)\right\|$ $\displaystyle\approx\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}(x|x_{1})}\left \|\dfrac{1}{1-t}\cdot\dfrac{1}{N}\sum_{i=1}^{N}\left(\text{align}_{x_{1},i}( \hat{x_{1}}(x,t;\theta))-\text{align}_{x_{1},i}(x_{1})\right)\right\|$ $\displaystyle\approx\mathbb{E}_{t,p_{\text{data}}(x_{1}),p_{t}(x|x_{1})}\left \|\dfrac{1}{1-t}\cdot\dfrac{1}{N}\sum_{i=1}^{N}\left(\text{align}_{x_{1},i}( \hat{x_{1}}(x,t;\theta))-x_{1}\right)\right\|$ $\displaystyle=\mathcal{L}_{\text{CFM-FAPE}}$ $$

###### Proposition 2 .

$\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}\iff\mathcal
{L}_{\text{CFM}}=0\iff\mathcal{L}_{\text{CFM-FAPE}}=0$.

###### Proof.

When $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}$, we have $\forall i,\text{align}_{\mathbf{x_{1}},i}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}$. As a result, $\mathcal{L}_{\text{CFM}}=\mathcal{L}_{\text{CFM-FAPE}}=0$. This establishes that:

$$ $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}\iff\mathcal {L}_{\text{CFM}}=0\quad\text{and}\quad\text{align}_{\mathbf{x_{1}}}(\hat{ \mathbf{x_{1}}})=\mathbf{x_{1}}\iff\mathcal{L}_{\text{CFM-FAPE}}=0.$ (15) $$

Now, assume $\mathcal{L}_{\text{CFM}}=0$. Suppose $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})\neq\mathbf{x_{1}}$. Then for all transformations $R$ and $t$, we have $R\hat{\mathbf{x_{1}}}+t\neq\mathbf{x_{1}}$, which implies: $\|\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})-\mathbf{x_{1}}\|\neq 0$, leading to $\mathcal{L}_{\text{CFM}}\neq 0$. This is a contradiction. Therefore, $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}$. This proves that

$$ $\mathcal{L}_{\text{CFM}}=0\iff\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1} }})=\mathbf{x_{1}}.$ (16) $$

Similarly, assume $\mathcal{L}_{\text{CFM-FAPE}}=0$. Suppose $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})\neq\mathbf{x_{1}}$. Then:
$\|\text{align}_{\mathbf{x_{1}},i}(\hat{\mathbf{x_{1}}})-\mathbf{x_{1}}\|\neq 0,$
which leads to $\mathcal{L}_{\text{CFM-FAPE}}\neq 0$, again a contradiction. Therefore, $\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{x_{1}}})=\mathbf{x_{1}}$. This proves that:

$$ $\mathcal{L}_{\text{CFM-FAPE}}=0\iff\text{align}_{\mathbf{x_{1}}}(\hat{\mathbf{ x_{1}}})=\mathbf{x_{1}}.$ (17) $$

The proposition is proved by combining equation 15,16,17.
∎

This means that both $\mathcal{L}_{\text{CFM}}$ and $\mathcal{L}_{\text{CFM-FAPE}}$ provides an optimization direction towards minimizing the SE(3) invariant structural difference between the predicted structure and the ground truth structure.
Thus, we adopt $\mathcal{L}_{\text{CFM-FAPE}}$ as a realistic approximation of $\mathcal{L}_{\text{CFM}}$ and adopt it as the training objection during evaluation.

###### Proposition 2 .

###### Proof.

### A.5 Evaluation Details

#### Specifications

Following RFDiffusionAA, we use FAD, SAM, IAI, OQO as the selected evaluation set. FAD and SAM is witnessed by both models as training data, while IAI and OQO is not and demonstrates the generalization ability. To further investigate the performance of our method, we conduct experiments on an extended set of 20 ligands (ligands from PDB id 6cjs, 6e4c, 6gj6, 5zk7, 6qto, 6i78, 6ggd, 6cjj, 6i67, 6iby, 6nw3, 6o5g, 6hlb, 6efk, 6gga, 6mhd, 6i8m, 6s56, 6tel, and 6ffe). The extended dataset includes ligand sizes (including hydrogen) ranging from 21 to 104 in length.

#### Extended Results

We illustrate the designability (scRMSD) and binding affinity (Vina energy) of AtomFlow-N in Figure 8. The extended evaluation shows that the performance of AtomFlow on the extended set is similar to the evaluation set shown in the main article, and demonstrates that AtomFlow is able to tackle with almost all kinds of ligands.

Figure: Figure 8: A: scRMSD of designs for each ligand in the extended set; B: Vina energy of designs for each ligand in the extended set.
Refer to caption: x8.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: My major question is about the novelty and insufficient experiments. I hope the author can help to address my concerns.
soundness: 2
strengths: S1. This method is based on full-atom level protein and molecule generation,  with both the protein binder and the molecule pose is designed.
S2. Experiments show the effectiveness of the method.
summary: This paper presents ATOMFLOW, a novel deep generative model under the flow-matching framework for the design of ligand-binding proteins from the 2D target molecular graph alone.  It co-design the protein binder structure and the target-molecule flexibility. Experimental shows the method's effectiveness.
weaknesses: W.1 I think the major weakness lies in its novelty. First of all, using Flow-matching with full atom generation in peptide [1] or protein [2] other bio-molecules is not a new idea. Besides, the equivariant generation process [3,4,5] and FAPE loss in AlphaFold still follow previous works. Therefore,this work appears to be a fusion of several previous studies, applied in a new task, lacking innovation in terms of the methodology. I hope the author could clarify the difference between the techniques used in this paper and the origianl ones, and the  improvements upon previous approaches. This would give the authors an opportunity to highlight any methodological innovations that may not be immediately apparent. Therefore, I have doubts about whether the level of novelty in this paper meets the standards expected for a top-tier conference like ICLR.

W.2 In the experiment, RFDiffusion based method like RFDiffusionAA is compared, but there should be more works which can fulfill such tasks, such as [6] and [7]. I am wondering that can such baselines be all included into experimental comparison? I hope the author  could discuss any technical or practical limitations that may have prevented including other baselines. Besides, I think AlphaFold3 can also fulfill this task, so what is the advantage of yours over AlphaFold3? Or there are several distinct differences between AtomFlow and AlphaFold3 in design protein binder with molecule targets? I hope the authors can clarify if AlphaFold3 is directly applicable to this specific task of designing ligand-binding proteins from 2D molecular graphs, and if so, I hope the authors to explain the key differences or advantages of their approach compared to AlphaFold3 for this particular application.

[1] Kong, Full-Atom Peptide Design with Geometric Latent Diffusion, https://arxiv.org/abs/2402.13555

[2] Avishek Joey Bose, SE(3)-Stochastic Flow Matching for Protein Backbone Generation, https://arxiv.org/abs/2310.02391

[3] Xu, GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, https://arxiv.org/abs/2203.02923

[4] Emiel, Equivariant Diffusion for Molecule Generation in 3D, https://arxiv.org/abs/2203.17003

[5] Lin, PPFlow: Target-aware Peptide Design with Torsional Flow Matching, https://arxiv.org/abs/2405.06642

[6] Zhang, Pocketgen: Generating Full-Atom Ligand-Binding Protein Pockets. bioRxiv, 2024.

[7] Hannes Stark, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic Self-Conditioned Flow Matching for joint Multi-Ligand Docking and Binding Site Design. In Forty-first International Conference on Machine Learning

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 3
questions: LCFM and LCFM-FAPE indeed give zero at the same point. But why their gradients will be similar? 

Can you explain more what does it mean "since the aligning object x varies upon training."? It contradicts all the theory you have just presented, no?

Can you comment on proper cross-validation experiments by removing certain fold classes from training? The same is valid for affinity experiments. Please also consider novel data with one protein and multiple ligands as validation, see four datasets from CASP16, for example.
soundness: 2
strengths: It's a well written technical paper about a novel architecture that can be useful for the bioinformatics community. I am thankful to authors for providing the code. Please remove the infinite loop and the check for the clash score, it takes too much time. 

We have executed it on a couple of example, including recent L targets from CASP. I have carefully examined the generated output for L1000 ligand set. With a significant probability (about 30%), AtomFlow finds the native experimental fold (we have ~20 PDB structures with ligands different from those in CASP) and the binding site, and is also rather stable with variations of the sequence length. If I vary the design length slightly, AtomFlow seems to prefer to shorten loops even those inside the binding site, so it seems it has memorized the space of all protein folds in the PDB and simply selects the one it has seen during training. It actually can be a useful tool for fold search and classification.

Individual protein-ligand contacts vary in quality significantly. So I have doubts in any usefulness of further affinity prediction providing a very low quality of experimental affinity data (see below).

We have also tried to generate folds for a purine ligand, which is a part of nucleotide bases (expecting many different folds to be generated). Surprisingly, at a constant generation length, the results were extremely conservative, with zero novelty and extremely low diversity, preferring the fold that dominantly occurs in the PDB. So, I suggest the authors to run more cross-validation tests hiding some fold classes from the training and trying to reproduce them in the test. Otherwise it seems there is a significant bias towards the folds that are abandon in the PDB.
summary: The authors present the ATOMFLOW architecture for the design of protein folds conditioned by the ligand SMILES description. The tool actually can be useful for the bioinformatics community, but many more validation tests are required here. I also thank the authors for providing the code, we had fun time generating different folds and assessing them.
weaknesses: There is little technical novelty in the architecture blocks, loss function of flow matching process (but the architecture is novel and useful).

Since you are using the FAPE loss, why do you need equations 3-5, references to Riemannian Flow Matching, and fancy words? It will be much cleaner to remove this part. 

Binding affinity experiments are far from complete and better to be removed. You need to split the data correctly and run multiple affinity benchmarks, for example, following CASF and DUDe protocols. It will also be useful to provide energies and folds for some well-studied proteins that have multiple structures with different ligands (they are currently coming in the PDB, you may look at the last CASP experiment). Please also provide Vina scores for native ligand poses. Vina is not the best proxy for the affinity, as it was design only to rank binding poses for the same protein-ligand complex.

Please also be aware of the quality of current affinity training sets -- e.g., https://pubs.acs.org/doi/full/10.1021/acs.jcim.4c00049.

"The scatter plot of scRMSD vs. pdbTM shown in Figure 6B reveals that ATOMFLOW has the ability to generate structures that are quite different from existing proteins with acceptable designability" - Well, all of your designs have scores > 0.5, a few are around 0.5. It means that the fold structure is the same as in the PDB, no new folds have been discovered. Please try to discover a fold with a TM score around 0.2.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: No specific questions.
soundness: 3
strengths: * The work is solidly executed, with sensible architectural choices, strong initial evaluations, and clear and concise writing. 
* The task tackled is significant and the competitive results signify well-executed model engineering and training practices.
* The authors re-derive the quotient-space flow matching fromework from AlphaFlow with more solid theoretical justification.
* The paper is quite clearly written. The figures are well made, visually appealing, and informative.
summary: This paper presents AtomFlow, a flow matching model for designing a protein structure to bind a small molecule ligand. The model jointly denoises the structure of the protein and ligand and thus does not require knowledge of the ligand pose, unlike RFDiffusionAA. The architecture is based on AlphaFold and predicts denoised structures from a distance map input of the input structure. The method is evaluated on (1) the set of 4 ligands studied in RFDiff-AA, and (2) an expanded set of ligands curated by the authors. AtomFlow is shown to have comparable or better designability than RFAA as well as similar Vina score.
weaknesses: **Originality**
* The methodology can be described as a flow-matching version of RFDiff-AA and does not score high on originality / novelty from a ML perspective. Further, the flow model architecture and noising process are based on AlphaFlow, with different justification but no difference in practice as far as I can tell. To improve on this axis, while it's not clear that more methodological novelty is needed for its own sake, the authors could focus on novel evaluations or applications of the proposed method.

**Quality**
* The computational evaluations are well executed, but limited in scope. Most of the analysis focuses on only 4 ligands, raising concerns about sample size and statistical significance. 
* The diversity and novelty evaluations are nice, but only AtomFlow is evaluated, not RFDiff-AA or the other baselines.


**Significance**
* The overall significance of the contribution is unclear as it represents an incremental methodological advance over RFDiff-AA with more or less the same model capabilities. The authors argue that not needing to specify the ligand pose is a big plus, but no meaningful evidence or use case is provided for this distinction. After all, RFDiff-AA has been experimentally validated, whereas AtomFlow-generated poses have not. It is of course not expected for a ML submission to experimentally validate the proteins, but it should be made a bit clearer why the main point of difference with RFDiff-AA is important to tackle as a ML problem.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: - As I mentioned in the weakness, I hope the author can add some ablation experiments to demonstrate the effectiveness of some training strategies used in the paper.
- The author can provide a more detailed explanation of the composition of the loss in Section 4.4. I haven't seen relevant information in the appendix yet.
soundness: 3
strengths: -  The paper is well-structured and logically organized. Readers can easily follow the authors’ thought process, and the content is presented in a way that is easy to understand.
- Conceptualizing both proteins and molecules as biotokens with representative atoms is an interesting approach. Although similar concepts are mentioned in Umol and RoseTTAFold All-Atom, using this framework to design specific proteins without relying on the initial ligand structure demonstrates impressive performance and is commendable.
- Compared to RFDiffusionAA, ATOMFLOW shows competitive results in specific protein-molecule design tasks as well as in binding affinity prediction. The supplementary experiments in Table 8 of the appendix provide additional support for the current experimental results.
summary: In this work, the authors present ATOMFLOW, a deep generative model based on a flow-matching framework for designing ligand-binding proteins from 2D target molecular graphs. ATOMFLOW operates on representative atoms of biological tokens to capture ligand flexibility, iteratively generating ligand conformations and protein backbone structures. Overall, this paper introduces an innovative model that offers a novel approach to designing proteins that bind to specific molecules.
weaknesses: - The intervals used in modeling the distance map for ligands and proteins are not uniformly divided. While the authors account for differing precision requirements between residues and atoms, given that the paper frames proteins and molecules as biotokens with representative atoms, could the distance partitioning criteria be standardized? Establishing a more universal distance characterization approach—without introducing additional biases—might provide a more consistent method for distance representation.
-  The authors mention that “the model is first trained on solely generating the protein structure for 40k steps.” In this phase, is ligand information completely omitted, or is it set as specific input information? Further details on the full training process would be beneficial.
- The paper indicates that, during training, the FAPE loss is divided into protein-protein interaction, protein-ligand interaction, and ligand-ligand interaction components. However, these three losses are not directly introduced or explained in detail. Could the authors elaborate on the definitions of these three loss components and how they are calculated?
- The paper appears to lack ablation studies that could support certain design choices in the model. For instance, what impact would removing the first phase of training (the 40k steps on protein structure generation only) have on model performance? Additionally, how were the initial training and fine-tuning parameters in the second training phase determined, and what criteria guided the transition steps?

Meta-review: The paper presents a method, AtomFlow, to design protein structures that bind to (a small molecule) input ligands. Contrary to prior work, it does so without the assumption that the ligand’s bound pose is given by adopting a flow that jointly generates the structure of the protein and the binding ligand. The general method is a parallel to the diffusion-based RFDiffAA that is based on flow matching. It evaluates the method against RFDiffAA baselines and shows comparable designability results, without knowing the bound structure, both on the RFDiffAA benchmark as well as on a new set of ligands. 

The reviewers acknowledged the importance of the task, the suitability of the architecture and novel use of flow matching, the clarity of the writing and presentation, the meaningful results in the experiments, the practical usefulness of the method,

On the other hand, the reviewers were concerned about the lack of technical novelty in light of prior work such as AlphaFlow and RFDiffAA (for a technical paper), and –as an application paper– they were concerned about  proper and conclusive comparison with baselines, the size of the test set, the extent and quality of evaluation tasks and metrics.

The authors did a thorough rebuttal which addressed some of the concerns and provided additional results and insights, through which the paper clearly improved as indicated by the reviewers. Some concerns were fully resolved such as comparison with methods beyond RF-DIFF family and several clarifications on how the evaluation is currently done. 

The AC believes the current paper has several merits and majorly improved during the rebuttal period as indicated by the reviewers. However, during the discussion among the reviewers and after the rebuttal period, the reviewers felt that the paper can further improve along the directions mentioned in the original reviews and during the rebuttal. The AC agrees with this evaluation and believes it is important that, after the suggested improvements, the paper goes through another round of review. Therefore, the AC suggests rejection at this time. A resubmission considering all the valuable feedback from the reviewers should make the paper much stronger for the next top venue. The main points that can improve the paper are the suggested directions to extend the empirical evidence so that it becomes a stronger application paper. Alternatively, or in tandem, the paper can better highlight what technical contribution the paper has that goes beyond the application of the paper.

Decision: Reject