SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Initialization and Structure Refinement. After rounds of backbone refinements, the predicted residue type distribution and the backbone coordinates become relatively stable (verified in experiments). We then apply the full-atom refinement procedures to determine the side-chain structures. The residue types are initialized by sampling the predicted residue type  $p_i$  from backbone refinement (hard assignment). The number of side-chain atoms and the atom types can then be determined. The coordinates of the side-chain atoms are initialized with the corresponding  $\alpha$ -carbon coordinates. The full-atom structure refinement is similar to backbone refinement (Sec. 3.3.1) with minor adjustments to accommodate the side-chain atoms. We include the details in the Appendix A.
Residue Type Update. Since the structures and interactions of side-chain atoms can further influence the residue types, proper residue type update algorithms are needed to achieve sequence-structure consistency. Considering the hard assignment of residue types in the full-atom refinement, we randomly sample and mask one residue in each protein pocket and update its residue type based on the surrounding environment in each iteration. Specifically, we mask the residue types, remove the sampled residues' side-chain atoms in the input to the encoder, and use the obtained corresponding residue types for residue type prediction. If the predicted result is inconsistent with the original residue type, we will update the residue type with the newly predicted one and reinitialize the sidechain structures. As our experiments indicate, the residue types and structures become stable after several refinement iterations, and sequence-structure consistency is achieved.
3.4 Model Training and Sampling
The loss function for FAIR consists of two parts. We apply cross-entropy loss  $l_{ce}$  for (sampled) pocket residues over all iterations for the sequence prediction. For the structure refinement, we adopt the Huber loss  $l_{huber}$  for the stability of optimization following previous works [28, 56]:
$$
\mathcal {L} _ {s e q} = \frac {1}{T} \sum_ {t} \sum_ {i} l _ {c e} \left(p _ {i} ^ {t}, \hat {p} _ {i}\right); \tag {9}
$$
$$
\mathcal {L} _ {\text {s t r u c t}} = \frac {1}{T} \sum_ {t} \left[ \sum_ {i} l _ {\text {h u b e r}} \left(\boldsymbol {x} \left(\boldsymbol {b} _ {i}\right) ^ {t}, \hat {\boldsymbol {x}} \left(\boldsymbol {b} _ {i}\right)\right) + \sum_ {j} l _ {\text {h u b e r}} \left(\boldsymbol {x} _ {j} ^ {t}, \hat {\boldsymbol {x}} _ {j}\right) \right], \tag {10}
$$
where  $T = T_{1} + T_{2}$  is the total rounds of refinement.  $\hat{p}_i, \hat{\pmb{x}}(\pmb{b}_i)$ , and  $\hat{\pmb{x}}_j$  are the ground-truth residue types, residue coordinates, and ligand coordinates.  $p_i^t, \pmb{x}(\pmb{b}_i)^t$ , and  $\pmb{x}_j^t$  are the predicted ones at the  $t$ -th round. In the training process, we aim to minimize the sum of the above two loss functions  $\mathcal{L} = \mathcal{L}_{seq} + \mathcal{L}_{struct}$  similar to previous works [56, 34]. After the training of FAIR, we can co-design pocket sequences and structures. Algorithm 1 and 2 in Appendix A outline model training and sampling.
3.5 Equivariance
FAIR is E(3)-equivariant, which is a desirable property of generative protein models [33, 64]:
Theorem 1. Denote the  $E(3)$ -transformation as  $T_{g}$  and the generative process of FAIR as  $\{(b_i, \pmb{x}(b_i))\}_{i=1}^m = p_\theta(\mathcal{A} \setminus \mathcal{B}, \mathcal{M})$ , where  $\{(b_i, \pmb{x}(b_i))\}_{i=1}^m$  indicates the designed pocket sequence and structure. We have  $\{(b_i, T_g \cdot \pmb{x}(b_i))\}_{i=1}^m = p_\theta(T_g \cdot (\mathcal{A} \setminus \mathcal{B}), T_g \cdot \mathcal{M})$ .
Proof. The main idea is that the E(3)-invariant encoder and E(3)-equivariant structure refinement lead to an E(3)-equivariant generative process of FAIR. We give the full proof in Appendix A.  $\square$
We proceed by describing the experimental setups (Sec. 4.1). We then assess our method for ligand-binding pocket design (Sec. 4.2) on two datasets. We also perform further in-depth analysis of FAIR, including sampling efficiency analysis and ablations in Sec. 4.3.
Datasets. We consider two widely used datasets for experimental evaluations. CrossDocked dataset [16] contains 22.5 million protein-molecule pairs generated through cross-docking. We filter out data points with binding pose RMSD greater than  $1\AA$ , leading to a refined subset with around 180k data points. For data splitting, we use mmseqs2 [58] to cluster data at  $30\%$  sequence identity, and randomly draw 100k protein-ligand structure pairs for training and 100 pairs from the remaining clusters for testing and validation, respectively. Binding MOAD dataset [21] contains around 41k experimentally determined protein-ligand complexes. We further filter and split the Binding MOAD dataset based on the proteins' enzyme commission number [4], resulting in 40k protein-ligand pairs for training, 100 pairs for validation, and 100 pairs for testing following previous work [54]. More data split schemes based on graph clustering [66, 65] may be explored.
Considering the distance ranges of protein-ligand interactions [43], we redesign all the residues that contain atoms within  $3.5\AA$  of any binding ligand atoms, leading to an average of around eight residues for each pocket. We sample 100 sequences and structures for each protein-ligand complex in the test set for a comprehensive evaluation.
Baselines and Implementation. FAIR is compared with five state-of-the-art representative baseline methods. PocketOptimizer [45] is a physics-based method that optimizes energies such as packing and binding-related energies for ligand-binding protein design. DEPACT [8] is a template-matching method that follows a two-step strategy [67] for pocket design. It first searches the protein-ligand complexes in the database with similar ligand fragments. It then grafts the associated residues into the protein pocket with PACMatch [8]. HSRN [28], Diffusion [39], and MEAN [33] are deep-learning-based models that use auto-regressive refinement, diffusion model, and graph translation method respectively for protein sequence-structure co-design. They were originally proposed for antibody design, and we adapted them to our pocket design task with proper modifications (see Appendix. B for more details).
We use AMBER ff14S force field [41] for energy computation and the Dunbrack rotamer library [55] for rotamer sampling in PocketOptimizer. For the template-based method DEPACT, the template databases are constructed based on the training datasets for fair comparisons. For the deep-learning-based models, including our FAIR, we train them for 50 epochs and select the checkpoint with the lowest loss on the validation set for testing. We use the Adam optimizer with a learning rate of 0.0001 for optimization. In FAIR, the default setting sets  $T_{1}$  and  $T_{2}$  as 5 and 10. Descriptions of baselines and FAIR are provided in the Appendix. B.
Table 1: Evaluation of different approaches on the pocket design task.
Model	CrossDocked	Binding MOAD
AAR (↑)	RMSD (↓)	Vina (↓)	AAR (↑)	RMSD (↓)	Vina (↓)
PocketOptimizer	27.89±14.9%	1.75±0.08	-6.905±2.39	28.78±11.3%	1.68±0.12	-7.829±2.41
DEPACT	22.58±8.48%	1.97±0.14	-6.670±2.13	26.12±8.97%	1.76±0.15	-7.526±2.05
HSRN	31.62±10.4%	2.15±0.17	-6.565±1.95	33.70±10.1%	1.83±0.18	-7.349±1.93
Diffusion	34.62±13.7%	1.68±0.12	-6.725±1.83	36.94±12.9%	1.47±0.09	-7.724±2.36
MEAN	35.46±8.15%	1.76±0.09	-6.891±1.86	37.16±14.7%	1.52±0.09	-7.651±1.97
FAIR	40.17±12.6%	1.42±0.07	-7.022±1.75	43.75±15.2%	1.35±0.10	-7.978±1.91
Performance Metrics. To evaluate the generated sequences and structures, we employ Amino Acid Recovery (AAR), defined as the overlapping ratio between the predicted 1D sequences and ground truths, and Root Mean Square Deviation (RMSD) regarding the 3D predicted structure of residue atoms. Due to the different residue types and numbers of side-chain atoms in the generated pockets, we calculate RMSD with respect to backbone atoms following [28, 33]. To measure the binding affinity between the protein pocket and ligand molecules, we calculate Vina Score with QVina [61, 2]. The unit is kcal/mol; a lower Vina score indicates better binding affinity. Before feeding to Vina, all the generated protein structures are firstly refined by AMBER force field [41] in OpenMM [13] to optimize the structures. For the baseline methods that predict backbone structures only (Diffusion and MEAN), we use Rosetta [1] to do side-chain packing following the instructions in their papers.
4.2 Ligand-Binding Pocket Design
Comparison with Baselines. Table 1 shows the results of different methods on the CrossDocked and Binding MOAD dataset. FAIR overperforms previous methods with a clear margin on AAR, RMSD, and Vina scores, which verifies the strong ability of FAIR to co-design pocket sequences and structures with high binding affinity. For example, the average improvements on AAR and RMSD are  $15.5\%$  and  $13.5\%$  respectively. Compared with the energy-based method PocketOptimizer and the template-matching-based method DEPACT, deep-learning-based methods such as FAIR have better performance due to their stronger ability to model the complex sequence-structure dependencies and pocket-ligand interactions. However, these deep-learning-based methods still have limitations that restrict their performance improvement. For example, the autoregressive generation manner in HSRN inevitably incurs error accumulation; Diffusion and MEAN fail to model the influence of side chain atoms; all these methods ignore the flexibility of ligands. In contrast, FAIR addresses the limitations with properly designed modules, further validated in Sec. 4.3. Moreover, FAIR performs relatively better on Binding MOAD than the CrossDocked dataset, which may be explained by the higher quality and diversity of Binding MOAD [21] (e.g. Vina score on CrossDocked -7.022 vs. -7.978 on Binding MOAD). For example, Binding MOAD has 40k unique protein-ligand complexes, while CrossDocked only contains 18k complexes.
Comparison of Generated and Native Sequences. Figure 4(a) shows the residue type distributions of the designed pockets and the native ones from two datasets. Generally, we observe that the distributions align very well. Some residue types, such as Tyrosine and Asparagine, have been used more frequently in designed than naturally occurring sequences. In contrast, residue types with obviously reduced usages in the designed sequences include Lysine, Isoleucine, and Leucine.
Convergence of Iterative Refinement Process. During the iterations of FAIR, we monitor the ratio of unconverged residues (residue types different from those in the final designed sequences) and AAR, shown in Figure 4(b) ( $T_{1} = 5$  and  $T_{2} = 15$ ). We observe that both curves converge quickly during refinement. We find that the AAR becomes stable in just several iterations (typically less than 3 iterations in our experiments) of backbone refinement and then gradually increases with the complete atom refinement, validating the effectiveness of our design. We show the curves of RMSD in Appendix C.3, which also converges rapidly. Therefore, we set  $T_{1} = 5$  and  $T_{2} = 10$  for efficiency as default setting.
Case Studies. In Figure. 3, we visualize two examples of designed pocket sequences and structures from the CrossDocked dataset (PDB: 2pc8) and the Binding MOAD dataset (PDB: 6p7r), respectively. The designed sequences recover most residues, and the generated structures are valid and realistic. Moreover, the designed pockets exhibit comparable or higher binding affinities (lower Vina scores) with the target ligand molecules than the reference pockets in the test set.
Pocket Set:EYHNNEYFQW  
Vina Score: -5.98
Pocket Set: DEHNDYYEQW  
Vina Score: -7.12
Packet Seq: YFYVKM  
Vina Score: -7.33
Figure 3: Case studies of Pocket design. We show the reference and designed structures/sequences of two protein pockets from the CrossDocked (PDB: 2pc8) and Binding MOAD datasets (PDB: 6p7r).
Pocket Seq: FYMKGKM  
Vina Score: -7.74
4.3 Additional Analysis of FAIR's Performance
Sampling Efficiency Analysis. To evaluate the efficiency of FAIR, we considered the generation time of different approaches using a single V100 GPU on the same machine. Average runtimes are shown in Figure 4(c). First, the generation time of FAIR is more than ten times faster (FAIR: 5.6s vs. DEPACT: 107.2s) than traditional PocketOptimizer and DEPACT methods. Traditional methods are time-consuming because they require explicit energy calculations and combinatorial enumerations. Second, FAIR is more efficient than the autoregressive HSRN and diffusion methods. The MEAN sampling time is comparable to FAIR since both use full-shot decoding schemes. However, FAIR has the advantage of full-atom prediction over MEAN.
Local Geometry of Generated Pockets. Since RMSD values in Tables 1 and 2 reflect only the global geometry of generated protein pockets, we also consider the RMSD of bond lengths and angles of predicted pockets to evaluate their local geometry [33]. We measure the average RMSD of the covalent bond lengths in the pocket backbone (C-N, C=O, and C-C). We consider three conventional dihedral angles in the backbone structure, i.e.,  $\phi, \psi, \omega$  [57] and calculate the average RMSD of their cosine values. The results in Table 3 of Appendix C.3 show that FAIR achieves better performance regarding the local geometry than baseline methods. The RMSD of bond lengths and angles drop  $20\%$  and  $8\%$  respectively with FAIR.
Ablations. In Table 2, we evaluate the effectiveness and necessity of the proposed modules in FAIR. Specifically, we removed the residue-level encoder, the masked residue type update procedure, the full-atom refinement, and the ligand molecule context, denoted as w/o the encoder/consistency/full-atom/ligand. The results in Table 2 demonstrate that the removal of any of these modules leads to a decrease in performance. We observe that the performance of FAIR w/o hierarchical encoder drops the most. This is reasonable, as the hierarchical encoder captures essential representations for sequence/structure prediction. Moreover, the masked residue type update procedure helps ensure sequence-structure consistency. Additionally, the full-atom refinement module and the binding ligand are essential for modeling side-chain and pocket-ligand interactions, which play important roles in protein-ligand binding.
Further, we replaced the iterative full-shot refinement modules in FAIR with the autoregressive decoding method used in HSRN [28] as a variant of FAIR. However, this variant's performance is also worse than FAIR's. As discussed earlier, autoregressive generation can accumulate errors and is more time-consuming than the full-shot refinement in FAIR. For example, AAR drops from  $40.17\%$  to  $31.04\%$  on CrossDocked with autoregressive generation.
Table 2: Ablation studies of FAIR.
Model	CrossDocked	Binding MOAD
AAR (↑)	RMSD (↓)	Vina (↓)	AAR (↑)	RMSD (↓)	Vina (↓)
w/o hier encoder	25.30±9.86%	2.36±0.24	-6.694±2.28	29.21±10.8%	1.82±0.14	-7.235±2.10
w/o consistency	36.75±9.21%	1.49±0.08	-6.937±1.72	37.40±15.0%	1.49±0.12	-7.752±1.94
w/o full-atom	35.52±10.4%	1.62±0.13	-6.412±2.37	35.91±13.5%	1.54±0.12	-7.544±2.23
w/o ligand	33.16±15.3%	1.71±0.11	-6.353±1.91	34.54±14.6%	1.50±0.09	-7.271±1.79
autoregressive	31.04±14.2%	1.56±0.16	-6.874±1.90	33.65±12.7%	1.46±0.11	-7.765±1.84
FAIR	40.17±12.6%	1.42±0.07	-7.022±1.75	43.75±15.2%	1.35±0.10	-7.978±1.91
(b)
Figure 4: (a) The distributions of residue types in the designed pockets and the test sets. (b) The ratio of unconverged residues (residue types different from those in the final designed sequences) and AAR (Amino Acid Recovery) during the iterations of FAIR on the CrossDocked dataset. (c) Comparisons of average generation time for 100 pockets with baseline methods.
5 Limitations and Future Research
While we have made significant strides in addressing the limitations of prior research, several avenues remain for future research. First, we preliminarily formulate the protein pocket co-design problem in this paper. In the future, more realistic problem formulations may be constructed according to the requirements of different scenarios. Second, beyond the design of protein pockets, optimizing properties such as the binding affinity of existing pockets is crucial, especially for therapeutic applications. In forthcoming research, we aim to combine reinforcement learning and Bayesian optimization with FAIR to refine pocket design. Third, leveraging physics-based and template-matching methodologies can pave the way for the development of hybrid models that are both interpretable and generalizable. For instance, physics-based techniques can be employed to sample various ligand conformations during ligand initialization. The inherent protein-ligand interaction data within templates can guide pocket design. Lastly, integrating FAIR into protein/enzyme design workflows has the potential to be immensely beneficial for the scientific community. Conducting wet-lab experiments to assess the efficacy of protein pockets designed using FAIR would be advantageous. Insights gleaned from these experimental outcomes can be instrumental in refining our model, thereby establishing a symbiotic relationship between computation and biochemical experiments.
We develop a full-atom iterative refinement framework (FAIR) for protein pocket sequence and 3D structure co-design. Generally, FAIR has two refinement steps (backbone refinement and full-atom refinement) and follows a coarse-to-fine pipeline. The influence of side-chain atoms, the flexibility of binding ligands, and sequence-structure consistency are well considered and addressed. We empirically evaluate our method on cross-docked and experimentally determined datasets to show the advantage over existing physics-based, template-matching-based, and deep generative methods in efficiently generating high-quality pockets. We hope our work can inspire further explorations of pocket design with deep generative models and benefit the broad community.
This research was partially supported by grants from the National Key Research and Development Program of China (No. 2021YFF0901003) and the University Synergy Innovation Program of Anhui Province (GXXT-2021-002). We wish to express our sincere appreciation to Dr. Yaoxi Chen and Yifei Hu for their constructive discussions, which greatly enrich this research.
[1] Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew J O'Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, P Douglas Renfrew, Vikram K Mulligan, Kalli Kappel, et al. The rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation, 13(6):3031–3048, 2017.
[62] Ashish Vavwani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[63] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
[64] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. *ICLR*, 2022.
[65] Xihong Yang, Jiaqi Jin, Siwei Wang, Ke Liang, Yue Liu, Yi Wen, Suyuan Liu, Sihang Zhou, Xinwang Liu, and En Zhu. Dealmvc: Dual contrastive calibration for multi-view clustering. In Proceedings of the 31th ACM International Conference on Multimedia, 2023.
[66] Xihong Yang, Cheng Tan, Yue Liu, Ke Liang, Siwei Wang, Sihang Zhou, Jun Xia, Stan Z Li, Xinwang Liu, and En Zhu. Convert: Contrastive graph clustering with reliable augmentation. In Proceedings of the 31th ACM International Conference on Multimedia, 2023.
[67] Alexandre Zanghellini, Lin Jiang, Andrew M Wollacott, Gong Cheng, Jens Meiler, Eric A Althoff, Daniela Röthlisberger, and David Baker. New algorithms and an in silico benchmark for computational enzyme design. Protein Science, 15(12):2785-2794, 2006.
[68] Zaixi Zhang, Qi Liu, Qingyong Hu, and Chee-Kong Lee. Hierarchical graph transformer with adaptive node sampling. Advances in Neural Information Processing Systems, 2022.
[69] Zaixi Zhang, Qi Liu, Chee-Kong Lee, Chang-Yu Hsieh, and Enhong Chen. An equivariant generative framework for molecular graph-structure co-design. Chemical Science, 14(31):8380-8392, 2023.
[70] Zaixi Zhang, Yaosen Min, Shuxin Zheng, and Qi Liu. Molecule generation for target protein binding with structural motifs. In ICLR, 2023.
[71] Zaixi Zhang, Jiaxian Yan, Qi Liu, Enhong Chen, and Marinka Zitnik. A systematic survey in geometric deep learning for structure-based drug design. arXiv preprint arXiv:2306.11768, 2023.
A Details on FAIR Model
A.1 Training and Sampling Algorithms
Here, we show the pseudo-codes of the training and sampling of FAIR in the Algorithm. 1 and 2.
Algorithm 1 Training of FAIR
Input: protein sequences  $\mathcal{A} = a_1\cdots a_{N_s}$  and structures  $\{\pmb{x}(a_i)\}_{i=1}^{N_s}$ , ligand molecule  $\mathcal{M}$ . Initialize the coordinates of pocket residues  $\{\pmb{x}(b_i)\}$ . Add Gaussian noise to the ligand molecular structure  $\{x_j\}$ . Initialize the sequences of pocket residues with uniform distributions over 20 residue categories.  $\mathcal{L}_{pred} = 0$ ;  $\mathcal{L}_{struct} = 0$ . Backbone refinement for  $t$  in  $1, \cdots, T_1$ : Obtain the residue embeddings  $\{f_i\}$  and atom embeddings  $\{h_i\}$  with the hierarchical encoder. Predict residue types and update coordinates with structure refinement modules.  $\mathcal{L}_{seq} + = \sum_{i} l_{ce}(p_i^t, \hat{p}_i)$ .  $\mathcal{L}_{struct} + = \sum_{i} l_{huber}(\pmb{x}(b_i)^t, \hat{\pmb{x}}(\pmb{b}_i)) + \sum_{j} l_{huber}(\pmb{x}_j^t, \hat{\pmb{x}}_j)$ . end for Sample residue types from the predicted residue type distributions. Initialize the coordinates of side-chain atoms with the coordinates of corresponding  $C_\alpha$ . # Full-atom refinement for  $t$  in  $1, \cdots, T_2$ : Randomly mask a subset of pocket residues. Obtain the residue embeddings  $\{f_i\}$  and atom embeddings  $\{h_i\}$  with the hierarchical encoder. Predict the residue types of masked residues.  $\mathcal{L}_{seq} + = \sum_{i} l_{ce}(p_i^t, \hat{p}_i)$ . Reinitialize the masked residue if its residue type conflicts with the predicted result. Update coordinates with structure refinement modules.  $\mathcal{L}_{struct} + = \sum_{i} l_{huber}(\pmb{x}(b_i)^t, \hat{\pmb{x}}(\pmb{b}_i)) + \sum_{j} l_{huber}(\pmb{x}_j^t, \hat{\pmb{x}}_j)$ . end for Minimize  $\frac{1}{T} (\mathcal{L}_{seq} + \mathcal{L}_{struct})$ .
Algorithm 2 Sampling of FAIR
Input: protein sequences  $\mathcal{A} = a_1\cdot \cdot \cdot a_{N_s}$  and structures  $\{\pmb {x}(a_i)\}_{i = 1}^{N_s}$  , ligand molecule  $\mathcal{M}$  Initialize the coordinates of pocket residues  $\{\pmb {x}(\pmb {b}_i)\}$  . Initialize the sequences of pocket residues with uniform distributions over 20 residue categories. # Backbone refinement   
for  $t$  in  $1,\dots T_{1}$  .. Obtain the residue embeddings  $\{f_i\}$  and atom embeddings  $\{h_i\}$  with the hierarchical encoder. Predict residue types and update coordinates with structure refinement modules.   
end for   
Sample residue types from the predicted residue type distributions.. Initialize the coordinates of side-chain atoms with the coordinates of corresponding  $C_\alpha$  # Full-atom refinement   
for  $t$  in  $1,\dots T_{2}$  .. Randomly mask a subset of pocket residues. Obtain the residue embeddings  $\{f_i\}$  and atom embeddings  $\{h_i\}$  with the hierarchical encoder. Predict the residue types of masked residues. Reinitialize the masked residue if its residue type conflicts with the predicted result. Update coordinates with structure refinement modules.   
end for
A.2 Additional Information on Hierarchical Graph Transformer Encoder