SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zaixi Zhang $^{1,2,4}$ , Zepu Lu $^{1,2}$ , Zhongkai Hao $^{3}$ , Marinka Zitnik $^{4}$ , Qi Liu $^{1,2*}$
<sup>1</sup> Anhui Province Key Lab of Big Data Analysis and Application, School of Computer Science and Technology, University of Science and Technology of China
$^{2}$  State Key Laboratory of Cognitive Intelligence, Hefei, Anhui, China
$^{3}$  Dept. of Comp. Sci. and Tech., Institute for AI, THBI Lab, BNRist Center,
Tsinghua-Bosch Joint ML Center, Tsinghua, Harvard University
{zaixi, zplu}  $@$  mail.ustc.edu.cn, hzj21  $@$  mails.tsinghua.edu.cn,
marinka@hms.harvard.edu, qiliuql@ustc.edu.cn
The design and optimization of functional proteins that bind specific ligand molecules is paramount in therapeutics and bio-engineering. A critical yet formidable task in this endeavor is the design of the protein pocket, which is the cavity region of the protein where the ligand binds. Current methods are plagued by inefficient generation, inadequate context modeling of the ligand molecule, and the inability to generate side-chain atoms. Here, we present the Full-Atom Iterative Refinement (FAIR) method, designed to address these challenges by facilitating the co-design of protein pocket sequences, specifically residue types, and their corresponding 3D structures. FAIR operates in two steps, proceeding in a coarse-to-fine manner (transitioning from protein backbone to atoms, including side chains) for a full-atom generation. In each iteration, all residue types and structures are simultaneously updated, a process termed full-shot refinement. In the initial stage, the residue types and backbone coordinates are refined using a hierarchical context encoder, complemented by two structure refinement modules that capture both inter-residue and pocket-ligand interactions. The subsequent stage delves deeper, modeling the side-chain atoms of the pockets and updating residue types to ensure sequence-structure congruence. Concurrently, the structure of the binding ligand is refined across iterations to accommodate its inherent flexibility. Comprehensive experiments show that FAIR surpasses existing methods in designing superior pocket sequences and structures, producing average improvement exceeding  $10\%$  in AAR and RMSD metrics.
Proteins are macromolecules that perform fundamental functions in living organisms. An essential and challenging step in functional protein design, such as enzymes [52, 27] and biosensors [5, 18], is to design protein pockets that bind specific ligand molecules. Protein pockets usually refer to residues binding with the ligands in the protein-ligand complexes. Nevertheless, such a problem is very challenging due to the tremendous search space of both sequence and structure and the small solution space that satisfies the constraints of binding affinity, geometric complementarity, and biochemical stability. Traditional approaches for pocket design use hand-crafted energy functions [42, 11] and template-matching methods [67, 8, 48], which can be less accurate and lack efficiency.
Deep learning approaches, especially deep generative models, have made significant progress in protein design [46, 17, 9, 20] using experimental and predicted 3D structure information to design amino acid sequences [26, 30, 7, 14]. Despite their strong performance, these approaches are not directly applicable to pocket design because the protein pocket structures to be designed are apriori unknown [8]. To this end, generative models with the capability to co-design both the sequence and structure of protein pockets are needed. In another area of protein design, pioneering research developed techniques to co-design sequence and structure of complementarity determining regions (CDRs) for antibodies (Y-shaped proteins that bind with specific pathogens) [29, 28, 39, 33, 56]. Nevertheless, these approaches are tailored specifically for antibodies and are unsuitable for the pocket design of general proteins. For example, prevailing methods generate only protein backbone atoms while neglecting side-chain atoms, which play essential roles in protein pocket-ligand interactions [43]. Further, binding targets (antigens) considered by these methods are fixed, whereas binding ligand molecules are more flexible and require additional modeling.
Here, we introduce a full-atom iterative refinement approach (FAIR) for protein pocket sequence-structure co-design, addressing the limitations of prior research. FAIR consists of two primary steps: initially modeling the pocket backbone atoms to generate coarse-grained structures, followed by incorporating side-chain atoms for a complete full-atom pocket design. (1) In the first step, since the pocket residue types and the number of side-chain atoms are largely undetermined, FAIR progressively predicts and updates backbone atom coordinates and residue types via soft assignment over multiple rounds. Sequence and structure information are jointly updated in each round for efficiency, which we refer to as full-shot refinement. A hierarchical encoder with residue- and atom-level encoders captures the multi-scale structure of protein-ligand complexes. To update the pocket structure, FAIR leverages two structure refinement modules to predict the interactions within pocket residues and between the pocket and the binding ligand in each round. (2) In the second step, FAIR initializes residue types based on hard assignment and side-chain atoms based on the results of the first step. The hierarchical encoder and structure refinement modules are similar to the first step for the full-atom structure update. To model the influence of structure on residue types, FAIR randomly masks a part of the pocket residues and updates the masked residue types based on the neighboring residue types and structures in each round. After several rounds of refinement, the residue type and structure update gradually converge, and the residue sequence-structure becomes consistent. The ligand molecular structure is also updated along with the above refinement processes to account for the flexibility of the ligand.
Overall, FAIR is an E(3)-equivariant generative model that generates. both pocket residue types and the pocket-ligand complex 3D structure. Our study presents the following main contributions:
- New task: We preliminarily formulate protein pocket design as a co-generation of both the sequence and full-atom structure of a protein pocket (i.e., co-design), conditioned on a binding ligand molecule and the protein backbone (see Sec. 3.1).
- Novel method: FAIR is an end-to-end generative framework that generates both the pocket sequence and structure through iterative refinement (see Sec. 3). It addresses the limitations of previous approaches and takes into account side-chain effects, ligand flexibility, and sequence-structure consistency to enhance prediction efficiency.
- **Strong performance:** FAIR outperforms existing methods on various protein pocket design quality metrics, producing average improvements of  $15.5\%$  (AAR) and  $13.5\%$  (RMSD). FAIR's pocket-generation process is over ten times faster than traditional methods (see Sec. 4).
Protein Design. Computational protein design has brought the capability of inventing proteins with desired structures and functions [22, 35, 60]. Existing studies focus on designing amino acid sequences that fold into target protein structures, i.e., inverse protein folding. Traditional approaches mainly rely on hand-crafted energy functions to iteratively search for low-energy protein sequences with heuristic sampling algorithms, which are inefficient, less accurate, and less generalizable [1, 6]. Recent deep generative approaches have achieved remarkable success in protein design [15, 26, 30, 7, 14, 51, 20, 3]. Some protein design methods also benefit from the latest advances in protein structure prediction techniques, e.g., AlphaFold2 [31]. Nevertheless, they usually assume the target protein structure is given, which is unsuitable for our pocket design case. Another
(b)
Full-atom structure including sidechains
Pocket residue type sequence: MTREIALE
Figure 1: (a) The residue structure, where the backbone atoms are  $C_{\alpha}, N, C, O$ .  $R$  represents a side chain that determines the residue types. (b) The protein pocket design problem. Pocket (colored in black) consists of a subsequence of residues closest to the binding ligand molecule (formal definition in Sec. 3.1). FAIR aims to co-design the full-atom structures and residue types of the pocket.
line of recent works proposes to co-design the sequence and 3D structure of antibodies as graphs [29, 28, 39, 33, 56]. However, they are designed explicitly for antibody CDR topologies and are primarily unable to consider side-chain atoms and the flexibility of binding targets. For more detailed discussions on protein design, interested readers may refer to the comprehensive surveys [22, 35].
Pocket Design. Protein pockets are the cavity region of protein where the ligand binds to the protein. In projects to design functional proteins such as biosensors [52, 27] and enzymes [5, 18], pocket design is the central task [8, 48]. Compared with the general protein design/antibody design problems, the pocket design brings unique challenges: (1) protein pocket has a smaller spatial scale than whole proteins/antibodies. Therefore, the structure and interactions of side-chain atoms play essential roles in pocket design and require careful modeling [36, 12]; (2) the influence and structural flexibility of target ligand molecules need to be considered to achieve high binding affinity [42]. Traditional methods for pocket design are physics-based methods [42, 11] or template-matching methods [67, 8, 48], which often suffer from the large computational complexity and limited accuracy. So far, the pocket design problem has rarely been tackled with deep learning.
Structure-Based Drug Design. Structure-based drug design is another line of research inquiry that explores 3D ligand molecule design inside protein pockets [49, 40, 37, 47, 70, 69, 19]. These studies can be regarded as the dual problem of pocket design studied in our paper. However, the methods primarily focus on generating 3D molecular graphs based on a fixed protein pocket structure and may not be easily adapted to the pocket sequence-structure co-design problem examined in our study. We refer interested readers to [71] for a comprehensive survey on structure-based drug design.
3 Full-Atom Protein Pocket Design
Next, we introduce FAIR, our method for pocket sequence-structure co-design. The problem is first formulated in Sec. 3.1. We then describe the hierarchical encoder that captures residue- and atom-level information in Sec. 3.2. Furthermore, we show the iterative refinement algorithms with two main steps in Sec. 3.3. Finally, we show the procedures for model training and sampling in Sec. 3.4 and demonstrate the E(3)-equivariance property of FAIR in Sec. 3.5.
3.1 Notation and Problem Statement
We aim to co-design residue types (sequences) and 3D structures of the protein pocket that can fit and bind with the target ligand molecule. As shown in Figure 1, a protein-ligand binding complex typically consists of a protein and a ligand molecule. The ligand molecule can be represented as a 3D point cloud  $\mathcal{M} = \{(x_k,v_k)\}_{k = 1}^{N_l}$  where  $x_{k}$  denotes the 3D coordinates and  $v_{k}$  denotes the atom feature. The protein can be represented as a sequence of residues  $\mathcal{A} = \pmb {a}_1\dots \pmb {a}_{N_s}$ , where  $N_{s}$  is the length of the sequence. The 3D structure of the protein can be described as a point cloud of protein atoms  $\{\pmb{a}_{i,j}\}_{1\leq i\leq N_s,1\leq j\leq n_i}$  and let  $\pmb {x}(\pmb{a}_{i,j})\in \mathbb{R}^3$  denote the 3D coordinate of the pocket atoms.  $n_i$  is the number of atoms in a residue determined by the residue types. The first four atoms in any residue correspond to its backbone atoms  $(C_{\alpha},N,C,O)$ , and the rest are its side-chain atoms. In our work, the protein pocket is defined as a set of residues in the protein closest to the binding ligand
molecule:  $\mathcal{B} = \pmb{b}_1\cdots\pmb{b}_m$ . The pocket  $\mathcal{B}$  can thus be represented as an amino acid subsequence of a protein:  $\mathcal{B} = \pmb{a}_{e_1}\cdots\pmb{a}_{e_m}$  where  $\pmb{e} = \{e_1,\dots ,e_m\}$  is the index of the pocket residues in the protein. The index  $\pmb{e}$  can be formally given as:  $\pmb{e} = \{i\mid \min_{1\leq j\leq n_i,1\leq k\leq N_l}\| \pmb{x}(\pmb{a}_{i,j}) - \pmb{x}_k\| _2\leq \delta \}$ , where  $\| \cdot \| _2$  is the  $L_{2}$  norm and  $\delta$  is the distance threshold. According to the distance range of pocket-ligand interactions [43], we set  $\delta = 3.5\AA$  in the default setting. Note that the number of residues (i.e.,  $m$ ) varies for pockets due to varying structures of protein-ligand complexes. To summarize this subsection, our preliminary problem statement is formulated as below:
Problem Statement: In FAIR, our objective is to learn a conditional generative model for full-atom protein pocket generation conditioned on the protein backbone and the binding ligand. Given the availability of existing protein pockets, we may leverage them for reference to construct residue graphs and calculate geometric features. Our preliminary problem statement is formulated considering practical situations such as enzyme optimization.
3.2 Hierarchical Encoder
Leveraging this intrinsic multi-scale protein structure [59], we adopt a hierarchical graph transformer [44, 38, 68] to encode the hierarchical context information of protein-ligand complexes for pocket sequence-structure co-design. The transformer includes both atom- and residue-level encoders, as described below.
3.2.1 Atom-Level Encoder
The atom-level encoder considers all atoms in the protein-ligand complex and constructs a 3D context graph connecting the  $K_{a}$  nearest neighboring atoms in  $\mathcal{A} \bigcup \mathcal{M}$ . The atomic attributes are firstly mapped to node embeddings  $h_{i}^{(0)}$  with two MLPs, respectively, for protein and ligand molecule atoms. The edge embeddings  $e_{ij}$  are obtained by processing pairwise Euclidean distances with Gaussian kernel functions [53]. The 3D graph transformer consists of  $L$  Transformer layers [62] with two modules in each layer: a multi-head self-attention (MHA) module and a position-wise feed-forward network (FFN). Specifically, in the MHA module of the  $l$ -th layer ( $1 \leq l \leq L$ ), the queries are derived from the current node embeddings  $h_{i}^{(l)}$  while the keys and values come from the relational representations:  $r_{ij}^{(l)} = [h_{j}^{(l)}, e_{ij}^{(l)}] ([\cdot, \cdot])$  denotes concatenation) from neighboring nodes:
$$
\boldsymbol {q} _ {i} ^ {(l)} = \mathbf {W} _ {Q} \boldsymbol {h} _ {i} ^ {(l)}, \boldsymbol {k} _ {i j} ^ {(l)} = \mathbf {W} _ {K} \boldsymbol {r} _ {i j} ^ {(l)}, \boldsymbol {v} _ {i j} ^ {(l)} = \mathbf {W} _ {V} \boldsymbol {r} _ {i j} ^ {(l)}, \tag {1}
$$
where  $\mathbf{W}_Q, \mathbf{W}_K$  and  $\mathbf{W}_V$  are learnable transformation matrices. Then, in each attention head  $c \in \{1, 2, \dots, C\}$  ( $C$  is the total number of heads), the scaled dot-product attention mechanism works as follows:
$$
\operatorname {h e a d} _ {i} ^ {c} = \sum_ {j \in \mathcal {N} (i)} \operatorname {S o f t m a x} \left(\frac {\boldsymbol {q} _ {i} ^ {(l) ^ {\top}} \cdot \boldsymbol {k} _ {i j} ^ {(l)}}{\sqrt {d}}\right) \boldsymbol {v} _ {i j} ^ {(l)}, \tag {2}
$$
where  $\mathcal{N}(i)$  denotes the neighbors of the  $i$ -th atom in the constructed graph and  $d$  is the dimension size of embeddings for each head. Finally, the outputs from different heads are further concatenated and transformed to obtain the final output of MHA:
$$
\mathrm {M H A} _ {i} = \left[ \operatorname {h e a d} _ {i} ^ {1}, \dots , \operatorname {h e a d} _ {i} ^ {C} \right] \mathbf {W} _ {O}, \tag {3}
$$
where  $\mathbf{W}_O$  is the output transformation matrix. The output of the atom-level encoder is a set of atom representations  $\{\pmb{h}_i\}$ . Appendix A includes the details of the atom-level encoder.
3.2.2 Residue-Level Encoder
The residue-level encoder only keeps the  $C_{\alpha}$  atoms of residues and a coarsened ligand node at the ligand's center of mass to supplement binding context information (the coarsened ligand node is appended at the end of the residue sequence as a special residue). Then, a  $K_{r}$  nearest neighbor graph is constructed at the residue level. We use the original protein pocket backbone atoms for reference to construct the  $K_{r}$  nearest neighbor graph and help calculate geometric features. The  $i$ -th residue  $\pmb{a}_{i}$  can be represented by a feature vector  $\pmb{f}_{\pmb{a}_{i}}^{res}$  describing its geometric and chemical characteristics such as volume, polarity, charge, hydropathy, and hydrogen bond interactions according to its residue
Figure 2: Overview of FAIR with two main steps ( $T_{1}$  rounds of backbone refinement and  $T_{2}$  rounds of full-atom refinement). FAIR co-designs pocket residue types and structures via iterative full-shot refinement. The structure refinement is illustrated with the atoms in the dotted ovals as examples.
type  $\pmb{a}_i$ . We concatenate the residue features with the sum of atom-level embeddings  $h_k$  within each residue to initialize the residue representations; the ligand node representation is initialized by summing up all the ligand atom embeddings:
$$
\hat {\boldsymbol {f}} _ {i} ^ {r e s} = \left[ \boldsymbol {f} _ {\boldsymbol {a} _ {i}} ^ {r e s}, \sum_ {j \in r e s _ {i}} \boldsymbol {h} _ {j} \right]; \quad \hat {\boldsymbol {f}} _ {i} ^ {l i g} = \sum_ {k \in \mathcal {M}} \boldsymbol {h} _ {k}. \tag {4}
$$
As for the edge features  $e_{ij}^{res}$ , local coordinate frames are built for residues, and  $e_{ij}^{res}$  are computed according to the distance, direction, and orientation between neighboring residues [26]. Lastly, the encoder takes the node and edge features into the residue-level graph transformer to compute the residue-level representations. The residue-level graph transformer architecture is similar to that of the atom-level encoder, and details are shown in Appendix A.
In summary, the output of our hierarchical encoder is a set of residue representations  $\{\pmb{f}_i\}$  and atom representations  $\{\pmb{h}_i\}$ , which are used in the following subsections for structure refinement and residue type prediction. Since our encoder is based on the atom/residue attributes and pairwise relative distances, the obtained representations are E(3)-invariant<sup>2</sup>.
3.3 Iterative Refinement
For efficiency, FAIR predicts pocket residue types and structures via an iterative full-shot refinement scheme: all the pocket residue types and structures are updated together in each round. Since the pocket residue types and the number of side-chain atoms are largely unknown at the beginning rounds, FAIR is designed to have two main steps that follow a coarse-to-fine pipeline: FAIR firstly only models the backbone atoms of pockets to generate the coarse-grained structures and then fine-adjusts full-atom residues to achieve sequence-structure consistency. Figure 2 shows the overview of FAIR.
3.3.1 Backbone Refinement
Initialization. Before backbone refinement, the residue types to predict are initialized with uniform distributions over all residue type categories. Since the residue types and the number of side-chain atoms are unknown, we only initialize the 3D coordinates of backbone atoms  $(C_{\alpha},N,C,O)$  with linear interpolations and extrapolations based on the known structures of nearest residues in the protein. Details are shown in the Appendix A.4. Following previous research [8, 28, 33], we initialize the ligand structures with the reference structures from the dataset in the generation process.
Residue Type Prediction. In each round of iterative refinement, we first encode the current protein-ligand complex with the hierarchical encoder in Sec. 3.2. We then predict and update the sequences and structures. To predict the probability of each residue type, we apply an MLP and SoftMax function on the obtained pocket residue embeddings  $\pmb{f}_i : p_i = \mathrm{Softmax}(\mathrm{MLP}(\pmb{f}_i))$ , where  $p_i \in \mathbb{R}^{20}$  is the predicted distribution over 20 residue categories. The corresponding residue features can then be updated with weighted averaging:  $\pmb{f}_{\pmb{b}_i}^{res} \gets \sum_j p_{i,j} \pmb{f}_j^{res}$ , where  $p_{i,j}$  and  $\pmb{f}_j^{res}$  are the probability and feature vector for the  $j$ -th residue category respectively.
Structure Refinement. Motivated by the force fields in physics [50] and the successful practice in previous works [28, 33], we calculate the interactions between atoms for the coordinate update.
$^{2}\mathrm{E}(3)$  is the group of Euclidean transformations: rotations, reflections, and translations.
According to previous works [32, 12, 10], the internal interactions within the protein and the external interactions between protein and ligand are different. Therefore, FAIR uses two separate modules for interaction predictions. Firstly, the internal interactions are calculated based on residue embeddings:
$$
\boldsymbol {g} _ {i j, k j} = g \left(\boldsymbol {f} _ {i}, \boldsymbol {f} _ {k}, \boldsymbol {e} _ {\mathrm {R B F}} \left(\| \hat {\boldsymbol {x}} \left(\boldsymbol {b} _ {i, 1}\right) - \boldsymbol {x} \left(\boldsymbol {a} _ {k, 1}\right) \|\right)\right) _ {j} \cdot \frac {\boldsymbol {x} \left(\boldsymbol {b} _ {i , j}\right) - \boldsymbol {x} \left(\boldsymbol {a} _ {k , j}\right)}{\left\| \boldsymbol {x} \left(\boldsymbol {b} _ {i , j}\right) - \boldsymbol {x} \left(\boldsymbol {a} _ {k , j}\right) \right\|}, \tag {5}
$$
where  $e_{\mathrm{RBF}}(\| \hat{\pmb{x}} (\pmb {b}_{i,1}) - \pmb {x}(\pmb {a}_{k,1})\|)$  is the radial basis distance encoding of pairwise  $C_\alpha$  distances, the function  $g$  is a feed-forward neural network with four output channels for four backbone atoms, and  $\frac{\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}(\pmb{a}_{k,j})}{\|\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}(\pmb{a}_{k,j})\|}\in \mathbb{R}^3$  is the normalized vector from the residue atom  $\pmb{a}_{k,j}$  to  $\pmb{b}_{i,j}$ . For the stability of training and generation, we use original protein backbone coordinates to calculate the RBF encodings. Similarly, for the external interactions between pocket  $\mathcal{B}$  and ligand  $\mathcal{M}$ , we have:
$$
\boldsymbol {g} _ {i j, s} = g ^ {\prime} \left(\boldsymbol {h} _ {\boldsymbol {b} _ {i, j}}, \boldsymbol {h} _ {s}, \boldsymbol {e} _ {\mathrm {R B F}} \left(\left\| \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) - \boldsymbol {x} _ {s} \right\|\right)\right) \cdot \frac {\boldsymbol {x} \left(\boldsymbol {b} _ {i , j}\right) - \boldsymbol {x} _ {s}}{\left\| \boldsymbol {x} \left(\boldsymbol {b} _ {i , j}\right) - \boldsymbol {x} _ {s} \right\|}, \tag {6}
$$
which is calculated based on the atom representations.  $g'$  is a feed-forward neural network with one output channel. For the conciseness of expression, here we use  $h_{\pmb{b}_{i,j}}$  and  $h_s$  to denote the pocket and ligand atom embeddings, respectively. Finally, the coordinates of pocket atoms  $x(\pmb{b}_{i,j})$  and ligand molecule atoms  $x_s$  are updated as follows:
$$
\boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) \leftarrow \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) + \frac {1}{K _ {r}} \sum_ {k \in \mathcal {N} (i)} \boldsymbol {g} _ {i j, k j} + \frac {1}{N _ {l}} \sum_ {s} \boldsymbol {g} _ {i j, s}, \tag {7}
$$
$$
\boldsymbol {x} _ {s} \leftarrow \boldsymbol {x} _ {s} - \frac {1}{| \mathcal {B} |} \sum_ {i, j} \boldsymbol {g} _ {i j, s}, 1 \leq i \leq m, 1 \leq j \leq 4, 1 \leq s \leq N _ {l}, \tag {8}
$$
where  $\mathcal{N}(i)$  is the neighbors of the constructed residue-level  $K_{r}$ -NN graph;  $|\mathcal{B}|$  denotes the number of pocket atoms. Details of the structure refinement modules are shown in the Appendix A.
3.3.2 Full-Atom Refinement