SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our hierarchical encoder includes the atom-level encoder and the residue-level encoder. The powerful 3D graph transformer is used as the model backbone. Here, we first provide additional information on the graph transformer. Residue-level node and edge feature representations follow prior research [26, 28].
Graph transformer architecture. The atom/residue-level encoder contains  $L$  graph transformer layers. Let  $h^{(l)}$  be the set of node representations at the  $l$ -th layer. In each graph transformer layer, there is a multi-head self-attention (MHA) and a feed-forward block (FFN). The layer normalization (LN) is applied before the two blocks [63, 68]. The details of MHA have been shown in Sec 3.2, and the graph transformer layer is formally characterized as:
$$
\boldsymbol {h} ^ {\prime (l - 1)} = \operatorname {M H A} \left(\ln \left(\boldsymbol {h} ^ {(l - 1)}\right)\right) + \boldsymbol {h} ^ {(l - 1)} \tag {11}
$$
$$
\boldsymbol {h} ^ {(l)} = \operatorname {F F N} \left(\ln \left(\boldsymbol {h} ^ {\prime (l - 1)}\right)\right) + \boldsymbol {h} ^ {\prime (l - 1)}, \quad (0 \leq l <   L). \tag {12}
$$
Residue-level node features. The residue-level encoder only keeps the  $C_{\alpha}$  atoms to represent residues and constructs a  $K_{r}$  nearest neighbor graph 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. Each residue node is represented by six features: polarity  $f_{p} \in \{0,1\}$ , hydropathy  $f_{h} \in [-4.5,4.5]$ , volume  $f_{v} \in [60.1,227.8]$ , charge  $f_{c} \in \{-1,0,1\}$ , and whether it is a hydrogen bond donor  $f_{d} \in \{0,1\}$  or acceptor  $f_{a} \in \{0,1\}$ . We expand hydropathy and volume features into radial basis with interval sizes 0.1 and 10, respectively. Overall, the dimension of the residue-level feature vector  $f_{i}$  is 112.
Residue-level edge features. For the  $i$ -th residue, we let  $\pmb{x}(\pmb{a}_{i,1})$  denote the coordinate of its  $\mathbf{C}_{\alpha}$  and define its local coordinate frame  $\pmb{O}_i = [c_i, n_i, c_i \times n_i]$  as:
$$
\boldsymbol {u} _ {i} = \frac {\boldsymbol {x} \left(\boldsymbol {a} _ {i , 1}\right) - \boldsymbol {x} \left(\boldsymbol {a} _ {i - 1 , 1}\right)}{\left\| \boldsymbol {x} \left(\boldsymbol {a} _ {i , 1}\right) - \boldsymbol {x} \left(\boldsymbol {a} _ {i - 1 , 1}\right) \right\|}, \quad \boldsymbol {c} _ {i} = \frac {\boldsymbol {u} _ {i} - \boldsymbol {u} _ {i + 1}}{\left\| \boldsymbol {u} _ {i} - \boldsymbol {u} _ {i + 1} \right\|}, \quad \boldsymbol {n} _ {i} = \frac {\boldsymbol {u} _ {i} \times \boldsymbol {u} _ {i + 1}}{\left\| \boldsymbol {u} _ {i} \times \boldsymbol {u} _ {i + 1} \right\|}. \tag {13}
$$
Based on the local frame, the edge features between residues  $i$  and  $j$  can be computed as:
$$
\boldsymbol {e} _ {i j} ^ {r e s} = \operatorname {C o n c a t} \left(E _ {\text {p o s}} (i - j), \boldsymbol {e} _ {\text {R B F}} (\| \boldsymbol {x} (\boldsymbol {a} _ {i, 1}) - \boldsymbol {x} (\boldsymbol {a} _ {j, 1}) \|), \boldsymbol {O} _ {i} ^ {\top} \frac {\boldsymbol {x} (\boldsymbol {a} _ {j , 1}) - \boldsymbol {x} (\boldsymbol {a} _ {i , 1})}{\| \boldsymbol {x} (\boldsymbol {a} _ {i , 1}) - \boldsymbol {x} (\boldsymbol {a} _ {j , 1}) \|}, \boldsymbol {q} (\boldsymbol {O} _ {i} ^ {\top} \boldsymbol {O} _ {j})\right). \tag {14}
$$
The edge feature  $e_{ij}^{res}$  contains four parts. The positional encoding  $E_{\mathrm{pos}}(i - j)$  encodes the relative sequence distance between two residues. The second term  $e_{\mathrm{RBF}}(\cdot)$  is a distance encoding with radial basis functions. The third term is a direction encoding corresponding to the relative direction of  $x(a_{j,1})$  in the local frame of the  $i$ -th residue. The last term  $q(O_i^\top O_j)$  is the orientation encoding of the quaternion representation  $q(\cdot)$  of the spatial rotation matrix  $O_i^\top O_j$  [25]. Overall, the dimension of the residue-level edge feature  $e_{ij}^{res}$  is 39.
A.3 More Details of the Full-atom Structure Refinement
The full-atom structure refinement is similar to the backbone structure refinement in Sec. 3.3.1 with minor modifications to accommodate side-chain atoms. We consider the internal interactions within the protein and the external interactions between ligand molecules and protein residues. Firstly, the internal interactions within the protein 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) - \bar {\boldsymbol {x}} \left(\boldsymbol {a} _ {k}\right)}{\left\| \boldsymbol {x} \left(\boldsymbol {b} _ {i , j}\right) - \bar {\boldsymbol {x}} \left(\boldsymbol {a} _ {k}\right) \right\|}, \tag {15}
$$
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 14 channels for all side-chain atoms (the maximum number of atoms in a residue is 14). Due to the different numbers of side-chain atoms in different kinds of residues, the normalized direction vector here is  $\frac{\pmb{x}(\pmb{b}_{i,j}) - \bar{\pmb{x}}(\pmb{a}_k)}{\|\pmb{x}(\pmb{b}_{i,j}) - \bar{\pmb{x}}(\pmb{a}_k)\|}\in \mathbb{R}^3$  where  $\bar{\pmb{x}} (\pmb {a}_k)$  is the average coordinates of residue  $\pmb{a}_k$ . The external interactions between pocket  $\mathcal{B}$  and ligand  $\mathcal{M}$  is the same as the backbone refinement:
$$
\boldsymbol {g} _ {i j, s} = g ^ {\prime} \left(\boldsymbol {h} _ {\boldsymbol {b} _ {i, j}}, \boldsymbol {h} _ {s}, \boldsymbol {e} _ {\mathrm {R B F}} \left(\| \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) - \boldsymbol {x} _ {s} \|\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 {16}
$$
which is calculated based on the atom representations. For the conciseness of expression, here we use  $h_{b_{i,j}}$  and  $h_s$  to denote the pocket and ligand atom embeddings, respectively. Finally, the coordinates of pocket atoms  $x(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 {17}
$$
$$
\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 n _ {i}, 1 \leq s \leq N _ {l}, \tag {18}
$$
where  $\mathcal{N}(i)$  indicates the neighbors in the constructed residue-level  $K_{r}$ -NN graph;  $|\mathcal{B}|$  denotes the number of pocket atoms.
A.4 Structure initialization based on interpolation and extrapolation
Figure 5: Structure initialization based on interpolation and extrapolation
Residues with known structure
Linear extrapolation
Figure 5 illustrates our structure initialization strategy. We initialize the residue coordinates with linear interpolations and extrapolations based on the nearest residues with known structures in the protein. Denote the sequence of residues as  $\mathcal{A} = \pmb{a}_1\cdots\pmb{a}_{N_s}$ , where  $N_{s}$  is the length of the sequence. Let  $\pmb{Z}(\pmb{a}_i)\in \mathbb{R}^{3\times 4}$  denote the backbone coordinates of the  $i$ -th residue. If we want to initialize the coordinates of the  $i$ -th residue, we take the following initialization strategy: (1) We use linear interpolation if there are residues with known coordinates at both sides of the  $i$ -th residue. Specifically, assume  $p$  and  $q$  are the indexes of the nearest residues with known coordinates at each side of the  $i$ -th residue ( $p < i < q$ ), we have:  $\pmb{Z}(\pmb{a}_i) = \frac{1}{q - p}[(i - p)\pmb{Z}(\pmb{a}_q) + (q - i)\pmb{Z}(\pmb{a}_p)]$ . (2) We conduct linear extrapolation if the  $i$ -th residue is at the ends of the chain, i.e., no residues with known structures at one side of the  $i$ -th residue. Specifically, let  $p$  and  $q$  denote the index of the nearest and the second nearest residue with known coordinates. The position of the  $i$ -th residue can be initialized as  $\pmb{Z}(\pmb{a}_i) = \pmb{Z}(\pmb{a}_p) + \frac{i - p}{p - q}(\pmb{Z}(\pmb{a}_p) - \pmb{Z}(\pmb{a}_q))$ . We initialize the side-chain atoms' coordinates with the coordinate of their corresponding  $C_{\alpha}$ . In practice, we can further apply local geometric constraints to the initialized backbone atoms in each residue.
A.5 Equivariance.
FAIR has the desirable property of E(3)-equivariance as follow:
Theorem 1. Denote the  $E(3)$ -transformation as  $T_{g}$  and the generative process of FAIR as  $\{(\pmb {b}_i,\pmb {x}(\pmb {b}_i))\}_{i = 1}^m = p_\theta (\mathcal{A}\setminus \mathcal{B},\mathcal{M})$ , where  $\{(\pmb {b}_i,\pmb {x}(\pmb {b}_i))\}_{i = 1}^m$  indicates the designed pocket sequence and structure. We have  $\{(b_i,T_g\cdot x(b_i))\}_{i = 1}^m = p_\theta (T_g\cdot (\mathcal{A}\setminus \mathcal{B}),T_g\cdot \mathcal{M})$
The E(3)-transformation on the euclidean coordinate  $\pmb{x} \in \mathbb{R}^3$  can be represented as:  $T_g \cdot \pmb{x} = O\pmb{x} + t$ , where  $O \in \mathbb{R}^3$  is the orthogonal transformation matrix,  $t \in \mathbb{R}^3$  is the translation vector. Before we give the proof, we present the following lemmas.
Lemma 1. The hierarchical graph transformer encoder is  $E(3)$ -invariant.
Proof. As shown in Sec. 3.2, the input to the hierarchical graph transformer encoders are scalar node features and pairwise distances, and the output is atom and residue scalar features. Therefore, the hierarchical encoder is E(3)-invariant by construction.  $\square$
Proof. Firstly, with the E(3)-invariant encoder, the obtained residue embeddings  $\{\pmb{f}_i\}$ , atom embeddings  $\{\pmb{h}_i\}$ , and the output of functions  $g$  and  $g'$  are E(3)-invariant. Furthermore,  $\forall T_g \in \mathrm{E}(3)$ , it is easy to have  $T_g \cdot (\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}(\pmb{a}_{k,j})) = O\pmb{x}(\pmb{b}_{i,j}) + t - O\pmb{x}(\pmb{a}_{k,j}) - t = O(\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}(\pmb{a}_{k,j}))$ . Similarly, we have  $\bar{T}_g \cdot (\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}_s) = O(\pmb{x}(\pmb{b}_{i,j}) - \pmb{x}_s)$  and  $T_g \cdot (\pmb{x}(\pmb{b}_{i,j}) - \bar{\pmb{x}}(\pmb{a}_k)) = O(\pmb{x}(\pmb{b}_{i,j}) - \bar{\pmb{x}}(\pmb{a}_k))$ . Let  $\pmb{x}(\pmb{b}_{i,j})'$  denote the updated pocket coordinates in Equation 7 and 17. It is then possible to derive that the updated pocket coordinates are E(3)-equivariant as follows:
$$
T _ {g} \cdot \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) + T _ {g} \cdot \frac {1}{K _ {r}} \sum_ {k \in \mathcal {N} (i)} \boldsymbol {g} _ {i j, k j} + T _ {g} \cdot \frac {1}{N _ {l}} \sum_ {s} \boldsymbol {g} _ {i j, s} \tag {19}
$$
$$
= O \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) + t + O \frac {1}{K _ {r}} \sum_ {k \in \mathcal {N} (i)} \boldsymbol {g} _ {i j, k j} + O \frac {1}{N _ {l}} \sum_ {s} \boldsymbol {g} _ {i j, s} \tag {20}
$$
$$
= O \left[ \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} \right] + t \tag {21}
$$
$$
= T _ {g} \cdot \boldsymbol {x} \left(\boldsymbol {b} _ {i, j}\right) ^ {\prime}. \tag {22}
$$
Similarly, for the updated ligand coordinates, we have:
$$
T _ {g} \cdot \boldsymbol {x} _ {s} - T _ {g} \cdot \frac {1}{| \mathcal {B} |} \sum_ {i, j} \boldsymbol {g} _ {i j, s} \tag {23}
$$
$$
= O \boldsymbol {x} _ {s} + t - O \frac {1}{| \mathcal {B} |} \sum_ {i, j} \boldsymbol {g} _ {i j, s} \tag {24}
$$
$$
= O \left[ \boldsymbol {x} _ {s} - \frac {1}{| \mathcal {B} |} \sum_ {i, j} \boldsymbol {g} _ {i j, s} \right] + t \tag {25}
$$
$$
= T _ {g} \cdot \boldsymbol {x} _ {s} ^ {\prime}. \tag {26}
$$
Here, we use  $x_{s}^{\prime}$  to denote the updated ligand coordinates in Equation. 8 and 18. Therefore, the structure refinement modules in backbone and full-atom refinement are E(3)-equivariant.
Proof. The predicted residue types are obtained by applying softmax and MLPs on the hidden representations from the E(3)-invariant encoder. With the above lemmas, Theorem 1 can be proved since the E(3)-invariant encoder and the E(3)-equivariant structure refinement constitute the E(3)-equivariant generative process of FAIR. Formally, we have:
$$
\left\{\left(\boldsymbol {b} _ {i}, T _ {g} \cdot \boldsymbol {x} \left(\boldsymbol {b} _ {i}\right)\right) \right\} _ {i = 1} ^ {m} = p _ {\theta} \left(T _ {g} \cdot (\mathcal {A} \backslash \mathcal {B}), T _ {g} \cdot \mathcal {M}\right), \tag {27}
$$
which concludes Theorem 1.
A.6 Huber Loss
In Equation 10, we use Huber loss [24] as the structure refinement loss for the stability of optimization, which is defined as follows:
$$
l _ {h u b e r} (x, y) = \left\{ \begin{array}{l} 0. 5 (x - y) ^ {2}, i f | x - y | <   \epsilon , \\ \epsilon \cdot (| x - y | - 0. 5 \cdot \epsilon), e l s e, \end{array} \right. \tag {28}
$$
where  $x$  and  $y$  represent the predicted and ground-truth coordinates. The Huber loss has the following property: if the L1 norm of  $|x - y|$  is smaller than  $\epsilon$ , it is MSE loss, otherwise it is L1 loss. At the beginning of the model training, the deviation between the predicted and ground-truth coordinates is large, and the L1 term makes the loss less sensitive to outliers than MSE loss. The deviation is small when the training is almost complete, and the MSE loss is near 0. In practice, we find that directly using MSE loss sometimes leads to NaN at the beginning of the training, while Huber loss makes the training procedure more stable. Following the suggestion of previous works [28, 33], we set  $\epsilon = 1$  in all our experiments.
Lemma 2. The backbone and full-atom structure refinement are  $E(3)$ -equivariant.
PocketOptimizer [45] is a physico-based computational protein design method that predicts mutations in the binding pockets of proteins to increase affinity for a specific ligand. Its previous version is [42]. Generally, there are four main steps in PocketOptimizer: structure preparation, flexibility sampling, energy calculations, and computation of design solutions. Specifically for the energy calculations, both packing-related energies and binding-related energies are considered. As for the output design solutions, we select the top 100 designs identified by PocketOptimizer based on integer linear programming for downstream metric calculations. Following the suggestions in the original paper, we use AMBER ff14S force field [41] for energy computation and the Dunbrack rotamer library [55] for rotamer sampling for PocketOptimizer in our implementation.
DEPACT [8]  $^{4}$  is a template-matching method that follows the two-step strategy of RosettaMatch [67] for pocket design and has better performance. It first searches the protein-ligand complexes in the database with similar ligand fragments and constructs a cluster model (a set of pocket residues). It then grafts the cluster model into the protein pocket with PACMatch, which places residues in pocket clusters on protein scaffolds by matching key atoms in the cluster model with key atoms of the protein scaffold. The qualities of the generated pockets are evaluated and ranked based on a statistical scoring function. We take the top 100 designed pockets for evaluation. The template databases are constructed in our experiments based on the corresponding training datasets for fair comparisons.
HSRN [28] was originally developed for antibody design. It autoregressively predicts the residue sequence and docks the generated structure. To adapt HSRN to our pocket design task, we replace the antigen with the target ligand molecule to provide the context information. We finetune HSRN's hyperparameters based on its performance on the validation set. The hidden size in the message-passing network is 256, the number of layers is 4, and the number of RBF kernels is 16.
Diffusion [39]  $^{6}$  jointly models both the sequence and structure of an antibody using diffusion probabilistic estimates and an equivariant neural network. The diffusion model first aggregates information from the antigen and the antibody during generation. It then iteratively updates each amino acid's type, position, and orientation. Since the side-chain atoms are not modeled, they are reconstructed at the last step with side-chain packing. To adapt Diffusion to our pocket design task, we replace the antigen with the target ligand molecule to provide the context information. We adjust the key hyperparameters based on the results of the validation sets. The number of layers, hidden dimension, and diffusion steps are 6, 128, and 100, respectively.
MEAN [33] 7 tackles antibody design as a conditional graph translation problem with the modified GMN [23] as encoders. Specifically, MEAN co-designs antibody sequence and structure via a multi-round progressive full-shot scheme, which is more efficient than auto-regressive or diffusion-based approaches. However, it only models protein backbone atoms and neglects side chains, which are important for protein-protein and protein-ligand interactions. To adapt MEAN to our pocket design task, we replace the antigen with the target ligand molecule to provide the context information. The hyperparameters are finetuned based on the validation set for better performance. We set the hidden size as 128, the number of layers as 3, and the number of iterations for decoding as 3.
FAIR is our full atom pocket sequence-structure co-design method. The number of layers for the atom and residue-level encoder are 6 and 2, respectively.  $K_{a}$  and  $K_{r}$  are set as 24 and 8 respectively. The number of attention heads is set as 4; The hidden dimension  $d$  is set as 128. The standard deviation of the Gaussian noise added to the ligand coordinates in model training is 0.1.
<sup>3</sup>https://github.com/Hoeyer-Lab/pocket optimizer
$^{4}$ https://github.com/chenyaoxi/DEPACT_PACMatch
$^{5}$ https://github.com/wengong-jin/abdockgen
$^{6}$ https://github.com/luost26/diffab
<sup>7</sup>https://github.com/THUNLP-MT/MEAN
(b)
Figure 6: (a) The ratio of unconverged residue types and AAR during the iterations of FAIR on Binding MOAD; (b) The convergence plot of RMSD on CrossDocked; (c) The convergence plot of RMSD on Binding MOAD.
C.1 Pocket Generation without Backbone Reference
In FAIR, the original protein backbone coordinates are leveraged for residue-level KNN graph construction for training and generation stability. However, FAIR can also work well without protein backbone coordinates as a reference. In our supplementary experiments, FAIR can still achieve comparable performance.
C.2 Convergence of Iterative Refinement
We show the RMSD curves in Table 6 (b) & (c). We observe the RMSD converges rapidly as the protein pocket is refined. The curve has fluctuations near step 5, potentially caused by the initialization of the side-chain atom at the beginning of full-atom refinement. The coordinates of side-chain atoms are randomly initialized near the  $C_{\alpha}$  and may distort the backbone atoms with the structural refinement. However, the RMSD of FAIR becomes steady quickly after step 10, which again shows the effectiveness of FAIR.
C.3 Local Geometry of the Generated Pockets
Since the RMSD used in Tables 1 and 2 reflects only the correctness of global geometry in generated protein pockets, we also calculate the RMSD of bond lengths and angles in the designed pocket to validate the local geometry following previous work [33]. We measure the average RMSD of the covalent bond lengths in the pocket backbone (C-N, C=O, and C-C). The three conventional dihedral angles in the backbone structure, i.e.,  $\phi, \psi, \omega$  [57] are considered, and the average RMSD of their cosine values are calculated. The results in Table 3 demonstrate that FAIR still achieves much better performance regarding the local geometry than the baseline methods.
Table 3: The average RMSD of bond lengths and dihedral angles of the generated backbone structures on the two datasets.
Models	Bond lengths	Angles (φ,ψ,ω)
PocketOptimizer	0.47	0.328
DEPACT	0.68	0.360
HSRN	0.76	0.437
Diffusion	0.49	0.344
MEAN	0.44	0.312
FAIR	0.35	0.287
C.4 Hyperparameter Analysis
Here, we evaluate the influence of hyperparameter choices on FAIR performance. We focus on two hyperparameters,  $T_{1}$  and  $T_{2}$ , for iterative refinements. In Figure 7, we observe that the performance
of FAIR is stable when  $T_{1}$  and  $T_{2}$  are larger than 5 and 10, respectively, indicating the refinement processes gradually converge. We set  $T_{1}$  and  $T_{2}$  to 5 and 10 throughout our experiments.
(f)
Figure 7: Hyperparameter analysis with respect to  $T_{1}$  (first row) and  $T_{2}$  (second row).