# CNN-Trans-Enc: A CNN-Enhanced Transformer-Encoder On Top Of Static BERT
representations for Document Classification

Charaf Eddine Benarab  School of Computer Science and Engineering
University of Electronics Science and Technology of China
Chengdu, China
<EMAIL_ADDRESS>Shenglin Gui  School of Computer Science and
Engineering
University of Electronics Science and Technology of China
Chengdu, China
<EMAIL_ADDRESS>
###### Abstract

BERT achieves remarkable results in text classification tasks, it is yet not
fully exploited, since only the last layer is used as a representation output
for downstream classifiers. The most recent studies on the nature of
linguistic features learned by BERT, suggest that different layers focus on
different kinds of linguistic features. We propose a CNN-Enhanced Transformer-
Encoder model which is trained on top of fixed BERT $[CLS]$ representations
from all layers, employing Convolutional Neural Networks to generate QKV
feature maps inside the Transformer-Encoder, instead of linear projections of
the input into the embedding space. CNN-Trans-Enc is relatively small as a
downstream classifier and doesn’t require any fine-tuning of BERT, as it
ensures an optimal use of the $[CLS]$ representations from all layers,
leveraging different linguistic features with more meaningful, and
generalizable QKV representations of the input. Using BERT with CNN-Trans-Enc
keeps $98.9\%$ and $94.8\%$ of current state-of-the-art performance on the
IMDB and SST-5 datasets respectably, while obtaining new state-of-the-art on
YELP-5 with $82.23$ ($8.9\%$ improvement), and on Amazon-Polarity with
$0.98\%$ ($0.2\%$ improvement) (K-fold Cross Validation on a 1M sample subset
from both datasets). On the AG news dataset CNN-Trans-Enc achieves $99.94\%$
of the current state-of-the-art, and achieves a new top performance with an
average accuracy of $99.51\%$ on DBPedia-14.

###### Index Terms:

Text Classification, Natural Language Processing, Convolutional Neural
Networks, Transformers, BERT

##  I Introduction

Transformer based architectures[32], and BERT[4] in particular achieve
remarkable results in natural language processing. Within document
classification, and although BERT achieves outstanding results on almost all
datasets available, there is still potential to further explore the effect
that the classifier added on top of the last layer has on the performance.
BERT contains multiple Transformer-Encoder layers [32], which output different
hidden states for the same input. The experiment conducted by Jawahar et al.
[8], provides an empirical diagnosis of the linguistic structures learned by
different BERT layers. Their experiment employs probing tasks [2], where
specific auxiliary classifiers are used to predict targeted properties in a
sentence representation. Taking the activation of a $[CLS]$ token, which is
representative of the whole sentence due to the attention mechanism embedded
in the Transformer-Encoder [32], and primarily used in classification tasks.
Jawahar et al. [8] show that after probing BERT layers separately, it is
observed that three different kinds of linguistic features are better encoded
by different layers. Their findings which are prerequisites in this paper are
as follows: Surface features are mostly captured by bottom layers, syntactic
features by middle layers, and semantic features by higher layers.

We propose a variant of the Transformer-Encoder layer introduced by Vaswani et
al. [32], replacing Key, Query, and Value projection matrices with
convolution-based submodules. The input to our CNN-Trans-Enc model consists of
$[CLS]$ representations from all layers of BERT [4], this ensures the
downstream model is receiving multiple representations which are rich in
different kinds of linguistic features, as suggested by Jawahar et al. [8]

We compare the proposed CNN-enhanced Transformer-Encoder to the CNN for
Sentence Classification [11], vanilla Transformer-Encoder [32], and a simple
Softmax classifier [29] (Which is the common approach). We train the Softmax
classifier exclusively on the last $[CLS]$ token representation, the CNN-based
models both on the last $[CLS]$ representation and on all $[CLS]$
representations from all layers, and the Transformer-Encoder and CNN-enhanced
Transformer-Encoder on all $[CLS]$ representations obtained from all layers.
The contributions of this paper are summarized as follows:

  * •

We demonstrate the importance of linguistic feature types in the task of
Sentiment/Document Classification. We build on a conjecture that suggests
using representations rich in different linguistic features would improve
classification performance, and we investigate our assumption by training
different models on fixed BERT representations learned by the last layer, and
all layers ($[CLS]$ representations).

  * •

We introduce an approach which employs $[CLS]$ representations from different
layers of BERT as channels to a CNN sub-architecture, congenially with a
Transformer-Encoder for the purpose of leveraging all BERT-layers to generate
an optimal representation for classification.

  * •

BERT+CNN-Trans-Enc exceeds the state-of-the-art results on YELP-5, Amazon-
Polarity, and DBPedia-14 datasets, and outperforms most models on all used
datasets in our experiment without any kind of fine-tuning of BERT.

##  II Related work

Text Classification [21] in the context of deep learning mainly includes:
Sentiment Analysis, News Categorization, and Topic Classification. To validate
and test our method, we use six different datasets spanning the three tasks
mainly constituting Text Classification. Survey [21] reviews multiple datasets
and models employed within the context of Text Classification (using deep
learning models).

BERT (Bidirectional Encoder Representations from Transformers) [4] is a
landmark architecture, which is fundamentally a stack of bidirectional
Transformer-Encoders [32], pre-trained for masked language and next sentence
prediction. BERT learns sentence representations which are semantically rich,
and achieves remarkable results on a variety of language-related benchmarks.
BERT takes a sequence of tokens as input, and outputs a representation for the
whole input after adding special tokens signifying the beginning and end of a
sentence. For classification, the most common practice is to feed the $[CLS]$
representation obtained from the last layer (hidden-state of the $[CLS]$ token
from the last Transformer-Encoder layer), into a Softmax [29], which outputs
the probability of a label $c$:

$p(c|h)=Softmax(\mathbf{W}\cdot h_{[CLS]_{12}})$ (1)

$\mathbf{W}$ is a task-specific parameter matrix, and $h_{[CLS]_{12}}$ is the
last hidden-state of the $[CLS]$ token.
Pre-training and fine-tuning BERT on downstream datasets can be
computationally expensive, given the large number of parameter it contains,
and the ambiguity behind why it works so well. ALBERT [13] proposes a
lightweight version of BERT employing parameter-reduction techniques to lower
memory consumption, this shows significant improvements in scaling on multi-
sentence inputs. RoBERTa [17] on the other hand presents a study of BERT pre-
training, and applies alternative pre-training objectives and tasks,
showcasing the effect they have (alongside pre-training data) have on the
performance. SMART [9] proposes a learning framework to achieve efficient
fine-tuning and attain better generalization, using smoothness-inducing
regularization and Bregman proximal point optimization. SMARTRoBERTa [9] and
SMARTBERT [9] show performance gains in comparison to BERT [4] and RoBERTa
[17] across a variety of NLP tasks. XLNet [37] improves over BERT using auto-
regressive pretraining, enabling leaning bidirectional contexts using the
expected likelihood over all permutations of the factorization order, and
solving BERT’s limitations due to neglecting dependency between masked
positions.

The methods and techniques which are BERT-based and proposed in the last
couple of years do not incorporate in any way the nature of representations
learned and generated by BERT, While for a task like classification, the
linguistic features fed to a classifier can be crucial, as they encode
relevant information about the input in the embedding space.
Convolutional Neural Networks have proven to be efficient feature extractors
which serve well for text classification. Kim Yoon [11] suggested training a
CNN architecture with filters of multiple sizes, on top of pre-trained
Word2Vec [20] embeddings. The feature maps from the different filters, are
concatenated after a max-over-time pooling operation, and passed to a Softmax
layer. Zhang et al. [39] expanded on Kim Yoon’s [11] model, with an empirical
identification of hyper-parameter settings. Kim’s CNN for text [11] motivates
a portion of our work, and inspires a critical part of the proposed method in
this paper. DPCNN [10] employs deepening of word level CNNs to capture global
representations of text, and proposes a deep pyramid convolutional neural
network that achieves the best accuracy by increasing the network depth,
without increasing computational cost.
Standalone hybrid frameworks which combine attention and CNNs were proposed,
in an attempt to leverage their capabilities simultaneously. In HAHNN [1],
both word and sentence-level context vectors are generated hierarchically
using separate attention blocks, in order to construct a document
representation based on the interaction between words and sentences, setting a
high-level context vector which serves as representation for a document. A
variation of the HAHNN [1] method employs a simple convolutional layer after
the embedding layer for the purpose of effective feature extraction before the
attention blocks.

##  III Model

[scale=0.9]figures/fig1

Figure 1:  Model architecture : CNN-Trans-Enc obtains Queries, Keys, and
Values matrices from convolution operations applied on the input $X$, which
represents 12 stacked BERT $[CLS]$ representations. The $CNN[CLS]$ submodule
is replicated three times. We apply an Adaptive-Max Pooling and concatenation
($\oplus$) on the resulting feature maps to obtain
$Q_{conv},K_{conv},V_{conv}\in\mathbb{R}^{12\times d_{m}}$ where $d_{m}$ is
the embedding dimension. $Q_{conv}$, $K_{conv}$, and $V_{conv}$ are then
directly fed to a stack of two Transformer-Encoder [32] layers.

The model architecture, shown in Fig 1, consists of a representation layer,
where BERT [4] is used to obtain fixed $[CLS]$ representations from its 12
layers. In our experiment, the pre-trained BERT-base model is exclusively
employed as an embedding layer, and $[CLS]$ representations from all layers
are concatenated to form a tensor $X\in R^{12\times 768}$ for each text sample
in a dataset $D$. This work presents a more practical downstream model for
classification tasks, added on top of BERT. CNN-Trans-Enc is a variant of the
original Transformer-Encoder [32], which uses convolution based submodules to
obtain Q, K, and V matrices instead of linear projections of the input to the
embedding space. To obtain more efficient and meaningful representations which
contain linguistic features learned by all BERT layers, we treat the 12
$[CLS]$ representations for a sample $s$, as $12$ channels. We use three
convolutional filters set to a length of $l$,
$W_{1},W_{2},W_{3}\in\mathbb{R}^{4\times l}$. Each filter takes in every $4$
adjacent channels, and each feature $h_{i}$, $s_{i}$, or $v_{i}$ could be
computed as follows:

$\displaystyle h_{i}=f(W^{T}_{1}\cdot X_{i:i+l-1}+B_{h})$ (2) $\displaystyle
s_{i}=f(W^{T}_{2}\cdot X_{i:i+l-1}+B_{s})$ $\displaystyle
v_{i}=f(W^{T}_{3}\cdot X_{i:i+l-1}+B_{v})$

Where $B$ is the bias term, $f$ is the hyperbolic tangent non-linear function,
$i$ is the position of the current convolution operation increase by a stride
of 2, and $X_{i:i+l-1}\in\mathbb{R}^{4\times l}$ is a window of length $l$
from convolution position $i$ to position $i+l-1$. Three feature maps $h$,
$s$, and $v$ are then formed with:

$\displaystyle h=[h_{0},h_{1},....,h_{n}]$ (3) $\displaystyle
s=[s_{0},s_{1},....,s_{n}]$ $\displaystyle v=[v_{0},v_{1},....,v_{n}]$

Where $n=\frac{768-kernel\\_size}{2}+1$. Adaptive-Max pooling is used to
reduce $h,s,v\in\mathbb{R}^{4\times 382}$ to
$h_{p},s_{p},v_{p}\in\mathbb{R}^{4\times d_{m}}$. The resulting pooled
feature-maps $h_{p},s_{p}$ and $v_{p}$ are then concatenated over the first
dimension to form the desired $Q_{conv}$, $K_{conv}$, or $V_{conv}$ matrix
$[h_{p};s_{p};v_{p}]$. We refer to this CNN-based module as $CNN[CLS]$. We
employ three replicas of $CNN[CLS]$ to generate
$Q_{conv},K_{conv},V_{conv}\in\mathbb{R}^{12\times d_{m}}$ simultaneously in
parallel, where $d_{m}$ is the chosen embedding dimension [32]. After
generating $Q_{conv},K_{conv}$, and $V_{conv}$, we feed them into a first
Transformer-Encoder [32], starting with a multi-head attention:

$\begin{split}&\operatorname{Multi-head}(Q_{conv},K_{conv},V_{conv})=\\\
&Concat(head_{1},...,head_{h})W^{O}\end{split}$ (4)
$head_{k}=Attention(Q_{conv}{W_{k}}^{Q},K_{conv}{W_{k}}^{K},V_{conv}{W_{k}}^{V})$
(5)

${W_{k}}^{Q}$, ${W_{k}}^{K}$, ${W_{k}}^{V}\in\mathbb{R}^{d_{m}\times d_{v}}$
and $W^{O}\in\mathbb{R}^{d_{m}\times 768}$ are parameter matrices and
$Attention$ is defined as:

$\begin{split}Attention(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})V\end{split}$
(6)

The main objective behind using a Transformer-Encoder [32] that takes in
$[CLS]$ representations from all layers, is to capture linguistic features
learned by all layers. And to integrate the findings reported by Jawahar et
al. [8], the Scaled-Dot Product attention is used in this scope to capture
similarities between $[CLS]$ representations from all layers. Intuitively,
using two Transformer-Encoder layers, and reducing the size of the output
projection matrix for the second one, results in a representation that encodes
multiple linguistic features captured by BERT layers, in one fairly small
vector. We safely remove the positional encoding since the Transformer-Encoder
receives the same order of $[CLS]$ representations from the exact same BERT
layers which are frozen and exclusively used as a feature extractor. The
output of the first Transformer-Encoder block is denoted as
$X_{0}\in\mathbb{R}^{12\times 768}$, and it is of the same size as the input
$X$.
The Second Transformer-Encoder layer is almost similar to the first one, they
only differ in the Multi-Head Attention operation. The projection matrix
$W_{outdim}^{O}\in\mathbb{R}^{(12d_{m})\times outdim}$ is applied to the
vectorized output of the $\operatorname{Multi-head}$ attention, where $outdim$
is the representation dimension we obtain from the second Transformer-Encoder
layer. We remove the residual connection linking the input to the Multi-head
Attention in the second layer and its output, due to the incompatible shapes
caused by the projection matrix $W_{outdim}^{O}$. The Multi-Head Attention
equation is therefore reformulated as:

$\operatorname{Multi-head}(Q_{2},K_{2},V_{2})=ZW_{outdim}^{O}$ (7)
$Z=\mathbf{vec}(Concat(head_{1},...,head_{h}))$ (8)
$head_{k}=Attention(Q_{2}{W_{k}}^{Q},K_{2}{W_{k}}^{K},V_{2}{W_{k}}^{V})$ (9)

Where $Q_{2},K_{2}$, $V_{2}\in\mathbb{R}^{12\times d_{m}}$ are obtained with
linear projections of the output of the first Transformer-Encoder layer
$X_{0}\in\mathbb{R}^{12\times 768}$, through projection matrices
$W_{Q},W_{K},W_{V}\in\mathbb{R}^{768\times d_{m}}$. $\mathbf{vec}$ is the
vectorization of $Concat(head_{1},...,head_{h})\in\mathbb{R}^{12\times d_{m}}$
to a vector $Z\in\mathbb{R}^{12d_{m}}$. $Output\in\mathbb{R}^{outdim}$ is a
vector representation of the whole concatenation of the $[CLS]$
representations obtained from BERT, which encodes linguistic features learned
by all BERT layers, chosen and pooled using the $CNN[CLS]$ modules, and
reduced to a small vector by exploiting the attention mechanism in the
Transformer-Encoder, which allows for similarity measurement, and
dimensionality reduction by altering the output projection matrix of the
second Transformer-Encoder layer. Using three separate CNN-based modules on
the output of BERT [4] allows for more effective feature extraction, taking in
consideration the features which are considered to be most relevant by the
convolution and pooling operations. The attention mechanism is then applied to
the resulting feature maps to generate a vector comprising different kinds of
linguistic features learned by BERT, and serves as a fairly small
representation for classification tasks which helps achieve competitive
results without fine-tuning BERT. The last layer is a Softmax function which
receives $Output$ and generates probability distributions over classes
according to:

$p(c|Output)=Softmax(\mathbf{W}\cdot Output)$ (10)

Where $\mathbf{W}$ is a parameter matrix, and $Output\in\mathbb{R}^{outdim}$
is the representation obtained from applying CNN-Trans-Enc on the 12 $[CLS]$
representations from all layers of BERT.

##  IV Datasets and Experimental Setup

Survey [22] is a holistic overview of $Text$ and $Sentiment$ $Classification$
methods and processes. We follow the scheme for achieving state-of-the-art
$Document$ $Classification$ as follows:

  * •

Datasets are chosen according to current state-of-the-art in sentiment and
topic classification, and news categorization. We used six supervised datasets
both for binary and multi-class classification, section IV-A and TABLE I.

  * •

Evaluation is conducted using K-fold Cross Validation using random splits, and
repeated over five random seeds, average Cross-Validation results are reported
in TABLE V. Alongside our proposed architecture, we evaluate other four models
on top of fixed representations from BERT, according to different design
choices, and use ASO [5] test for statistical Significance Testing between
pairs of models in our experiment.

###  IV-A Datasets

We evaluate our work on multiple document classification datasets, including
four sentiment analysis datasets, one news classification dataset and one
topic classification dataset, shown in TABLE I:

TABLE I: Dataset statistics, C: Number of classes, T: Task, N: Number of samples Dataset |  C |  T |  N
---|---|---|---
IMDB |  2 |  SA |  50k
Amazon |  2 |  SA |  1M
YELP |  5 |  SA |  1M
SST-5 |  5 |  SA |  11861
AG |  4 |  NC |  127600
DBPedia |  14 |  TC |  630,000

  * •

IMDB [19]: IMDB movie reviews dataset, is a collection of 50,000 reviews, with
a maximum of 30 reviews per movie. A review is negative if it has a score $<$
4 out of 10, and a positive review has score $>$ 7 out of 10.

  * •

Yelp-Reviews [38]: This is a dataset for business reviews and ratings,
currently containing 8,635,403 samples with 5 labels corresponding to the
scores associated with the sentiments expressed in the reviews. We use a
subset of 1 million reviews from the Yelp-Reviews dataset and we make sure we
have a balanced set, with 200 thousand reviews for each label.

  * •

SST-5 [28]: Stanford Sentiment Treebank is a Sentiment Scoring corpus, which
contains parse trees for all samples, We use 11861 samples, with labels binned
in 5 bins from 0 to 4.

  * •

Amazon-Polarity [38]: The Amazon Polarity dataset, contains product reviews
from Amazon, up to march 2013 it contained up to 35 million reviews with
binary labels representing the corresponding sentiments as positive or
negative classes. We use 1 million samples with a balanced distribution of 500
thousand samples for each label.

  * •

AG’s news corpus
111http://groups.di.unipi.it/~gulli/AG_corpus_of_news_articles.html: AG is a
corpus of news articles containing 127,600 samples, distributed over 4
classes, obtained from more than 2,000 news sources by ComeToMyHead, an
academic news search engine. We use the entire number of samples for k-fold
cross validation in our experiment.

  * •

DBPedia 14 [15]: DBPedia is formed by extracting structured information from
Wikipedia, it contains 630,000 samples, and 14 ontology classes. We use the
entirety under k-fold cross validation in our experiment.

###  IV-B Training and Hyperparameters

We use the BERT-base model [4], which is pre-trained and available online by
the HuggingFace team222https://huggingface.co/bert-base-uncased, with a hidden
size of 768, 12 layers, and 12 attention heads. We use the pre-trained BERT
model once for each dataset and generate $[CLS]$ representations from all 12
layers, for all samples in our datasets. This means that if a dataset has a
number of samples $N$, this is results in dataset $D\in\mathbb{R}^{N\times
12\times 768}$, where the second dimension $12$, represents $12$ $[CLS]$
representations from $12$ layers.
Our proposed architecture, which is trained on top of fixed $[CLS]$
representations from BERT [4], combines convolution-based modules, and two
stacked Transformer-Encoder layers. In this work,we employ $h=20$ parallel
attention heads for both Trans-Enc layers and we set $d_{k}=d_{v}=d_{m}/h=19$.
The embedding dimension $d_{m}$ is set to $380$ and $outdim$ to $320$.
The $CNN[CLS]$ convolutional modules used to obtain $Q_{conv}$, $K_{conv}$,
and $V_{conv}$, are trained with hyperparameter choices that are heavily
influenced by the the work conducted by Kim Yoon [11] and Zhang et al. [39].
We give the baseline CNN configuration described in TABLE II.

TABLE II: Baseline configuration. Adam: Adam Optimizer [12], filter-sizes: both architectures have three convolutional filters. Xavier-Init: Xavier-Initialization [23], Tanh: hyperbolic tangent non-linear function. Description |  Values
---|---
optimizer |  Adam
learning-rate |  0.0001
filter-length $l$ |  5
input-channels |  12
initialization |  Xavier-Init
activation function |  Tanh
Dropout ratio |  0.3

The Post-LN Transformer-Encoder [32] [36] from which our two Trans-Enc layers
are inspired, suffers from instability when optimized with large learning
rates and requires careful optimization choices. For an accurate evaluation,
we follow the training paradigm used with the original Transformer
architecture [32], and emphasized by Xiong et al. [36]. For both Transformer-
Encoder layers, we use Adam optimizer [12], with $\beta_{1}=0.9$,
$\beta_{2}=0.98$, and $\epsilon=10^{-9}$. The learning rate is increased
during a warmup-stage according to:

$lr(t)=\frac{t}{T_{warmup}}lr_{max},t\leq T_{warmup}$ (11)

We set $T_{warmup}=1000$, $lr_{max}=0.001$, and adapt the batch-size for all
datasets to have $6000$ training steps for each one. After $T_{warmup}$,
inverse square-root decay is used as a scheduler to set the learning rate.
Xavier initialization [23] is consistently used throughout our experiment.
We replicate the experiment five times using different random seeds. Each
replication constitutes five runs of a 5-fold Cross-Validation. We report
averages of Cross-Validation accuracy means over the five seeds in TABLE III.

###  IV-C Model Variations

TABLE III: Results of our CNN-Trans-Enc model against other methods. Results recorded on the first five rows concern the models implemented by the authors and tested according to the methodology described in this paper (recorded 5-fold Cross-Validation accuracies as described in IV-B). $*$ means results are reported on their papers. $-$ means not reported. Results without $*$ are obtained from our experiments. RoBERTa-base [17] and ALBERT [13] have no reported performance on the used datasets, and their performance on the IMDB dataset is obtained from our experiment, due to the small size of IMDB and difficulty to train both of them. Model |  IMDB |  YELP-5 |  SST-5 |  AMZ-2 |  AG |  DBP
---|---|---|---|---|---|---
BERT-last+Softmax [29] |  93.21 |  70.1 |  49.15 |  93.2 |  92.05 |  99.10
BERT-last+Kim-CNN [11] |  94.3 |  72.05 |  51.23 |  96.45 |  93.21 |  99.28
BERT-12+CNN[CLS] |  88.0 |  65.3 |  48.1 |  91.2 |  86.98 |  98.73
BERT-12+Trans-Enc [32] |  94.6 |  75.33 |  53.01 |  96.9 |  93.40 |  99.43
BERT-12+CNN-Trans-Enc(Ours) |  95.7 |  82.23 |  56.03 |  98.05 |  95.49 |  99.51
BERT-ITPT-FiT [29] |  95.63* |  70.58* |  - |  - |  - |  -
DistilBERT [27] |  92.82* |  - |  - |  - |  - |  -
RoBERTa-base [17] |  93.78 |  - |  - |  - |  - |  -
ALBERT [13] |  88.69 |  - |  - |  - |  - |  -
HAHNN(CNN) [1] |  92.26* |  73.28* |  - |  - |  - |  -
HAHNN (TCN) [1] |  95.17* |  72.63* |  - |  - |  - |  -
XLNet [37] |  96.8* |  72.95* |  - |  97.89* |  95.55* |  99.40*
BERT-base [24, 16] |  95.63* |  70.80* |  53.2* |  - |  95.20* |  99.35*
BERT-large [24, 35, 29] |  95.79* |  70.68* |  55.5* |  97.37* |  95.34* |  99.39*
BERT-large-finetune-UDA [35] |  95.8* |  67.92* |  - |  96.5* |  - |  -
RoBERTa-large+Self-Explaining [30] |  - |  - |  59.1* |  - |  - |  -
DPCNN [10] |  - |  69.42* |  - |  96.68* |  93.13* |  99.12*
CCCapsNet [26] |  - |  - |  - |  94.96* |  92.39* |  98.72*
EXAM [6] |  - |  - |  - |  95.5* |  - |  -
DRNN [33] |  - |  - |  - |  96.49* |  94.47* |  99.19
NB-weighted-BON + dv-cosine [31] |  97.42* |  - |  - |  - |  - |  -
EFL [34] |  96.1* |  - |  - |  - |  86.10* |  -
GraphStar [18] |  96.0* |  - |  - |  - |  - |  -
ULMFiT [7] |  95.40* |  70.02* |  - |  - |  94.99* |  99.20*
VDCNN [3] |  - |  64.72* |  - |  95.72* |  91.33* |  98.71*
ULMFiT (small-data) [25] |  - |  67.60* |  - |  96.10* |  93.70* |  99.20*
CharCNN [38] |  - |  62.05* |  - |  95.07* |  91.45* |  98.63*
DNC + CUW [14] |  - |  - |  - |  - |  93.90* |  99.00*

Alongside CNN-Trans-Enc, we experiment with the $CNN[CLS]$, and the two
Transformer-Encoder layers (Trans-Enc) separately to understand the effect of
each one. They are both trained on top of 12 $[CLS]$ representations from
BERT, while freezing BERT layers during training.

  * •

$\mathbf{CNN[CLS]}$: implemented as described in III, feeding the pooled
features from $h,s$, and $v$ feature maps into a softmax classifier, which
outputs probabilities for correct class labels.

  * •

Trans-Enc: implemented and trained in the same way as the two Transformer-
Encoder layers in III, only the $Q$, $K$, and $V$ matrices are obtained using
linear projections as proposed by Vaswani et al. [32] unlike in CNN-Trans-Enc.
$Output$ is directly fed to a softmax classifier as described in equation 10.
Training procedure conforms to hyperparameter choices mentioned in IV-B for
both Transformer-Encoder layers.

We also conduct an experiment using the last $[CLS]$ representation from BERT,
fed to a Convolutional Neural Network (Kim Yoon [11] and Zhang et al. [39]). A
Softmax classifier (equation 1) is also trained and evaluated using the last
$[CLS]$ representation. the later pair of models is used to evaluate a single
$[CLS]$ representation using a CNN, and a simple softmax classifier. This
comparison serves as an empirical proof of the feature extraction abilities of
CNNs, and how they affect the performance in document classification.

  * •

Kim-CNN: Instead of the concatenation of word embeddings $x_{1}\oplus
x_{2}\oplus...\oplus x_{i}\oplus...\oplus x_{n}$, where
$x_{i}\in\mathbb{R}^{d}$, and $n$ is the sentence length, we feed the last
$[CLS]$ representation from BERT to three convolution filters
$W_{1}\in\mathbb{R}^{l_{1}}$, $W_{2}\in\mathbb{R}^{l_{2}}$, and
$W_{3}\in\mathbb{R}^{l_{3}}$. Let $X\in\mathbb{R}^{768}$ be the last $[CLS]$
sentence representation, the application of the convolutional filters is
performed on windows of lengths $l_{1}$, $l_{2}$, and $l_{3}$ according to:

$\displaystyle c_{i}=f(W^{T}_{1}\cdot X_{i:i+l_{1}-1}+B_{c})$ (12)
$\displaystyle d_{i}=f(W^{T}_{2}\cdot X_{i:i+l_{2}-1}+B_{d})$ $\displaystyle
k_{i}=f(W^{T}_{3}\cdot X_{i:i+l_{3}-1}B_{k})$

Where $B$ is the corresponding bias term and $i$ is decided by a sliding
operation adding a stride of 2 at each convolution step, and f is the
hyperbolic tangent non-linear function. $l_{1},l_{2}$ and $l_{3}$ are window
lengths with values $5$, $10$, and $15$ respectively. The resulting feature
maps $c$, $d$, and $k$ are formed with simple concatenations:

$\displaystyle c=[c_{0},c_{1},....,c_{n}]$ (13) $\displaystyle
d=[d_{0},d_{1},....,d_{n}]$ $\displaystyle k=[k_{0},k_{1},....,k_{n}]$

Where $n=\frac{768-kernel\\_size}{2}+1$. Then we apply an Adaptive-Max pooling
to capture higher values corresponding to essential features in the resulting
feature maps $c$, $d$, and $k$. The pooled features are concatenated, fed to a
linear classifier followed by a softmax (equation 1). Xavier-Init [23], Adam
[12] with a learning rate of $0.001$, and Hyperbolic Tangent Non-Linear
Function are used to train Kim-CNN in this experiment.

  * •

Softmax: The last fixed $[CLS]$ representation is directly fed to the softmax
equation 1, and optimized to predict correct class probabilities.

##  V Results and Discussion

We built on a conjecture, that different linguistic features which are
distributed over different layers, makes the $[CLS]$ representation from the
last layer only partially representative of the sentence. The Transformer-
Encoder [32] prominently uses Multi-Head attention to derive an output based
on the similarities between the input elements. Also, the work conducted by
Kim Yoon [11], and Zhang et al. [39] proves Convolutional Neural Networks
serve as optimal feature extractors for text represented with static
embeddings like Word2Vec [20]. CNN-Trans-Enc harnesses both the feature
extraction abilities of convolutional neural networks, and the similarity
measurement hallmark of the attention mechanism prominently used in
Transformers [32]. By altering the last projection matrix size $W^{O}$ of the
second Transformer-Encoder layer, we optimize the architecture to learn a
representation which encodes the different types of linguistic features
learned by different layer regions of BERT in varying aggregations (according
to [8]).
On the IMDB dataset, BERT-12+CNN-Trans-Enc does not show performance gains
over XLNet [37] as it scores exactly the same as BERT-ITPT-FiT [29], achieving
$98.23\%$ of the current state-of-the-art performance (NB-weighted-BON + dv-
cosine [31]). BERT-12 + CNN-Trans-Enc outperforms DistilBERT [27], RoBERTa-
base [17], and ALBERT [13].
On the YELP-5 dataset, both BERT-12+Trans-Enc and BERT-12+CNN-Trans-Enc
outperform BERT-ITPT-Fit [29], XLNet [37], and HAHNN(CNN) [1], and BERT+CNN-
Trans-Enc achieves new state-of-the-art with $82.23\%$, outperforming the
current state-of-the-art HAHNN(CNN) [1] by a margin of $\approx$ $8.9\%$.
On the SST-5 dataset, our method outperforms BERT-large [24] [35] by $\approx$
$1.0\%$, obtaining $56.03\%$. BERT-12+CNN-Trans-Enc achieves $94.8\%$ of the
current state-of-the-art performance on the SST-5 dataset, RoBERTa-large+Self-
Explaining [30] with $59.1\%$.
BERT+CNN-Trans-Enc also achieves new state-of-the-art on the Amazon-Polarity
dataset, exceeding XLNet [37] by a margin of $\approx$ $0.16\%$, with a 5-fold
cross-validation accuracy of $98.05\%$ on only a 1M sample subset of the
dataset.
On the AG news dataset BERT+CNN-Trans-Enc achieves $99.94\%$ of the current
state-of-the-art performance (XLNet [37]), and outperforms it on DBPedia [15]
by $\approx$ $0.11\%$ points, achieving new state-of-the-art with a $99.51\%$
accuracy.

### Effect of CNNs to generate QKV matrices:

TABLE IV: $\epsilon^{min}$ scores between all pairs of architectures used on top of static BERT representations. To avoid confusion, the $-$ is used for cases concerning the same model. The values indicate the comparison between the model names in the leftmost column, with the ones in the upper row. $\epsilon^{min}=0.0$ means that the model on the leftmost column is better than the one on the upper row, and $\epsilon^{min}=1.0$ means the inverse. Model |  CNN-Trans-Enc |  Trans-Enc |  Softmax |  Kim-CNN |  CNN[CLS]
---|---|---|---|---|---
CNN-Trans-Enc |  - |  0.0 |  0.0 |  0.0 |  0.0
Trans-Enc |  1.0 |  - |  0.0 |  0.0 |  0.0
Softmax |  1.0 |  1.0 |  - |  1.0 |  0.0
Kim-CNN |  1.0 |  1.0 |  0.0 |  - |  0.0
CNN[CLS] |  1.0 |  1.0 |  1.0 |  1.0 |  -

The results depicted in TABLE III show that for the last $[CLS]$
representation, BERT-Feat+Kim-CNN [11] outperforms BERT-Feat+Softmax [29] on
all six datasets. This pairwise comparison suggests that for a $[CLS]$
representation from one layer, which is a vector of length 768, Kim-CNN [11]
serves as a better feature extractor than a simple Dense Network which outputs
logits to a Softmax function. We based our initial assumption, on recorded
results from previous work on topics that compare CNN-based architectures to
simple approaches such as Dense Networks, and the results depicted in TABLE
III show that projecting a $[CLS]$ representation from BERT to an embedding
space using a matrix of trainable parameters with a predefined shape, is
comparatively inadequate. In our model CNN-Trans-Enc, we use replicas of
CNN[CLS] separately, to produce $Q_{conv}$, $K_{conv}$, and $V_{conv}$ feature
maps which are used equivalently in the same way as in the original
Transformer-Encoder [32]. As discussed before, CNN[CLS] doesn’t perform well
on its own, but using its output feature maps as inputs to an Attention-based
model, seems to exploit similarities between $[CLS]$ representations from the
12 layers more efficiently. CNN-Trans-Enc outperforms the Trans-Enc [32], and
achieves higher accuracies by appreciable differences on all six datasets.
In particular, CNN-Trans-Enc (Our model), exceeds the Softmax classifier by
$\approx$ 2 points on the IMDB, $\approx$ 5 on Amazon-Polarity, $\approx$ 12
on the YELP-5 datasets, $\approx$ 7 points on SST-5, $\approx$ 3.24 on AG;s
news, and $\approx$ 0.4 on DBPedia-14. The CNN[CLS] blocks, used to produce
$Q_{conv}$, $K_{conv}$, and $V_{conv}$, contribute with a significant
improvement on all datasets, compared to the Trans-Enc results shown in TABLE
III. Since both models (CNN-Trans-Enc and Trans-Enc) share the same design and
hyperparameter settings, and only differ in the projection philosophy of the
input to the embedding space, it is safe to say that using a Convolutional
Neural Network inside a Transformer-Encoder [32], results in significant
performance gains for text classification using $[CLS]$ representations from
all layers of BERT [4].

### Statistical Significance Testing

In order to investigate the statistical significance of our results, we used
ASO (Almost Stochastic Order) [5], to compare all pairs of models based on
five random seeds, stacking CV average means obtained from empirical
evaluation from all datasets. Each comparison uses ASO with a confidence level
of $\alpha=0.05$ (before adjusting for all pair-wise comparisons using the
Bonferroni correction). $\epsilon^{min}$ values for all pairs are depicted in
TABLE IV. Looking at the Significance Test results, we can see that CNN-Trans-
Enc is the model with the overall best performance. We also notice that the
Softmax approach only outperforms the standalone CNN[CLS] model. Kim-CNN [11],
when used on top of the last $[CLS]$ representations from BERT [4], improves
over the Softmax classifier on large datasets.

##  VI Conclusion

In this paper, we investigate the effect of using $[CLS]$ representations from
all layers of BERT (which are kept frozen during training) on the
classification performance. We propose a CNN-Enhanced Transformer-Encoder to
obtain more generalizable representations through convolutional layers and
capture similarities between representations from all BERT layers and compute
an average using the Multi-Head Attention mechanism. Our model outperforms the
vanilla Transformer-Encoder and other approaches used in text classification
on six benchmark datasets. BERT+CNN-Trans-Enc achieves new state-of-the-art
results on YELP-5, Amazon-Polarity, and DBPedia-14 datasets exceeding large
models which are fine-tuned on the mentioned datasets, using limited hardware
resources. We validated our results using a Statistical Significance Test and
highlight our findings with empirical analysis of the architectures compared
in our experiment. In the future, we aim to explore architectures with more
generalization and characterization abilities, both for language and vision,
and incorporate hybrid techniques, that combine vision, and language-
representation methods, to facilitate and improve results for both domains in
multiple tasks.

## References

  * [1] Jader Abreu, Luis Fred, David Macêdo, and Cleber Zanchettin.  Hierarchical attentional hybrid neural networks for document classification.  CoRR, abs/1901.06610, 2019.
  * [2] Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni.  What you can cram into a single $&!#* vector: Probing sentence embeddings for linguistic properties.  In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2126–2136, Melbourne, Australia, July 2018. Association for Computational Linguistics.
  * [3] Alexis Conneau, Holger Schwenk, Loïc Barrault, and Yann Lecun.  Very deep convolutional networks for text classification.  In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 1, Long Papers, pages 1107–1116, Valencia, Spain, April 2017. Association for Computational Linguistics.
  * [4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  * [5] Rotem Dror, Segev Shlomov, and Roi Reichart.  Deep dominance-how to properly compare deep neural models.  In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2773–2785, 2019.
  * [6] Cunxiao Du, Zhaozheng Chen, Fuli Feng, Lei Zhu, Tian Gan, and Liqiang Nie.  Explicit interaction model towards text classification.  In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, AAAI’19/IAAI’19/EAAI’19. AAAI Press, 2019.
  * [7] Jeremy Howard and Sebastian Ruder.  Universal language model fine-tuning for text classification.  In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 328–339, Melbourne, Australia, July 2018. Association for Computational Linguistics.
  * [8] Ganesh Jawahar, Benoît Sagot, and Djamé Seddah.  What does BERT learn about the structure of language?  In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3651–3657, Florence, Italy, July 2019. Association for Computational Linguistics.
  * [9] Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao.  SMART: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2177–2190, Online, July 2020. Association for Computational Linguistics.
  * [10] Rie Johnson and Tong Zhang.  Deep pyramid convolutional neural networks for text categorization.  In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 562–570, Vancouver, Canada, July 2017. Association for Computational Linguistics.
  * [11] Yoon Kim.  Convolutional neural networks for sentence classification.  In EMNLP, pages 1746–1751. ACL, 2014.
  * [12] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In ICLR (Poster), 2015.
  * [13] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.  ALBERT: A lite BERT for self-supervised learning of language representations.  In ICLR. OpenReview.net, 2020.
  * [14] Hung Le, Truyen Tran, and Svetha Venkatesh.  Learning to remember more with less memorization.  In International Conference on Learning Representations, 2019.
  * [15] Jens Lehmann, Robert Isele, Max Jakob, Anja Jentzsch, Dimitris Kontokostas, Pablo N. Mendes, Sebastian Hellmann, Mohamed Morsey, Patrick van Kleef, Sören Auer, and Christian Bizer.  Dbpedia - a large-scale, multilingual knowledge base extracted from wikipedia.  Semantic Web, 6(2):167–195, 2015.
  * [16] Yuxiao Lin, Yuxian Meng, Xiaofei Sun, Qinghong Han, Kun Kuang, Jiwei Li, and Fei Wu.  BertGCN: Transductive text classification by combining GNN and BERT.  In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 1456–1462, Online, August 2021. Association for Computational Linguistics.
  * [17] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  Roberta: A robustly optimized bert pretraining approach.  arXiv preprint arXiv:1907.11692, 2019.
  * [18] H. Lu, Seth H. Huang, Tian Ye, and Xiuyan Guo.  Graph star net for generalized multi-task learning.  ArXiv, abs/1906.12330, 2019.
  * [19] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts.  Learning word vectors for sentiment analysis.  In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142–150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
  * [20] Tomás Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean.  Efficient estimation of word representations in vector space.  In ICLR (Workshop Poster), 2013.
  * [21] Shervin Minaee, Nal Kalchbrenner, Erik Cambria, Narjes Nikzad, Meysam Chenaghlu, and Jianfeng Gao.  Deep learning–based text classification: A comprehensive review.  ACM Comput. Surv., 54(3), apr 2021.
  * [22] Marcin Michał Mirończuk and Jarosław Protasiewicz.  A recent overview of the state-of-the-art elements of text classification.  Expert Systems with Applications, 106:36–54, 2018.
  * [23] Dmytro Mishkin and Jiri Matas.  All you need is a good init.  In ICLR (Poster), 2016.
  * [24] Manish Munikar, Sushil Shakya, and Aakash Shrestha.  Fine-grained sentiment classification using bert, 2019.
  * [25] Ameya Prabhu, Charles Dognin, and Maneesh Singh.  Sampling bias in deep active classification: An empirical study.  In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4058–4068, Hong Kong, China, November 2019. Association for Computational Linguistics.
  * [26] Hao Ren and Hong Lu.  Compositional coding capsule network with k-means routing for text classification.  arXiv preprint arXiv:1810.09177, 2018.
  * [27] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  $EMC^{2}$: 5th Edition Co-located with NeurIPS’19, abs/1910.01108, 2019.
  * [28] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts.  Recursive deep models for semantic compositionality over a sentiment treebank.  In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1631–1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.
  * [29] Chi Sun, Xipeng Qiu, Yige Xu, and Xuanjing Huang.  How to fine-tune bert for text classification?  In Maosong Sun, Xuanjing Huang, Heng Ji, Zhiyuan Liu, and Yang Liu, editors, Chinese Computational Linguistics, pages 194–206, Cham, 2019. Springer International Publishing.
  * [30] Zijun Sun, Chun Fan, Qinghong Han, Xiaofei Sun, Yuxian Meng, Fei Wu, and Jiwei Li.  Self-explaining structures improve NLP models.  CoRR, abs/2012.01786, 2020.
  * [31] Tan Thongtan and Tanasanee Phienthrakul.  Sentiment classification using document embeddings trained with cosine similarity.  In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: Student Research Workshop, pages 407–414, Florence, Italy, July 2019. Association for Computational Linguistics.
  * [32] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Advances in neural information processing systems, pages 5998–6008, 2017.
  * [33] Baoxin Wang.  Disconnected recurrent neural networks for text categorization.  In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2311–2320, Melbourne, Australia, July 2018. Association for Computational Linguistics.
  * [34] Sinong Wang, Han Fang, Madian Khabsa, Hanzi Mao, and Hao Ma.  Entailment as few-shot learner.  ArXiv, abs/2104.14690, 2021.
  * [35] Qizhe Xie, Zihang Dai, Eduard H. Hovy, Minh-Thang Luong, and Quoc V. Le.  Unsupervised data augmentation for consistency training.  34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020.
  * [36] 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.
  * [37] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le.  Xlnet: Generalized autoregressive pretraining for language understanding.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [38] Xiang Zhang, Junbo Zhao, and Yann LeCun.  Character-level convolutional networks for text classification.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
  * [39] Ye Zhang and Byron Wallace.  A sensitivity analysis of (and practitioners’ guide to) convolutional neural networks for sentence classification.  In Proceedings of the Eighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 253–263, Taipei, Taiwan, November 2017. Asian Federation of Natural Language Processing.


# Score-Based Generative Models
for Designing Binding Peptide Backbones

John D Boom
University of Cambridge
Cambridge, UK
jdb206[at]cam.ac.uk
&Matthew Greenig††footnotemark:
University of Cambridge
Cambridge, UK
mg989[at]cam.ac.uk
&Pietro Sormanni
University of Cambridge
Cambridge, UK
ps589[at]cam.ac.uk
Pietro Liò
University of Cambridge
Cambridge, UK
pl219[at]cam.ac.uk
These authors contributed equally to this work

###### Abstract

Score-based generative models (SGMs) have proven to be powerful tools for
designing new proteins. Designing proteins that bind a pre-specified target is
highly relevant to a range of medical and industrial applications. Despite the
flurry of new SGMs in the last year, there has been little systematic
exploration of the impact of design choices in SGMs for protein design. Here
we present LoopGen, a flexible SGM framework for the design of short binding
peptide structures. We apply our framework to design antibody binding loop
structures conditional on a target epitope and evaluate a variety of modelling
choices in SGM-based protein design. We demonstrate that modelling residue
orientations in addition to positions improves not only the quality of the
output structures but also their diversity. Additionally, we identify variance
schedules that result in significant performance improvements and observe
patterns that may motivate the development of better schedules for protein
design. Finally, we develop three novel tests to evaluate whether the model
generates structures that are appropriately conditioned on an epitope,
demonstrating that LoopGen’s generated structures are dependent on the
structure, sequence, and position of the epitope. Our findings will help guide
future development and evaluation of generative models for binding proteins.

##  1 Introduction

Rationally designing proteins in silico has the potential to unlock new
treatments for disease, accelerate scientific research, and enable new green
manufacturing technologies. Score-Based Generative Models (SGMs) have proven
to be effective tools for protein design [8, 21], dramatically outperforming
previous methods like generative adversarial networks [16]. While there was an
explosion of new models in the last year [21, 8, 15, 22, 23, 2, 13], there has
been little systematic exploration of the design choices behind SGMs in this
domain. Therefore, we sought to develop a framework for evaluating different
model design choices in the context of an important task: generating proteins
that bind a specified target. We apply our framework to design binding loop
structures for antibodies, a key class of biomolecules widely applied as
therapeutics, diagnostics, and research tools [19, 14]. Antibody binding is
mediated primarily through interactions between the target and short loop
regions called complementarity-determining-regions (CDRs), which often lack
secondary structure and can exhibit extreme structural diversity [6], posing
challenges in the generative modelling setting. Designing CDR structures for
binding is also challenging due to limited data availability; just under 8000
total antibody structures are available in the PDB, many of which are
redundant or lack a binding partner [5]. These challenges, combined with the
general utility of antibodies [19, 14], motivate the development of novel
methods for generating and evaluating CDR structures.

##  2 Methodology

We present LoopGen, an SGM framework for the generation of binding peptides.
Our framework allows for direct comparison of SGM design choices in the
context of protein design, such as the generation of residue orientations and
coordinates ("frames") versus coordinates alone, different estimator
architectures, and different choices of variance schedule. We apply our
framework to generate CDR loops conditioned on a target epitope. For our
experiments, we use a heterogeneous variant of the Geometric Vector Perceptron
(GVP) GNN architecture [9] to estimate scores for CDR residues in each
CDR/epitope complex. The complex is represented as a heterogeneous graph with
edges drawn between each residue and its $K=6$ nearest neighbors in both the
CDR and epitope. CDR residues are represented as either frames (rotations and
coordinates) or coordinates alone, and epitope residues are represented using
node features describing their sequence identity and backbone geometry.
Training was first conducted using a large dataset of CDR-like fragments
obtained from the PDB90 database [1] and subsequent finetuning was performed
using a smaller, higher-fidelity dataset (SAbDab) of real CDRs in complex with
epitopes [5], filtering for structures with <90% sequence identity. Self-
conditioning was performed as in Watson et al. [21] at a rate of 0.5.
Generation of novel loop structures was applied to 687 test set CDR/epitope
complexes from SAbDab, generating 10 loops conditioned on the epitope
structure and of the same length as the ground truth CDR, as well as
generating 10 loops for a transformed version of each epitope: permuted
(swapped), sequence scrambled, and translated. Please see Appendix A for more
information on our methods. Code and data are publicly available at:
https://github.com/mgreenig/loopgen.

##  3 Experiments

###  3.1 Evaluating the Impact of Frames

We first investigated the importance of incorporating orientation information
for each residue. We compared sets of CDR structures produced by generative
models using rotations and positions (frames) versus models using C$\alpha$
coordinates alone. Interestingly, the RMSD111Generally, our structures had
much higher RMSD than similar models [15]. We suspect this occurs for many
reasons, including not using sequence information, pretraining on a more
diverse dataset, and removing CDRs with ¿90% sequence similarity. of the
generated structures compared to the ground truth was indistinguishable
between the models; however, structures generated by the coordinates-only
model displayed a more significant deviation from the ground truth distances
between adjacent C$\alpha$ atoms along the chain (Fig. 1, middle). To analyze
the diversity of the generated structures, we computed the mean pairwise RMSD
between the 10 generated structures against each test set epitope for both
models, removing duplicates. The coordinates-only model had dramatically lower
diversity as measured by pairwise RMSD (Figure 1, right). These findings
suggest that the coordinates-only model produces less biochemically plausible,
more homogeneous structures, and, more generally, that RMSD to the ground
truth structure fails to capture salient features of generated CDRs. We
include examples of generated structures in Appendix B.

Figure 1: Analyzing how incorporating rotational information improves model
performance. From left to right: A comparison of the RMSD of the generated
structures compared to the ground truth, A histogram of the distances between
Cα, and a comparison of the pairwise RMSD between sets of generated structures
(higher values indicate greater structural diversity).

###  3.2 Exploring how Diffusion Schedules Affect Performance

Table 1: Choice of Variance Schedule Dramatically Affects Structure Quality

Trans. Sched. |  Rot. Sched. | RMSD (Å) |  Internal Clashes (%) |  Bond Length (%) |  Bond Angle (%) |  Epi.-CDR Clash |
---|---|---|---|---|---|---|---
|  Any Struct. Flaw (%) |  |  |  |  | (%) |
Lin. | Log. |  4.98 $\pm$ 2.14 | 0.3 | 20.6 | 3.9 | 3.3 | 22.4
Quad. | Log. |  4.93 $\pm$ 2.15 | 0.0 | 6.3 | 0.6 | 2.7 | 8.7
Log. | Log. |  4.85 $\pm$ 2.08 | 88.9 | 96.2 | 43.9 | 5.7 | 97.1
Sig. | Log. |  4.93 $\pm$ 2.18 | 0.6 | 6.1 | 1.8 | 3.4 | 9.2
Lin. | Lin. |  5.04 $\pm$ 2.08 | 1.0 | 29.4 | 4.0 | 3.4 | 30.9
Lin. | Quad. |  5.13 $\pm$ 2.17 | 0.8 | 18.6 | 1.7 | 3.5 | 20.5

Score-based generative models of frames for protein design are unusual because
noise must be applied over both the orientation and coordinate of each
residue. These two processes introduce complexity because the relative
position between a residue and its neighbors informs what orientations are
possible and vice versa. Variance schedules have been shown to have a large
impact on diffusion models for image generation [11]; however, to our
knowledge, there is almost no published data on how different noising
schedules for translations and rotations affect frame generative models.

RMSD to the ground truth structure is the standard metric for evaluating the
quality of generated CDR structures [15, 22]. We found that models with
different variance schedules showed minimal variation in ground truth RMSD but
varied significantly in their ability to generate physicochemically plausible
structures (Table LABEL:tab:var_schedules), where we defined structural
violations using the violation loss functions from AlphaFold2 [10].
Strikingly, ranking the translation variance schedules from the "slowest"
destruction of information to the "fastest" results in the exact ordering of
performance (Figure 6, Appendix). The quadratic translation schedule combined
with a logarithmic schedule for rotations exhibited the lowest rate of
structural violations. Despite showing poor correspondence between generated
and ground truth CDR structures, the generated loops still satisfy key
criteria as both valid loops and potential binders, obeying the correct
dihedral angle distribution at the correct distance from the epitope (Figure
2). These findings suggest that RMSD may not be an appropriate metric for
assessing generative models of CDR loops, which have notably high structural
diversity.

Figure 2: Left: Ramachandran Distribution of the generated CDRs compared to
ground truth. Right: The minimum distance between an $C_{\alpha}$ on the CDR
and the epitope for the ground truth and generated CDRs. Stratification by
length reveals better performance on shorter CDRs.

###  3.3 Evaluating Generative Models for Binding Peptides

Given that ground truth RMSD appears to have many limitations as a metric, we
searched for additional metrics to evaluate the quality of the model’s output.
Ultimately, we determined that a major and neglected objective is that the
generated structures should be clearly conditioned on the provided epitope. To
evaluate the dependence on the epitope, we generated sets of 10 CDRs after
swapping each epitope (aligning its principal components to the original
epitope), scrambling the epitope sequence (permutation of residue identities),
and translating the epitope 20 Å away from the CDR. Then, we computed the mean
pairwise RMSD between each of these sets of generated CDRs and the set of CDRs
generated for the wild-type epitope. Larger mean pairwise RMSDs indicate
greater structural differences between sets of CDRs. As shown in Figure 3, we
see that all three perturbations result in a large increase in pairwise RMSD,
confirming that, in general, the model’s outputs are highly dependent on the
epitope’s structure, sequence, and positioning.

Figure 3: Analyzing the dependence of the generated CDR structures on the
epitope. A set of 10 structures are generated for 687 test set epitopes, i.e.
wild-type (WT) epitopes. Another 10 structures are generated under three types
of perturbation to each epitope: permuted (swapping the WT epitope with a
random epitope), scrambled (permuting residue identities within the WT epitope
structure), and translated (translating the epitope 20 Å away from the CDR
along the vector connecting their centers of masses). Pairwise RMSDs are
calculated between the set of CDRs generated for the WT epitope and the sets
generated under all conditions for each epitope in the test set.

##  4 Discussion and Conclusion

Score-based generative models have facilitated in silico protein design;
however, relatively little is known about the key components affecting their
performance. In this work we conduct, to our knowledge, the first direct
comparison between modelling entire residue frames versus C$\alpha$
coordinates alone. We hypothesised that modelling frames would provide
specific benefits in designing binding peptides, which have highly variable
structures and thus significant degrees of freedom in their dihedral angles,
and indeed we show that structures generated as frames are significantly more
diverse. Next, we compare models trained under different variance schedules
and show that, while schedules do not have a significant effect on the
commonly-used ground truth RMSD metric, they do significantly influence the
rate of structural violations. We also introduce additional metrics to assess
the quality of generated binder structures by evaluating the extent to which
they are conditioned on the epitope.

Our experiments identify multiple promising avenues for further investigation.
While we study the incorporation of orientation information, we only do so for
the IGSO(3) diffusion framework [23, 21, 15] (see Appendix A for more
details). However, future research may benefit from exploring different forms
of generative models over rotations (e.g. [2]). Furthermore, although we
showed significant performance variation across both translation and rotation
variance schedules, a more extensive evaluation of the entire space of
schedule combinations may identify configurations that yield even better
results. Finally, we note that our model alone is not sufficient for de novo
CDR design because it only generates CDR structures, independent of sequences.
Previous work tested conducting structure and sequence generation concurrently
(at the cost of lower structural diversity) [15], and future research would
benefit from comparing methods for incorporating sequence generation directly
into the SGM framework to performing sequence design post-hoc [21].

## References

  * Aguilar Rangel et al. [2022] Mauricio Aguilar Rangel, Alice Bedwell, Elisa Costanzi, Ross J. Taylor, Rosaria Russo, Gonçalo J. L. Bernardes, Stefano Ricagno, Judith Frydman, Michele Vendruscolo, and Pietro Sormanni.  Fragment-based computational design of antibodies targeting structured epitopes.  _Science Advances_ , 8(45):eabp9540, November 2022\.  doi: 10.1126/sciadv.abp9540.  URL https://www.science.org/doi/10.1126/sciadv.abp9540.  Publisher: American Association for the Advancement of Science.
  * Anand and Achim [2022] Namrata Anand and Tudor Achim.  Protein Structure and Sequence Generation with Equivariant Denoising Diffusion Probabilistic Models, May 2022.  URL http://arxiv.org/abs/2205.15019.  arXiv:2205.15019 [cs, q-bio].
  * Badaczewska-Dawid et al. [2020] Aleksandra E. Badaczewska-Dawid, Andrzej Kolinski, and Sebastian Kmiecik.  Computational reconstruction of atomistic protein structures from coarse-grained models.  _Computational and Structural Biotechnology Journal_ , 18:162–176, 2020.  ISSN 2001-0370.  doi: https://doi.org/10.1016/j.csbj.2019.12.007.  URL https://www.sciencedirect.com/science/article/pii/S2001037019305537.
  * De Bortoli et al. [2022] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet.  Riemannian Score-Based Generative Modelling, November 2022.  URL http://arxiv.org/abs/2202.02763.  arXiv:2202.02763 [cs, math, stat].
  * Dunbar et al. [2014] James Dunbar, Konrad Krawczyk, Jinwoo Leem, Terry Baker, Angelika Fuchs, Guy Georges, Jiye Shi, and Charlotte M. Deane.  SAbDab: the structural antibody database.  _Nucleic Acids Research_ , 42(D1):D1140–D1146, January 2014.  ISSN 0305-1048.  doi: 10.1093/nar/gkt1043.  URL https://doi.org/10.1093/nar/gkt1043.
  * Fernández-Quintero et al. [2020] Monica L. Fernández-Quintero, Martin C. Heiss, Nancy D. Pomarici, Barbara A. Math, and Klaus R. Liedl.  Antibody CDR loops as ensembles in solution vs. canonical clusters from x-ray structures.  _mAbs_ , 12(1), January 2020.  doi: 10.1080/19420862.2020.1744328.  URL https://doi.org/10.1080/19420862.2020.1744328.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising Diffusion Probabilistic Models, December 2020.  URL http://arxiv.org/abs/2006.11239.  arXiv:2006.11239 [cs, stat].
  * Ingraham et al. [2022] John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent Xue, Fritz Obermeyer, Andrew Beam, and Gevorg Grigoryan.  Illuminating protein space with a programmable generative model, December 2022.  URL https://www.biorxiv.org/content/10.1101/2022.12.01.518682v1.  Pages: 2022.12.01.518682 Section: New Results.
  * Jing et al. [2021] Bowen Jing, Stephan Eismann, Pratham N. Soni, and Ron O. Dror.  Equivariant Graph Neural Networks for 3D Macromolecular Structure, July 2021.  URL http://arxiv.org/abs/2106.03843.  arXiv:2106.03843 [cs, q-bio].
  * Jumper et al. [2021] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis.  Highly accurate protein structure prediction with AlphaFold.  _Nature_ , 596(7873):583–589, August 2021.  ISSN 1476-4687.  doi: 10.1038/s41586-021-03819-2.  URL https://www.nature.com/articles/s41586-021-03819-2.  Number: 7873 Publisher: Nature Publishing Group.
  * Karras et al. [2022] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the Design Space of Diffusion-Based Generative Models, October 2022.  URL http://arxiv.org/abs/2206.00364.  arXiv:2206.00364 [cs, stat].
  * Leach et al. [2022] Adam Leach, Sebastian M Schmon, Matteo T Degiacomi, and Chris G Willcocks.  DENOISING DIFFUSION PROBABILISTIC MODELS ON SO(3) FOR ROTATIONAL ALIGNMENT, 2022.
  * Lin and AlQuraishi [2023] Yeqing Lin and Mohammed AlQuraishi.  Generating Novel, Designable, and Diverse Protein Structures by Equivariantly Diffusing Oriented Residue Clouds, June 2023.  URL http://arxiv.org/abs/2301.12485.  arXiv:2301.12485 [cs, q-bio].
  * Lu et al. [2020] Ruei-Min Lu, Yu-Chyi Hwang, I-Ju Liu, Chi-Chiu Lee, Han-Zen Tsai, Hsin-Jung Li, and Han-Chung Wu.  Development of therapeutic antibodies for the treatment of diseases.  _Journal of Biomedical Science_ , 27(1):1, January 2020.  ISSN 1423-0127.  doi: 10.1186/s12929-019-0592-z.  URL https://doi.org/10.1186/s12929-019-0592-z.
  * Luo et al. [2022] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma.  Antigen-Specific Antibody Design and Optimization with Diffusion-Based Generative Models for Protein Structures, October 2022\.  URL https://openreview.net/forum?id=jSorGn2Tjg.
  * Repecka et al. [2021] Donatas Repecka, Vykintas Jauniskis, Laurynas Karpus, Elzbieta Rembeza, Irmantas Rokaitis, Jan Zrimec, Simona Poviloniene, Audrius Laurynenas, Sandra Viknander, Wissam Abuajwa, Otto Savolainen, Rolandas Meskys, Martin K. M. Engqvist, and Aleksej Zelezniak.  Expanding functional protein sequence spaces using generative adversarial networks.  _Nature Machine Intelligence_ , 3(4):324–333, April 2021.  ISSN 2522-5839.  doi: 10.1038/s42256-021-00310-5.  URL https://www.nature.com/articles/s42256-021-00310-5.  Number: 4 Publisher: Nature Publishing Group.
  * Solà et al. [2021] Joan Solà, Jeremie Deray, and Dinesh Atchuthan.  A micro lie theory for state estimation in robotics, 2021.
  * Song et al. [2021] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-Based Generative Modeling through Stochastic Differential Equations, February 2021.  URL http://arxiv.org/abs/2011.13456.  arXiv:2011.13456 [cs, stat].
  * Sormanni et al. [2018] Pietro Sormanni, Francesco A. Aprile, and Michele Vendruscolo.  Third generation antibody discovery methods: in silico rational design.  _Chemical Society Reviews_ , 47(24):9137–9157, 2018.  doi: 10.1039/c8cs00523k.  URL https://doi.org/10.1039/c8cs00523k.
  * Trippe et al. [2023] Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola.  Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem, March 2023.  URL http://arxiv.org/abs/2206.04119.  arXiv:2206.04119 [cs, q-bio, stat].
  * Watson et al. [2022] Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker.  Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models, December 2022.  URL https://www.biorxiv.org/content/10.1101/2022.12.09.519842v1.  Pages: 2022.12.09.519842 Section: New Results.
  * Xie et al. [2023] Xuezhi Xie, Jin Sub Lee, Dongki Kim, Jaehyeong Jo, Jisun Kim, and Philip M Kim.  Antibody-SGM: Antigen-Specific Joint Design of Antibody Sequence and Structure using Diffusion Models, July 2023.
  * Yim et al. [2023] Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola.  SE(3) diffusion model with application to protein backbone generation, May 2023.  URL http://arxiv.org/abs/2302.02277.  arXiv:2302.02277 [cs, q-bio, stat].

##  Appendix A LoopGen Technical Details

Score-based generative modelling via stochastic differential equations (SDEs)
[18] models the process of adding noise to data with the following forward
SDE:

$\mathrm{d}{\bf x}=f({\bf x},t)\mathrm{d}t+g(t)\mathrm{d}{\bf w}$ (1)

Where ${\bf x}$ are the data and ${\bf w}$ is a brownian motion. The forward
SDE has a corresponding reverse SDE, which models how data are generated from
noise:

$\mathrm{d}{\bf x}=\left[f({\bf x},t)-g(t)^{2}\nabla\log{p_{t}({\bf
x})}\right]\mathrm{d}t+g(t)\mathrm{d}\bar{{\bf w}}$ (2)

Where $\mathrm{d}t$ is an infinitesimal negative time step and $\bar{{\bf w}}$
is a brownian motion for time moving in the reverse direction. The score
$\nabla\log{p_{t}({\bf x})}$ is approximated with an estimator
$s_{\theta}({\bf x},t)$ by optimising a score-matching objective, typically of
the form:

$\mathcal{L}(\theta)=\mathbb{E}_{{\bf
x},t}\left[\lambda(t)\lVert\nabla\log{p_{t}({\bf x})}-s_{\theta}({\bf
x},t)\rVert_{2}^{2}\right]$ (3)

Where the coefficients $\lambda(t)$ are usually chosen to balance the expected
magnitudes of the ground truth scores, i.e.
$\lambda(t)=1/\mathbb{E}\left[\lVert\nabla\log{p_{t}({\bf
x})}\rVert^{2}\right]$ [18].

An important extension of SGM that has found applications in protein design is
the Riemannian SGM framework, which extends the forward and reverse SDEs to
Riemannian manifolds on which appropriate Brownian motion and probability
densities can be defined [4], including the manifold of three-dimensional
proper rotations $SO(3)$. In proteins, the orientation of each amino acid can
be represented using a three-dimensional rotation matrix [10], and recent
generative modelling approaches have attempted to generate novel protein
structures by formulating SGM approaches for both the rotational and
translational component of each residue [21, 23]. The combination of a
rotation and a translation in this context is known as a frame and is a member
of the Lie group $SE(3)$. Yim et al. [23] formalised the theory of SGM for
frames in their recent work in which they propose a protein generative
modelling approach known as FrameDiff.

To define generative models on $SE(3)$, Yim et al. [23] note that $SE(3)$ may
be identified as $SO(3)\times\mathbb{R}^{3}$ [23] under a proper choice of
inner product, allowing diffusion processes over translations
($\mathbb{R}^{3}$) and rotations ($SO(3)$) to be performed separately. Let
$\mathbf{T}^{(t)}=({\bf R}^{(t)},\mathbf{X}^{(t)})$, where
$\mathbf{X}^{(t)}\in\mathbb{R}^{3}$ represents a translation and ${\bf
R}^{(t)}\in SO(3)$ represents a rotation at time $t$. Let
$\mathbf{B}_{\mathcal{M}}^{(t)}$ represent Brownian motion on a manifold
$\mathcal{M}$.222Please see De Bortoli et al. [4] and section 2 of Yim et al.
[23] for more details on deriving the forward and reverse processes on
$\mathcal{M}$. Then, they derive the following equation for the forward
diffusion process:

$\mathrm{d}\mathbf{T}^{(t)}=[0,-\frac{1}{2}\mathrm{P}\mathbf{X}^{(t)}]\mathrm{d}t+\left[\mathrm{d}\mathbf{B}_{\mathrm{SO}(3)}^{(t)},\mathrm{dP}\mathbf{B}_{\mathbb{R}^{3}}^{(t)}\right]$
(4)

where $\mathrm{P}$ is a projection matrix centering the translations at the
origin. Removing the center of mass maintains SE(3)-invariance. Notably, there
is no drift term for the rotations ${\bf R}^{(t)}$. Similarly, the reverse
process is defined as

$\displaystyle\mathrm{d}{\mathbf{R}}^{(t)}=\nabla_{\bf R}\log
p_{t}({\mathbf{T}}^{(t)})\mathrm{d}t+\mathrm{d}\mathbf{B}_{\mathrm{SO}(3)}^{(t)}$
(5)
$\displaystyle\mathrm{d}{\mathbf{X}}^{(t)}=\mathrm{P}\left(\frac{1}{2}{\mathbf{X}}^{(t)}+\nabla_{\bf
x}\log
p_{t}({\mathbf{T}}^{(t)})\right)\mathrm{d}t+\mathrm{Pd}\mathbf{B}_{\mathbb{R}^{3}}^{(t)}$

SGM on the manifold of frames allows for nearly complete generation of protein
backbones - specifically the positions of the N, C$\alpha$, and C atoms - but
leverages a significant amount of additional machinery that poses technical
challenges. On the other hand, generation of coordinates (translations) alone
is more straightforwardly defined within the SGM framework, and hence many
generative modelling frameworks for proteins to date have focused solely on
generating C$\alpha$ atom positions [13, 20], which can serve as reasonable
structural templates in many cases. In the context of peptide binders, and
particularly binding loops, we hypothesised that generative modelling of
frames would provide unique benefits due to the conformational flexibility and
lack of secondary structure associated with loop regions; accurate
reconstruction of the entire backbone from C$\alpha$ positions often relies on
secondary structure modelling [3]. Hence we sought to conduct a controlled
comparison of generative models for binding loop design, comparing SGM
approaches that operate on frames to those that model coordinates only. We
combine these two modalities into an SGM framework known as LoopGen.

LoopGen takes inspiration from FrameDiff [23] and focuses specifically on
adapting their proposed SGM framework to make score-based generative models
for frames more flexible and more straightforward to train. FrameDiff relies
on the Invariant Point Attention (IPA) architecture [10] to predict the ground
truth frame ${\bf T}^{(0)}=({\bf R}^{(0)},{\bf x}^{(0)})$ for each residue,
and calculates the predicted score at each time step by backpropagating
through the score function:

$\displaystyle s_{\theta}({\bf T}^{(t)},t)$
$\displaystyle=\nabla\log{p_{t}({\bf T}^{(t)}|\hat{{\bf T}}^{(0)})}$ (6)
$\displaystyle=\left[\nabla\log{p_{t}({\bf R}^{(t)}|\hat{{\bf
R}}^{(0)})},\nabla\log{p_{t}({\bf x}^{(t)}|\hat{{\bf x}}^{(0)})}\right]$ (7)

Where the score for a frame is defined separately for rotations and
translations. Importantly, this formulation restricts the form of the
estimator $s_{\theta}({\bf T}^{(t)},t)$ to functions that directly predict
frames $\hat{{\bf T}}^{(0)}=(\hat{{\bf R}}^{(0)},\hat{{\bf x}}^{(0)})$.
However, we note that for translations ${\bf x}\in\mathbb{R}^{3}$, the score
$\nabla\log{p_{t}({\bf x})}$ for a Gaussian density can be predicted directly
as an SO(3)-equivariant 3-vector, and likewise for rotations ${\bf R}\in
SO(3)$, the isomorphism between the Lie algebra $\mathfrak{so}(3)$ and
$\mathbb{R}^{3}$ can be exploited to represent and predict the rotational
score $\nabla\log{p_{t}({\bf R})}$ directly as an element of $\mathbb{R}^{3}$.
For the IGSO(3) distribution [12], we note that this formulation has the key
advantage of not requiring continuous re-estimation of the infinite sum in the
density function (8) during the training process; instead, scores can be
computed once for a discretized support, cached, and returned as constants
during training since backpropagation through the score function is no longer
required. Furthermore, the reparameterization of the scores as elements of
$\mathbb{R}^{3}$ allows a much wider variety of estimator architectures to be
used, beyond functions that directly output frames.

###  A.1 Representing the IGSO(3) score

The IGSO(3) distribution is defined over axis-angle parameterizations of
three-dimensional rotations. These have the form $\omega\hat{\theta}$, where
$\hat{\theta}$ is a unit-length axis of rotation and $\omega$ is the angle of
rotation. The density is uniform over axes $\hat{\theta}$, and the density
over angles of rotation $p(\omega)$ has the following form:

$\displaystyle f(\omega)$
$\displaystyle=\sum_{l=0}^{\infty}{(2l+1)\exp{[-l(l+1)\sigma^{2}]}\frac{\sin{([l+\frac{1}{2}]\omega)}}{\sin{(\frac{\omega}{2})}}}$
(8) $\displaystyle p(\omega)$
$\displaystyle=\frac{1-\cos{(\omega)}}{\pi}f(\omega)$ (9)

where $\sigma^{2}$ is a variance parameter and the normalising constant
$\frac{1-\cos{\omega}}{\pi}$ ensures that the joint distribution over all axes
and angles integrates to 1 [12]. Adding noise at time $t$ to a sample rotation
${\bf R}^{(0)}$ then involves sampling an incremental rotation from the
corresponding time step distribution parameterised by variance
$\varsigma^{2}_{t}$, and then right multiplying with the sample rotation to
obtain a sample from $p({\bf R}^{(t)}|{\bf R}^{(0)})$:

$\displaystyle\mathcal{R}$ $\displaystyle\sim\text{IGSO3}(\mu={\bf
I},\sigma^{2}=\varsigma^{2}_{t})$ (10) $\displaystyle{\bf R}^{(t)}$
$\displaystyle=\mathcal{R}{\bf R}^{(0)}$ (11)

The score of the noised data distribution $p({\bf R}^{(t)}|{\bf R}^{(0)})$
then has the following form:

$\nabla\log{p({\bf R}^{(t)}|{\bf R}^{(0)})}=\frac{\mathrm{Log}{\left([{\bf
R}^{(0)}]^{\top}{\bf R}^{(t)}\right)}}{\omega\left([{\bf R}^{(0)}]^{\top}{\bf
R}^{(t)}\right)}\frac{f^{\prime}\left(\omega([{\bf R}^{(0)}]^{\top}{\bf
R}^{(t)})\right)}{f\left(\omega([{\bf R}^{(0)}]^{\top}{\bf R}^{(t)})\right)}$
(12)

Where we use $\omega({\bf R})$ to denote the angular component $\omega$ of an
axis-angle representation $\omega\hat{\theta}$ of a rotation matrix ${\bf R}$,
and follow the convention of Solà et al. [17], using Log to indicate the
logarithmic matrix map from a Lie group to Euclidean space, i.e.
$\text{Log}:SO(3)\rightarrow\mathbb{R}^{3}$ here. Note that we do not follow
previous works [23, 21] which represent $\nabla\log{p({\bf R}^{(t)}|{\bf
R}^{(0)})}$ as an element of the tangent space at the noised rotation
$\mathcal{T}_{{\bf R}^{(t)}}=\\{{\bf
R}^{(t)}\Theta:\Theta\in\mathfrak{so}(3)\\}$; instead, the score is always
represented at the tangent space at the identity rotation (i.e.
$\mathfrak{so}(3)$), hence no further processing is needed before passing
predicted scores through the exponential map to update rotations in the
generation process. We see that $\nabla\log{p({\bf R}^{(t)}|{\bf R}^{(0)})}$
in this representation is SO(3)-invariant since $({\bf R}^{(0)})^{\top}{\bf
R}^{(t)}$ is SO(3)-invariant.

###  A.2 Training details

To stabilise the score-matching loss for rotations we use coefficients
$\lambda(t)=1/\mathbb{E}\left[\lVert\nabla\log{p_{t}({\bf R}^{(t)}|{\bf
R}^{(0)})}\rVert^{2}\right]$ that we estimate via a monte carlo procedure,
calculating the empirical mean norm of the scores for each time point over
50,000 uniformly-sampled rotations ${\bf R}^{(0)}$. For translations, we found
that scaling the ground truth score by a factor of $-\sigma_{t}$ \- where
$\sigma_{t}=\sqrt{\mathrm{Var}[{\bf x}^{(t)}|{\bf x}^{(0)}]}$ is the standard
deviation of noised translations at time $t$ \- stabilised training
significantly, noting that:

$\displaystyle-\sigma_{t}\nabla\log{p_{t}({\bf x}^{(t)}|{\bf x}^{(0)})}$
$\displaystyle=-\sigma_{t}\frac{{\bf\mu}-{\bf x}^{(t)}}{{\sigma_{t}}^{2}}$
(13) $\displaystyle=\frac{{\bf x}^{(t)}-{\bf\mu}}{{\sigma_{t}}}$
$\displaystyle\sim\mathcal{N}(0,{\bf I})$

A single training step involving sampling and loss calculation is shown below
(13) for an input set of frames $\\{({\bf R}_{i}^{(0)},{\bf
x}_{i}^{(0)})\\}_{i=1}^{N}$, representation of the epitope ${\bf Z}$, score
estimator $s_{\theta}$, rotation and translation coefficient schedules
$\beta^{r}$ and $\beta^{x}$ (from which variances are calculated), rotation
score matching coefficient $\lambda^{r}$, and time step $t\in[0,1)$. Integrals
are all approximated using the trapezoid rule. Our noising regime corresponds
to the variance-preserving SDE from Song et al. [18].

Algorithm 1 LoopGen Training Step

$\\{({\bf R}_{i}^{(0)},{\bf x}_{i}^{(0)})\\}_{i=1}^{N},{\bf
Z},s_{\theta},\beta^{r},\beta^{x},\lambda^{r},t$

$\varsigma^{2}_{t}\leftarrow 1-\exp{(-\int_{0}^{t}{\beta^{r}(s)\
\mathrm{d}s})}$ $\triangleright$ $\mathrm{Var}[{\bf R}^{(t)}|{\bf R}^{(0)}]$

$\sigma^{2}_{t}\leftarrow 1-\exp{(-\int_{0}^{t}{\beta^{x}(s)\ \mathrm{d}s})}$
$\triangleright$ $\mathrm{Var}[{\bf x}^{(t)}|{\bf x}^{(0)}]$

for i = 1, …, N do

$\mathcal{R}\sim\text{IGSO3}(\mu={\bf I},\sigma^{2}=\varsigma^{2}_{t})$
$\triangleright$ Sample rotational noise

${\bf R}^{(t)}_{i}\leftarrow\mathcal{R}{\bf R}^{(0)}_{i}$ $\triangleright$ Get
noised rotation

${\bf x}^{(t)}_{i}\sim\mathcal{N}(\mu={\bf
x}^{(0)}_{i}\exp{(-\frac{1}{2}\int_{0}^{t}{\beta^{x}(s)\
\mathrm{d}s})},\sigma^{2}=\sigma^{2}_{t})$ $\triangleright$ Sample noised
translation

end for

$\\{(\hat{\bf y}^{r}_{i},\hat{\bf y}^{x}_{i})\\}_{i=1}^{N}\leftarrow
s_{\theta}(\\{({\bf R}_{i}^{(t)},{\bf x}_{i}^{(t)})\\}_{i=1}^{N},{\bf Z},t)$
$\triangleright$ Predict scores for each residue

$\mathcal{L}^{r}\leftarrow\lambda^{r}(t)\sum_{i=1}^{N}{\lVert\nabla\log{p_{t}}({\bf
R}^{(t)}_{i}|{\bf R}^{(0)}_{i})-\hat{\bf y}^{r}_{i}\rVert^{2}}$
$\triangleright$ Rotation score loss

$\mathcal{L}^{x}\leftarrow\sum_{i=1}^{N}{\lVert-\sigma_{t}\nabla\log{p_{t}}({\bf
x}^{(t)}_{i}|{\bf x}^{(0)}_{i})-\hat{\bf y}^{x}_{i}\rVert^{2}}$
$\triangleright$ Translation score loss

$\mathcal{L}\leftarrow\mathcal{L}^{r}+\mathcal{L}^{x}$ $\triangleright$
Combined loss

return $\mathcal{L}$

For generation (19), we broadly follow the procedures outlined by Ho et al.
[7] for translations and Yim et al. [23] for rotations, using the notation Exp
to denote the exponential map from euclidean space to a Lie group, where
$\text{Exp}:\mathbb{R}^{3}\rightarrow SO(3)$ in this case. We found
empirically that scaling the score term in the rotational update $\gamma
g_{t}^{2}\hat{\bf y}_{i}^{r}$ by a factor of 2 improved the quality of
generated structures. For our experiments we used number of time steps $T=100$
and a noise scaling term of $\zeta=0.2$. We found that reasonable structures
were generated for values up to $\zeta=0.5$.

Algorithm 2 LoopGen Generation

${\bf Z},s_{\theta},\beta^{r},\beta^{x},\zeta,T,N$

$\gamma\leftarrow\frac{1}{T}$

for i = 1, … , N do $\triangleright$ Sample frames for N residues

Sample ${\bf R}_{i}^{(1)}\sim\text{UniformSO3}$

Sample ${\bf x}_{i}^{(1)}\sim\mathcal{N}(0,{\bf I})$

end for

for $t=1-\gamma,1-2\gamma,...,0$ do

$g_{t}\leftarrow\sqrt{\beta^{r}(t)}$ $\triangleright$ Diffusion coefficient
for VP-SDE

$\alpha_{t}\leftarrow 1-\exp{(-\int_{t}^{t+\gamma}{\beta^{x}(s)\
\mathrm{d}s})}$ $\triangleright$ $\mathrm{Var}[{\bf x}^{(t+\gamma)}|{\bf
x}^{(t)}]$

$\sigma^{2}_{t}\leftarrow 1-\exp{(-\int_{0}^{t}{\beta^{x}(s)\ \mathrm{d}s})}$
$\triangleright$ $\mathrm{Var}[{\bf x}^{(t)}|{\bf x}^{(0)}]$

$\\{(\hat{\bf y}^{r}_{i},\hat{\bf y}^{x}_{i})\\}_{i=1}^{N}\leftarrow
s_{\theta}(\\{({\bf R}_{i}^{(t+\gamma)},{\bf
x}_{i}^{(t+\gamma)})\\}_{i=1}^{N},{\bf Z},t)$ $\triangleright$ Predict scores
for each residue

for i = 1, … N do

Sample $\epsilon^{x}\sim\mathcal{N}(0,{\bf I})$ $\triangleright$ Noise for
translations

Sample $\epsilon^{r}\sim\mathcal{N}(0,{\bf I})$ $\triangleright$ Noise for
rotations

${\bf x}^{(t)}_{i}\leftarrow\frac{{\bf
x}^{(t+\gamma)}_{i}-\alpha_{t}}{\sigma_{t}\sqrt{1-\alpha_{t}}}{\hat{\bf
y}}_{i}^{x}+\zeta\sqrt{\alpha_{t}}\epsilon^{x}$ $\triangleright$ Translation
update

${\bf R}^{(t)}_{i}\leftarrow{\bf R}^{(t+\gamma)}_{i}\text{Exp}(2\gamma
g_{t}^{2}\hat{\bf y}_{i}^{r}+\zeta\sqrt{\gamma}g_{t}\epsilon^{r})$
$\triangleright$ Rotation update

end for

end for

return $\\{({\bf R}_{i}^{(0)},{\bf x}_{i}^{(0)})\\}_{i=1}^{N}$

###  A.3 Tests for epitope dependence

Here we outline the procedures for evaluating the conditionality of generated
structures on the epitope. The fundamental principle we use is that dependence
on the epitope can be measured by estimating similarity between sets of
structures generated under different perturbations of the epitope. Let $\bf X$
be a random variable representing a generated structure and $\bf Y$ another
generated structure. The quantity of interest is then the expected RMSD
between $\bf X$ and $\bf Y$, i.e. the mean pairwise RMSD (mpRMSD):

$\displaystyle\mathrm{mpRMSD}=\mathbb{E}_{\bf X}[\mathbb{E}_{\bf
Y}[\mathrm{RMSD}({\bf X},{\bf Y})]]$ (14)

For a given epitope ${\bf Z}$ and a perturbed version of that epitope ${\bf
Z}^{\prime}$, we compare the mpRMSD for samples drawn from $p({\bf X}|{\bf
Z})$ and $p({\bf Y}|{\bf Z}^{\prime})$ to those generated under identical
epitope conditions, i.e. $p({\bf X}|{\bf Z})$ and $p({\bf Y}|{\bf Z})$. For a
given set of structures generated for a WT epitope ${\bf Z}$, with sets of
residue coordinates $\left\\{\\{{\bf x}^{(i)}_{1},...,{\bf
x}^{(i)}_{N}\\}\right\\}_{i=1}^{M}$, the mpRMSD within the set of structures
is estimated as:

$\mathrm{mpRMSD}\approx\frac{2}{M(M-1)}\sum_{i=1}^{M}{\sum_{j=i+1}^{M}\sqrt{\frac{1}{N}{\sum_{n=1}^{N}{\lVert{\bf
x}^{(i)}_{n}-{\bf x}^{(j)}_{n}}\rVert^{2}}}}$ (15)

Where the scaling factor is the reciprocal of $M$ choose 2, calculating the
mean RMSD over all possible unique pairings of structures in the set. We use
the within-group mpRMSD to measure the expected variation between CDR loop
structures generated for a constant epitope ("WT epitope vs WT epitope",
figure 3). The mpRMSD between different sets of structures - where one set
$\left\\{\\{{\bf x}^{(i)}_{1},...,{\bf x}^{(i)}_{N}\\}\right\\}_{i=1}^{M}$ is
generated for the WT epitope ${\bf Z}$ and another set $\left\\{\\{{\bf
x^{\prime}}^{(j)}_{1},...,{\bf x^{\prime}}^{(j)}_{N}\\}\right\\}_{j=1}^{K}$ is
generated for the perturbed epitope ${\bf Z}^{\prime}$ \- is similarly
estimated:

$\mathrm{mpRMSD}\approx\frac{1}{MK}\sum_{i=1}^{M}{\sum_{j=1}^{K}\sqrt{\frac{1}{N}{\sum_{n=1}^{N}{\lVert{\bf
x}^{(i)}_{n}-{\bf x^{\prime}}^{(j)}_{n}}\rVert^{2}}}}$ (16)

We compare the within-group mpRMSD values for structures generated on the WT
epitope to the mpRMSD values between structures generated for the WT epitope
and for epitopes perturbed using three methods: permutation (swapping with
another random epitope), sequence scrambling, and translation. For
permutation, to ensure that global orientation does not contribute to
differences in mpRMSD, we perform a rotational alignment of the principal
components of the two epitopes after centering both epitopes at the origin.
Specifically, for a WT epitope with atomic coordinates ${\bf
Z}_{i}\in\mathbb{R}^{N\times 3}$ and a permuted epitope with coordinates ${\bf
Z}_{j}\in\mathbb{R}^{M\times 3}$, we obtain the aligned permuted epitope
coordinates $\hat{\bf Z}_{j}$ via singular value decomposition of the two
matrices:

$\displaystyle{\bf USV}^{\top}$ $\displaystyle={\bf Z}_{i}$ (17)
$\displaystyle{\bf U^{\prime}S^{\prime}(V^{\prime})}^{\top}$
$\displaystyle={\bf Z}_{j}$ $\displaystyle\hat{{\bf Z}}_{j}$
$\displaystyle={\bf Z}_{j}{\bf V^{\prime}V}^{\top}$

Epitope scrambling was performed by shuffling the residue sequence labels of
each epitope, setting C$\beta$ coordinates of residues whose residue identity
becomes glycine to their C$\alpha$ coordinates, preventing information leakage
(glycines are represented the same way during training). C$\beta$ coordinates
for glycines that convert into non-glycine residues after scrambling are
obtained by uniformly sampling a C$\beta$-C$\alpha$ bond from another epitope
residue and imputing the location of the C$\beta$ atom in the previously-
glycine residue using that bond vector. Epitope translation is performed by
translating all epitope atoms 20 Å along the axis connecting the epitope
centroid to the CDR centroid. Translation is performed along this axis in the
direction opposite to the CDR centroid’s position relative to the epitope
centroid.

##  Appendix B Examples of Structures

Figure 4: Generated loops for 3ULU using a generative model of frames (both
rotations and positions). Loops seem to conform to the structure of the
epitope and adopt a wide variety of shapes, both in the ends of the loop and
in the parts in contact with the epitope. Figure 5: Generated loops for 3ULU
using a generative model of C$\alpha$ coordinates only. Loops are smaller and
more similar than for the frames.

##  Appendix C Evaluating Design Choices For Diffusion Models

Figure 6: The cumulative variance $\mathrm{Var}[x^{(t)}]$ of the different
noising schedules as a function of timestep $t\in[0,1]$. We find that using
the slower schedules (quadratic and sigmoid) for noising the translations
improves performance.