SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We use 2 layer MLPs throughout, with the embedding dimension of  $h = 1920$ . The column MLP is applied on three columns at a time to ensure easier modeling of 3D interactions. Also, to facilitate easier modeling of multiplicative interactions (e.g. forces) these MLPs use gated activation in the hidden layer  $x = x_{0:h} * \sigma(x_{h:2h})$ , where  $x$  is the output of the linear layer,  $h$  is the embedding dimension (the MLPs hidden dimension is effectively  $2 \times h$ ) and  $\sigma$  is the activation function. We chose to use the SiLU [25] activation function as it appears to perform the best on geometric tasks, both on APMixer and the baseline GNN architectures.
M.2 Baselines
M.2.1 Transformer
For the sequence-only generation, we use a standard transformer encoder[12] as implemented by PyTorch [65], with 3 blocks, embedding dimension of 2048, and pre-normalization. The transformer uses the GELU [25] activation function as standard.
M.2.2 EGNN
We use the EGNN [75] as implemented for the equivariant molecule diffusion by Hoogeboom et al. [28]. We extend it to work with residues that have many atoms by using the formulation proposed by Huang et al. [31] and also used by Kong et al. [43], which for updating each node (residue) make use of the distances  $Z_{ij}^{l} = X_{i} - X_{j}$  between atoms in each residue:
$$
m _ {i j} = \phi_ {e} \left(h _ {i} ^ {l}, h _ {j}, Z _ {i j} ^ {l} \left(Z _ {i j} ^ {l}\right) ^ {T}\right),
$$
$$
m _ {i} = \sum_ {j} m _ {i j},
$$
$$
h _ {i} ^ {l + 1} = \phi_ {h} \left(h _ {i} ^ {l}, m _ {i}\right),
$$
$$
X _ {i} ^ {l + 1} = X _ {i} ^ {l} + \sum_ {j} \frac {Z _ {i , j} ^ {l}}{\| Z _ {i , j} ^ {l} \| _ {2} + 1} \phi_ {x} \left(h _ {i} ^ {t}, h _ {j} ^ {t}, Z _ {i j} ^ {l} (Z _ {i j} ^ {l}) ^ {T}\right),
$$
where  $X_{i}$  is a matrix of  $a\times 3$  if we have  $a$  atoms per residue.
To improve the model stability, we apply a layer normalization on  $h^t$  at the beginning of each such block. We wrap every block in EGNN in our constraint layer Section 3.4. The EGNN has 6 such blocks and uses an embedding dimension of 576. As standard, we use SiLU activation and, similarly to AgentNet, we found it slightly beneficial to use a hidden dimension in the MLP of twice the embedding size. The model uses a readout MLP to make the final embedding which is used to predict the amino acid type. This layer is also preceded by layer normalization. An embedding MLP is used for the inputs.
M.2.3 FA GNN
We use the same GNN architecture as used by Puny et al. [68]:
$$
m _ {i, j} = \phi_ {e} \left(h _ {i} ^ {l}, h _ {j}\right)
$$
$$
m _ {i} = \sum_ {j} m _ {i, j}
$$
which is essentially the same architecture as EGNN, just without the distance-based edge features and explicit position updates. To process the positions, we wrap each of these GNN blocks using frame averaging, and for the very first embedding MLP we supply residue atom positions as inputs together with the other features.
The FA-GNN has 4 of these GNN blocks and uses an embedding dimension of 384. SiLU activation and the hidden MLP dimension of twice the embedding dimension are used. The FA-GNN uses the same setup as APMixer with two readout heads for amino acid type prediction and atom position prediction.
M.3 Training
We use AdamW for training [52] with a weight decay of 0.01 and with a learning rate of  $2 \cdot 10^{-4}$  for all structure models, while the transformer used a learning rate of  $1 \cdot 10^{-4}$ . We experimented with weight decay of  $1e - 12$  and learning rates in the range of  $1 \cdot 10^{-3}$  to  $1 \cdot 10^{-4}$  for all models, to determine the chosen values. We also normalize the gradient norm to unit length during training. All models use a batch size of 4. This is the maximum batch size that allows training the baseline models on our 80GB GPU. The APMixer due to the better memory complexity allows for batch sizes up to
32, but to keep the setup similar, we also used a batch size of 4 for it. For paired OAS we train the models for 50 epochs, while for HER2 binder generation we train for 600 epochs, for around 1.25M training steps in both cases. The Paired OAS dataset [61] was split into 1000 test samples and 1000 validation samples. Similarly, the HER2 binder dataset [55] was split to have 1000 test samples and 100 validation samples.
We keep track of the exponential moving average of the model weights, for all models, with a coefficient of 0.995, and use these weights at test time. No model selection is performed.
During training, we use idealized atom positions as returned by our projection layer (Section 3.4) as the target. This ensures that the exact target configuration is reachable by the model.
The CDR redesign baselines of RefineGNN [34] and MEAN [43] were trained for the same amount of steps using the default parameters and early stopping.
All experiments were performed on an 80GB A100 GPU.
M.4 APMixer Classifier
We adapt APMixer for HER2 binder classification (Section 5.2) by including a global readout. This readout is constructed by applying mean pooling over the frames and residues after each APMixer block and then applying a layer normalization and a linear layer. All of the linear layer outputs are summed to obtain the final prediction. This follows a popular graph classification readout used in GNNs [93].
The classifier is trained for 100 epochs using a batch size of 16. No exponential moving average of weights is used, and the model with the best validation accuracy is selected. The full HER2 dataset was split to have 3000 test samples and 1000 validation samples.
N Conditional CDR Redesign
While co-crystal structures usually come quite late in the drug discovery process, performing CDR structure and sequence redesign for a given, unseen co-crystal structure is a popular benchmark for deep generative models [53; 43; 44; 34]. Here, we use the data split from Luo et al. [53], derived from the SAbDab database of antibody-antigen structures [13]. This split ensures that all antibodies similar to those of the test set ( $>50\%$  CDR H3 identity) are removed from the training set. In contrast to the experimental setup of Luo et al. [53], we report the results for the more difficult task of redesigning all CDRs at once $^2$ , rather than redesigning a single CDR at a time. Following Luo et al. [53], the results are averaged over 100 samples for each of the 19 test structures. We observed that in this setup the minimization protocol used by Luo et al. [53] increased the RMSD values (Table 9), therefore, we report the results without it.
To adapt AbDiffuser to this conditional case, we build a bipartite graph between the antibody CDR residues and the nearest antigen residues that are at most  $10\AA$  away, in terms of  $C_{\beta} - C_{\beta}$  distance. Then we simply replace the initial MLP embedding layer with a one-layer FA-GNN. The rest of the APMixer architecture remains unchanged. We trained AbDiffuser with a batch size of 8 for 30k steps, with the rest of hyperparameters staying the same as before.
Table 9 shows that in this conditional setup AbDiffuser outperforms the state-of-the-art diffusion baseline (DiffAb) in amino acid recovery (AA) by a wide margin. This good performance could be attributed to the usefulness of AHo numbering, and APMixer being better at modeling sequences than a GNN. Our ablation experiments on paired OAS generation (Table 1) corroborate this hypothesis. Although the CDR RMSD results are generally quite comparable between DiffAb and AbDiffuser, importantly, AbDiffuser performs better in CDR H3, which has the most variability in its structure and can influence the binding the most. Both methods strongly outperform RosettaAntibodyDesign (RAbD) [2], the state-of-the-art computational biology antibody design program, in terms of amino acid recovery.
2We use the checkpoint for this task provided by the original authors.
Model	CDR H1	CDR H2	CDR H3	CDR L1	CDR L2	CDR L3
AA↑	RMSD↓	AA↑	RMSD↓	AA↑	RMSD↓	AA↑	RMSD↓	AA↑	RMSD↓	AA↑	RMSD↓
Rosetta (RAbD)* [2]	22.85%	2.261	25.5%	1.641	22.14%	2.9	34.27%	1.204	26.3%	1.767	20.73%	1.624
DiffAb (Minimized)* [53]	65.75%	1.188	49.31%	1.076	26.78%	3.597	55.67%	1.388	59.32%	1.373	46.47%	1.627
DiffAb (Minimized) [53]	66.37%	1.371	42.82%	1.337	28.27%	3.798	62.91%	1.520	62.59%	1.653	49.38%	1.616
DiffAb [53]	66.37%	0.802	42.82%	0.722	28.27%	3.550	62.91%	1.120	62.59%	1.025	49.38%	1.066
AbDiffuser	79.09%	1.120	72.33%	0.995	36.14%	2.921	85.08%	1.138	85.02%	1.273	75.68%	0.990
AbDiffuser (side chains)	76.30%	1.584	65.72%	1.449	34.10%	3.346	81.44%	1.462	83.22%	1.397	73.15%	1.591
AbDiffuser (τ = 0.75)	79.76%	1.083	72.96%	0.950	36.53%	2.805	85.60%	1.098	85.71%	1.237	76.29%	0.955
AbDiffuser (s.c., τ = 0.75)	76.06%	1.713	66.38%	1.512	34.71%	3.214	81.93%	1.373	83.54%	1.351	73.32%	1.509
AbDiffuser (τ = 0.01)	81.11%	1.075	74.27%	0.946	37.27%	2.795	86.26%	1.115	86.85%	1.238	77.06%	0.966
AbDiffuser (s.c. τ = 0.01)	75.36%	2.463	66.89%	2.010	35.56%	3.124	83.10%	1.525	82.95%	1.623	74.19%	1.502
Table 9: Generating CDR structure and sequence for existing co-crystal structures from SAbDab [13]. All CDRs are generated at once. The amino acid recovery rate (AA) and RMSD are reported for each individual region. Baseline models from Luo et al. [53] marked by * generated CDRs one at a time.