SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haitao Lin
linhaitao@westlake.edu.cn
Yufei Huang
huangyufei@westlake.edu.cn
Odin Zhang
haotianzhang@zju.edu.cn
Lirong Wu
wulirong@westlake.edu.cn
Siyuan Li
lisiyuan@westlake.edu.cn
Zhiyuan Chen
Deep Potential
chenzhiyuan@dp.tech
Stan Z. Li *
stan.zq.li@westlake.edu.cn
In recent years, AI-assisted drug design methods have been proposed to generate molecules given the pockets' structures of target proteins. Most of them are atom-level-based methods, which consider atoms as basic components and generate atom positions and types. In this way, however, it is hard to generate realistic fragments with complicated structures. To solve this, we propose D3FG, a functional-group-based diffusion model for pocket-specific molecule generation and elaboration. D3FG decomposes molecules into two categories of components: functional groups defined as rigid bodies and linkers as mass points. And the two kinds of components can together form complicated fragments that enhance ligand-protein interactions. To be specific, in the diffusion process, D3FG diffuses the data distribution of the positions, orientations, and types of the components into a prior distribution; In the generative process, the noise is gradually removed from the three variables by denoisers parameterized with designed equivariant graph neural networks. In the experiments, our method can generate molecules with more realistic 3D structures, competitive affinities toward the protein targets, and better drug properties. Besides, D3FG as a solution to a new task of molecule elaboration, could generate molecules with high affinities based on existing ligands and the hotspots of target proteins.
The established notion in drug-target identification is that similar structures perform similar functions. This principle allows classical computer-aided drug design (CADD) to abstract protein-ligand interactions as pharmacophores, aligning similar functional groups and extracting pharmaceutical
information from these structures. Fitting suitable functional groups into the pharmacophores can enhance ligand-protein interactions, thus improving drug efficiency [1, 2].
Artificial intelligence has already achieved outstanding success in protein design [3-8] and thus led to a new round of attention in drug design focused on AI-assisted drug design (AIDD). A number of methods of AI-assisted drug design on pocket-specific molecule generation are emerging [9-15], thanks to developments in deep generative models [16-21] and graph neural networks (GNNs) [22-31]. These methods usually focus on atom-level generation, which first generates atom types and positions and then assembles atoms into molecules that can bind to the protein pockets. Although significant progress has been made, they are still weak in two aspects. For one thing, it is hard for them to generate realistic functional groups that contribute pharmacological effects to the target as classical CADD methods are able to. It is shown that generating benzene rings is uncommon compared with the reference molecules, not to mention some large functional groups with complex structure constraints such as purine, indole, etc (Table. 1 in Sec. 5.2). For another, trade-offs between efficiency and sufficiency of protein context place them in a dilemma. For example, TARGETDIFF [13] employs sufficient protein context, which disassembles the amino acids into atoms, but leads to inefficiency due to a large node number (412.14 on average) in GNN's message passing. DIFFSBDD [14] simplifies the representation of protein context, by only using  $C_{\alpha}$ 's positions and residue types, resulting in a reduction of GNNs' node number (68.10 on average) but the insufficiency of context information.
To address the above issues, we establish a functional-group-based diffusion model (D3FG), including the following contributions: 1. Method Novelty. We denote the molecules' functional groups and proteins' amino acids as the same level's fragments, in which the intra-relative positions of atoms are fixed like rigid bodies, and represent single atoms as linkers. The positions and orientation of local structures and the atom type variables are generated gradually through denoising processes. The fragment-linker designation leads the binding systems to heterogeneous graphs, and thereby, two schemes are proposed as solutions, which achieve competitive performance in terms of molecule structures, binding affinity, and drug properties, and sufficiency in encoding protein context by employing more features. 2. Dataset Establishment. Though the CrossDocked2020 [32] has been a widely-used dataset for evaluation methods' performance on the task, the analysis of the molecules' functional groups of it is missing. We deeply explore the details of the inter-relative positions and types of functional groups of the molecules and establish an extendable database of common functional groups. 3. New task. Besides molecule generation, we propose molecule elaboration as another task that our model can fulfill. Fragment hotspot maps (FHM) [33, 34] are used to preprocess paired protein-molecule in CrossDocked2020 for the task. As a result, our model generates molecules with high binding affinity based on the reference.
For a binding system composed of a protein-molecule pair (also called protein-ligand pair) as  $\mathcal{C}$ , which contains  $N_{\mathrm{aa}}$  amino acids of proteins and  $N_{\mathrm{fg}}$  functional groups and  $N_{\mathrm{at}}$  single atoms of molecules, we represent the index set of the molecule's single atoms as  $\mathcal{I}_{\mathrm{at}}$ , functional groups as  $\mathcal{I}_{\mathrm{fg}}$ , and the protein's amino acids as  $\mathcal{I}_{\mathrm{aa}}$ , where  $|\mathcal{I}_{\mathrm{fg}}| = N_{\mathrm{fg}}$ ,  $|\mathcal{I}_{\mathrm{at}}| = N_{\mathrm{at}}$  and  $|\mathcal{I}_{\mathrm{aa}}| = N_{\mathrm{aa}}$ . Note that a molecule can be disassembled into functional groups and single atoms other than functional groups, which we also call linkers. For a protein, the amino acids can be represented by its type,  $\mathrm{C}_{\alpha}$  atom coordinate, and the orientation, denoted as  $s_i \in \{1, \ldots, 20\}$ ,  $x_i \in \mathbb{R}^3$ ,  $O_i \in \mathrm{SO}(3)$ , where  $i \in \mathcal{I}_{\mathrm{aa}}$ . For a molecule, assuming there are  $M_{\mathrm{fg}}$  and  $M_{\mathrm{at}}$  possible types in total functional groups and linker atoms respectively, the functional groups can be represented as the three elements if the inter-relative positions are fixed, as  $s_j$ ,  $x_j$  and  $O_j$ , where  $s_j \in \{21, \ldots, 21 + M_{\mathrm{fg}}\}$  is the type,  $x_j \in \mathbb{R}^3$  is the predefined center atom's coordinate, and  $O_j \in \mathrm{SO}(3)$  can also be obtained in the same way as amino acids (See Sec. 5.1 and Appendix D.) for  $j \in \mathcal{I}_{\mathrm{fg}}$ ; And its linkers can be represented as  $s_k$ ,  $x_k$  and  $O_k$ , with  $s_k \in \{22 + M_{\mathrm{fg}}, \ldots, 22 + M_{\mathrm{fg}} + M_{\mathrm{at}}\}$ ,  $k \in \mathcal{I}_{\mathrm{at}}$  and  $O_k = \operatorname{diag}\{1, 1, 1\} = I$ .
Therefore,  $\mathcal{C} = \{(s_i,\pmb {x}_i,\pmb {O}_i)\}_{i = 1}^{N_{\mathrm{aa}} + N_{\mathrm{fg}} + N_{\mathrm{at}}}$  can be split into two sets as  $\mathcal{C} = \mathcal{P}\cup \mathcal{M}$ , where  $\mathcal{P} = \{(s_i,\pmb {x}_i,\pmb {O}_i):i\in \mathcal{I}_{\mathrm{aa}}\}$  and  $\mathcal{M} = \{(s_j,\pmb {x}_j,\pmb {O}_j):j\in \mathcal{I}_{\mathrm{fg}}\cup \mathcal{I}_{\mathrm{at}}\}$ . For protein-specific molecule generation, our goal is to establish a probabilistic model to learn the distribution of molecules conditioned on the target proteins, i.e.  $p(\mathcal{M}|\mathcal{P})$ . In the following, we omit  $i\in \mathcal{I}_{\mathrm{aa}}$  and  $j\in \mathcal{I}_{\mathrm{fg}}\cup \mathcal{I}_{\mathrm{at}}$  by default unless specified.
3D molecule generation. Previous methods on molecule generation fucus on 1D-smiles-string [35-38] or 2D-molecule-graph [39-43]. In recent years, more works concentrate on 3D molecule generation, thanks to fast development in equivariant graph neural networks [44-46] and generative models [16-21]. Molecular conformation generation aims to generate 3D structures of molecules with stability, given 2D molecule graph structure, [47-51]. Further, De novo molecular generation attempts to generate both 2D chemical formulas and 3D structures from scratch [52-54].
Fragment-based drug design. Previously, works on fragment-based molecule generation are proposed. For example, JT-VAE [55] generates a tree-structured scaffold over chemical substructures and combines them into a 2D-molecule. PS-VAE [56] can automatically discover frequent principal subgraphs from the dataset, and assemble generated subgraphs as the final output molecule in 2D. Further, DEEPFRAG [57] predicts fragments conditioned on parents and the pockets, SQUID [58] generates molecules in a fragment level conditioned on molecule's shapes. FLAG[59] auto-regressively generates fragments as motifs based on the protein structures in 3D.
Structure-based drug design. Success in 3D molecule generation and an increasing amount of available structural data of protein and molecules raises scientific interests in structure-based drug design (SBDD), which aims to generate both 2D molecule graphs and 3D structures conditioned on target protein structure as contextual constraints. For example, LIGAN [60] and 3DSBDD [12] are grid-based models which predict whether the grid points are occupied by specific atoms. By harnessing 3D equivariant graph neural networks, POCKET2MOL [9] and GRAPHBP [10] generate atoms auto-regressively and model the probability of the next atom's type as discrete categorical attribute and position as continuous geometry. FLAG[59] is also fragment-based, but still generates motifs in an auto-regressive way. Recently, utilizing the diffusion denoising probability models [61-63, 53, 64], a series of SBDD methods generate ligands conditioned on the target pockets at full atom levels [13-15].
D3FG firstly decomposes molecules into two categories of components: functional groups and linkers, and them use the diffusion generative model to learn the type and geometry distributions of the components. In this section, we describe the D3FG by four parts: (i) the diffusion model as the generative framework, in which the three variables are generated; (ii) the denoiser parameterized by graph neural networks, satisfying certain symmetries so that the generative model is SE-(3) invariant; (iii) the sampling process in which the molecules are generated by the trained models; (iv) the further problems resulted from heterogenous graph with two solutions.
4.1 Diffusion Models
Diffusion models construct two Markov processes to learn the data distributions. The first called the forward diffusion process adds noises gradually until the noisy data's distribution reaches the prior distribution; The other called the generative denoising process, gradually removes the noise from the data sampled from the prior distribution until they are recovered to the desired data distribution. Assume there are  $T$  steps in both processes, and we denote  $\mathcal{M}^t = \{(s_j^t,\pmb{x}_j^t,\pmb{O}_j^t)\}$  as the  $t$ -th noisy state in the forward process, with  $\mathcal{M}^T\sim \mathrm{Prior}(\mathcal{M}^T)$ , and  $\mathcal{M}^0 = \mathcal{M}$ , where the transition distribution is denoted as  $q(\cdot |\cdot)$ ; In the generative process, sample goes from  $T$  to 0, in which the transition distribution is denoted as  $p(\cdot |\cdot)$ . Here we define the forward and generative processes of  $s_j^t$ ,  $\pmb{x}_j^t$ , and  $\pmb{O}_j^t$ .
Absorbing diffusion for functional group and linker types. Let  $s_j^t$  as the one-hot encoding of the type of a single functional group or linker. The forward process followed by D3PM [65] randomly transits  $s_j^t$  into the absorbing state  $K$  ( $K = 23 + M_{\mathrm{fg}} + M_{\mathrm{at}}$ ) with
$$
\begin{array}{l} q \left(s _ {j} ^ {t} \mid s _ {j} ^ {t - 1}\right) = \text {M u l t i n o m i a l} \left(\boldsymbol {s} _ {j} ^ {t - 1} \boldsymbol {Q} ^ {t}\right) \tag {1} \\ q \left(s _ {j} ^ {t} \mid s _ {j} ^ {0}\right) = \text {M u l t i n o m i a l} \left(s _ {j} ^ {0} \bar {\boldsymbol {Q}} ^ {t}\right) \\ \end{array}
$$
Figure 1: An illustration of the workflows of D3FG of the two schemes.
where  $\bar{\mathbf{Q}}^t = \mathbf{Q}^1\mathbf{Q}^2\dots \mathbf{Q}^t$  and  $[\mathbf{Q}^t]_{mn} = q(s_j^t = n|s_j^{t - 1} = m)$  denotes diffusion transition probabilities, with
$$
\left[ \boldsymbol {Q} ^ {t} \right] _ {m n} = \left\{ \begin{array}{l l} 1 & \text {i f} m = n = K \\ 1 - \beta_ {\text {t y p e}} ^ {t} & \text {i f} m = n \neq K \\ \beta_ {\text {t y p e}} ^ {t} & \text {i f} m \neq K, n = K \end{array} \right.. \tag {2}
$$
$\beta_{\mathrm{type}}^{t}$  monotonically increases from 0 to 1, means that when  $t = T$ , all the type variables are absorbed into the  $K$ -th category. In the generative process, it first samples  $N_{\mathrm{at}}$  linkers and  $N_{\mathrm{fg}}$  functional groups whose types are all in the absorbing states, selects  $(1 - \beta_{\mathrm{type}}^{t}) \times 100\%$  of them respectively, and transforms their types from the absorbing states to the predicted ones by
$$
p \left(s _ {j} ^ {t - 1} \mid \mathcal {M} ^ {t}, \mathcal {P}\right) = \text {M u l t i n o m i a l} \left(F \left(\mathcal {M} ^ {t}, \mathcal {P}\right) [ j ]\right). \tag {3}
$$
where  $F$  is the type denoiser parameterized with a neural network, and  $F(\cdot ,\cdot)[j]$  predicts the probability of the types for the  $j$ -th selected functional groups or linkers. With the effectiveness of BERT-style training [66], the denoiser directly predicts  $p(s_j^0 |\mathcal{M}^t,\mathcal{P})$ , leading to a training objective as
$$
L _ {\text {t y p e}} ^ {t} = \mathbb {E} _ {\mathcal {M} ^ {t}} \left[ \sum_ {j} \log p \left(s _ {j} ^ {0} \mid \mathcal {M} ^ {t}, \mathcal {P}\right) \right]. \tag {4}
$$
Gaussian diffusion for center atom positions. By defining the center atom in a functional group as shown in Appendix D, and itself as the center in a linker, the Cartesian coordinate of center nodes  $x_{j}$  represents its position. The forward transition distributions followed by DDPM [61] read
$$
q \left(\boldsymbol {x} _ {j} ^ {t} \mid \boldsymbol {x} _ {j} ^ {t - 1}\right) = \mathcal {N} \left(\boldsymbol {x} _ {j} ^ {t} \mid \sqrt {1 - \beta_ {\text {p o s}} ^ {t}} \cdot \boldsymbol {x} _ {j} ^ {t - 1}, \beta_ {\text {p o s}} ^ {t} \boldsymbol {I}\right); \tag {5}
$$
$$
q (\pmb {x} _ {j} ^ {t} | \pmb {x} _ {j} ^ {0}) = \mathcal {N} \left(\pmb {x} _ {j} ^ {t} | \sqrt {\bar {\alpha} _ {\mathrm {p o s}} ^ {t}} \cdot \pmb {x} _ {j} ^ {0}, (1 - \bar {\alpha} _ {\mathrm {p o s}} ^ {0}) \pmb {I}\right),
$$
in which  $\beta_{\mathrm{pos}}^t$  increases from 0 to 1, means that the noise levels are increasing and the data's coordinate signals fade out during the forward diffusion, with  $\alpha_{\mathrm{pos}}^t = 1 - \beta_{\mathrm{pos}}^t$ ,  $\bar{\alpha}_{\mathrm{pos}}^t = \alpha_{\mathrm{pos}}^0\alpha_{\mathrm{pos}}^1\dots \alpha_{\mathrm{pos}}^t$ , and finally  $\pmb{x}_j^T\sim \mathcal{N}(\pmb {0},\pmb {I})$ . Note that Eq. (5) is equivalent to the Markov process of  $\pmb{x}_j^t = \sqrt{\bar{\alpha}_{\mathrm{pos}}^t}\cdot \pmb{x}_j^0 +$ $\sqrt{1 - \bar{\alpha}_{\mathrm{pos}}^{0}}\cdot \epsilon_{j}$ , where  $\epsilon_{j}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . Rather than predicting the mean value of the reverse transition distribution in the generative process, the position denoiser  $G$  approximates the added noise  $\epsilon_{j}$  with the reparameterization trick as
$$
\begin{array}{l} p (\boldsymbol {x} _ {j} ^ {t - 1} | \mathcal {M} ^ {t}, \mathcal {P}) = \mathcal {N} \left(\boldsymbol {x} _ {j} ^ {t - 1} | \boldsymbol {\mu} _ {\mathrm {p o s}} (\mathcal {M} ^ {t}, \mathcal {P}), \beta_ {\mathrm {p o s}} ^ {t} \boldsymbol {I}\right); \\ \boldsymbol {\mu} _ {\text {p o s}} \left(\mathcal {M} ^ {t}, \mathcal {P}\right) = \frac {1}{\sqrt {\alpha_ {\text {p o s}} ^ {t}}} \left(\boldsymbol {x} _ {j} ^ {t} - \frac {\beta_ {\text {p o s}} ^ {t}}{\sqrt {1 - \bar {\alpha} _ {\text {p o s}} ^ {t}}} G \left(\mathcal {M} ^ {t}, \mathcal {P}\right) [ j ]\right). \tag {6} \\ \end{array}
$$
The training objective is thus established in a score-based way, as
$$
L _ {\text {p o s}} ^ {t} = \mathbb {E} _ {\mathcal {M} ^ {t}} \left[ \sum_ {j} \| G (\mathcal {M} ^ {t}, \mathcal {P}) [ j ] - \epsilon_ {j} \| _ {2} ^ {2} \right]. \tag {7}
$$
SO(3) diffusion for functional group orientations. By regarding the functional groups as rigid bodies, orientations together with the center atoms' positions determine all atoms' positions. Here we represent the orientation geometry as elements in SO(3). Following [67], we use isotropic Gaussian distribution on SO(3) [68] to formulate the process, i.e.  $\mathcal{IG}_{\mathrm{so}(3)}(\cdot |\pmb{\mu}_{\mathrm{ori}},\sigma_{\mathrm{ori}})$ , in which  $\pmb{\mu}_{\mathrm{ori}}$  and  $\sigma_{\mathrm{ori}}$  are viewed as mean orientation and variance, in analogy with Gaussian distribution. The transition distribution for orientation matrices  $O_{j}$  reads
$$
q \left(\boldsymbol {O} _ {j} ^ {t} \mid \boldsymbol {O} _ {j} ^ {0}\right) = \mathcal {I} \mathcal {G} _ {\text {s o (3)}} \left(\boldsymbol {O} _ {j} ^ {t} \mid \lambda_ {\text {o r i}} \left(\bar {\alpha} _ {\text {o r i}} ^ {t}, \boldsymbol {O} _ {j} ^ {0}\right), \left(1 - \bar {\alpha} _ {\text {o r i}} ^ {t}\right)\right), \tag {8}
$$
$\lambda_{\mathrm{ori}}(\bar{\alpha}_{\mathrm{ori}}^{t}, \mathbf{O}_{j}^{0})$  is the geodesic flow from  $\pmb{I}$  to  $\pmb{O}_{j}^{t}$  by the amount  $\bar{\alpha}_{\mathrm{ori}}^{t}$ , as  $\lambda_{\mathrm{ori}}(\bar{\alpha}_{\mathrm{ori}}^{t}, \mathbf{O}_{j}^{0}) = \exp(\bar{\alpha}_{\mathrm{ori}}^{t} \log(\mathbf{O}_{j}^{0}))$ , where  $\exp(\cdot)$  and  $\log(\cdot)$  are exponential and logarithm map on the SO(3) manifold. As  $\alpha_{\mathrm{ori}}^{t} \to 0$ ,  $\lambda_{\mathrm{ori}}(\bar{\alpha}_{\mathrm{ori}}^{t}, \mathbf{O}_{j}^{0}) \to \pmb{I}$ .  $\{\beta_{\mathrm{ori}}^{t}\}_{t=0}^{T}$  is the predefined noise level schedule ranging from 0 to 1 as  $t$  increases,  $\alpha_{\mathrm{ori}}^{t} = 1 - \beta_{\mathrm{ori}}^{t}$  and  $\bar{\alpha}_{\mathrm{ori}}^{t} = \alpha_{\mathrm{ori}}^{0} \alpha_{\mathrm{ori}}^{1} \dots \alpha_{\mathrm{ori}}^{t}$ .
In the generative process, an orientation denoiser  $H$  is used to predict the mean orientation in the isotropic Gaussian distribution, which reads
$$
p \left(\boldsymbol {O} _ {j} ^ {t - 1} \mid \mathcal {M} ^ {t}, \mathcal {P}\right) = \mathcal {I} \mathcal {G} _ {\text {s o} (3)} \left(\boldsymbol {O} _ {j} ^ {t - 1} \mid H \left(\mathcal {M} ^ {t}, \mathcal {P}\right) [ j ], \beta_ {\text {o r i}} ^ {t}\right). \tag {9}
$$
We use the same loss function as in [69] to minimize the expected discrepancy measured by the inner product between the data orientation matrices and the predicted ones, which reads
$$
L _ {\mathrm {o r i}} ^ {t} = \mathbb {E} _ {\mathcal {M} ^ {t}} \left[ \sum_ {j} \| \left(H \left(\mathcal {M} ^ {t}, \mathcal {P}\right) [ j ]\right) ^ {\top} \boldsymbol {O} _ {j} ^ {0} - \boldsymbol {I} \| _ {F} ^ {2} \right]. \tag {10}
$$
4.2 Parametrization of Denoisers with Neural Networks
Amino acid context encoding. In order to decrease the computational complexity, we denote the protein context at amino-acid levels. Besides the amino acid types,  $\mathrm{C}_{\alpha}$  atom coordinate and the orientation, each atom's coordinates in the local system and three torsion angles including angles around 'N-  $\mathrm{C}_{\alpha}$ ' bond, 'C  $\alpha$ -C' bond and 'C-N' bond, are also used as intra-amino-acid features, which are concatenated and embedded by an MLP to create the intra-amino-acid embedding vector  $e_{i}$ . For inter-amino-acid features, the pair of amino acid types, sequential relationships (if the two amino acids are adjacent in the protein sequence), pairwise distances between  $\mathrm{C}_{\alpha}$  and inter-residue dihedrals are all embedded as inter-amino-acid embedding vector  $z_{i,j}$  with  $i,j \in \mathcal{T}_{\mathrm{aa}}$ . Note that these embedding vectors are all translational and rotational invariant (Appendix. C).
Denoisers with equivariance. In the setting of generative models, the learned distribution  $p(\mathcal{M}|\mathcal{P})$  should be equivariant to translation and rotation, such that  $p(\mathbf{T}_g(\mathcal{M})|\mathbf{T}_g(\mathcal{P})) = p(\mathcal{M}|\mathcal{P})$  for any  $g \in \mathrm{SE}(3)$ , where  $\mathbf{T}_g$  is the corresponding roto-translational transformations, and  $\mathbf{T}_g(\mathcal{M})$  means each atom in the molecule is rotated and translated by  $\mathbf{T}_g$ . In the setting of diffusion models, the following proposition indicates the translational and rotational equivariance of each denoisers.
Proposition 1. Let  $p(\boldsymbol{x}^T)$ ,  $p(\boldsymbol{O}^T)$ , and  $p(s^T)$  be SE(3)-invariant distribution, and the transition distributions  $p(\boldsymbol{x}^{t-1}|\mathcal{M}^t,\mathcal{P})$  be SE(3)-equivariant,  $p(\boldsymbol{O}^{t-1}|\mathcal{M}^t,\mathcal{P})$  be T(3)-invariant and SO(3)-equivariant, and  $p(s^{t-1}|\mathcal{M}^t,\mathcal{P})$  be SE(3)-invariant, then the density  $p(\mathcal{M}|\mathcal{P})$  modeled by the reverse Markov Chains in the generative process of diffusion models is SE(3)-equivariant.
According to the Proposition 1, while the denoisers for functional group and linker types are roto-translational invariant, the denoiser for positions should be roto-translational equivariant, and it for orientations should be translational invariant and rotational equivariant. Therefore, we employ our denoiser's network structures with IPA [3] by harnessing the expressivity of TRANSFORMER [70] and roto-translational equivariance of LoCS [71]. Denote the binding system as a graph, in which the nodes are composed of amino acids, functional groups, and linkers, and the edges are established with
Algorithm 1 Joint Generation for Molecules using D3FG
Input: Zero-centered protein  $\{\pmb{s}_i,\pmb{x}_i,\pmb{O}_i\}_{i\in \mathcal{I}_{\mathrm{aa}}}$ , and graph denoiser  $F,G,H$ , and node number sampler  $\phi_{\mathrm{fg}},\phi_{\mathrm{at}}$
Sample  $N_{\mathrm{fg}} \sim \phi_{\mathrm{fg}}, N_{\mathrm{at}} \sim \phi_{\mathrm{at}}$ , leading to the index set  $\mathcal{I}_{\mathrm{fg}}$  and  $\mathcal{I}_{\mathrm{at}}$ .
Sample initial states of functional groups,  $\{s_j^T,\pmb{x}_j^T,\pmb{O}_j^T\}_{j\in \mathcal{I}_{\mathrm{fg}}\cup \mathcal{I}_{\mathrm{at}}}$ , where  $s_j^t = K$ ,  $\pmb{x}_j^T\sim \mathcal{N}(\pmb {0},\pmb {I})$ $\pmb{O}_j^T\sim \mathrm{Uniform}_{\mathrm{SO}(3)}$  for  $j\in \mathcal{I}_{\mathrm{fg}}$  else  $\pmb{O}_j^T = \pmb{I}$
for  $t$  in  $T - 1, T - 2, \ldots, 1$  do
Sample  $\{s_j^{t - 1},\pmb {x}_j^{t - 1},\pmb {O}_j^{t - 1}\}_{j\in \mathcal{I}_{\mathrm{fg}}\cup \mathcal{I}_{\mathrm{at}}}$  as Eq. (3), (6) and (9) and update  $\mathcal{M}^{t - 1}$ .
Output:  $\mathcal{M} = \{s_j^0,\pmb {x}_j^0,\pmb {O}_j^0\}_{j\in \mathcal{I}_{\mathrm{fg}}\cup \mathcal{I}_{\mathrm{at}}}$
K-nearest neighbor. Let  $\{h_i : i \in \mathcal{I}_{\mathrm{aa}} \cup \mathcal{I}_{\mathrm{fg}} \cup \mathcal{I}_{\mathrm{at}}\}$  be the node embedding which is SE(3)-invariant,  $\{e_i : i \in \mathcal{I}_{\mathrm{aa}} \cup \mathcal{I}_{\mathrm{fg}} \cup \mathcal{I}_{\mathrm{at}}\}$  and  $\{\pmb{z}_{i,j} : i,j \in \mathcal{I}_{\mathrm{aa}} \cup \mathcal{I}_{\mathrm{fg}} \cup \mathcal{I}_{\mathrm{at}}\}$  be the previously defined intra- and inter-amino acid embedding, with  $e_j = 0$  and  $z_{i,j} = 0$  if  $i \notin \mathcal{I}_{\mathrm{aa}}$  or  $j \notin \mathcal{I}_{\mathrm{aa}}$ . The attention mechanism in IPA updates the embedding of node  $i$  as