# Insights Into the Inner Workings of Transformer Models for
Protein Function Prediction

Markus Wenzel1 Erik Grüner1 Nils Strodthoff2

(1 Department of Artificial Intelligence, Fraunhofer Heinrich-Hertz-Institute,
Berlin, Germany
2 School of Medicine and Health Sciences, Carl von Ossietzky University of
Oldenburg, Oldenburg, Germany
<EMAIL_ADDRESS><EMAIL_ADDRESS>)

###### Abstract

Motivation: We explored how explainable AI (XAI) can help to shed light into
the inner workings of neural networks for protein function prediction, by
extending the widely used XAI method of integrated gradients such that latent
representations inside of transformer models, which were finetuned to Gene
Ontology term and Enzyme Commission number prediction, can be inspected too.
Results: The approach enabled us to identify amino acids in the sequences that
the transformers pay particular attention to, and to show that these relevant
sequence parts reflect expectations from biology and chemistry, both in the
embedding layer and inside of the model, where we identified transformer heads
with a statistically significant correspondence of attribution maps with
ground truth sequence annotations (e.g., transmembrane regions, active sites)
across many proteins. Availability and Implementation: Source code can be
accessed at this https URL.

##  1 Introduction

###  1.1 Protein function prediction

Proteins – constituents of life. Proteins are versatile molecular machines,
performing various tasks in basically all cells of every organism, and are
modularly constructed from chains of amino acids. Inferring the function of a
given protein merely from its amino acid sequence is a particularly
interesting problem in bioinformatics research.

Function prediction can help to rapidly provide valuable pointers in the face
of so far unfamiliar proteins of understudied species, such as of emerging
pathogens. Moreover, it makes the analysis of large, unlabeled protein data
sets possible, which becomes more and more relevant against the backdrop of
the massive and evermore growing databases of unlabeled nucleic acid
sequences, which again can be translated into amino acid sequences. Next-
generation DNA sequencers can read the nucleic acid sequences present in a
sample or specimen at decreasing costs [Mardis, 2017; Shendure et al., 2017],
much faster than experimenters can determine the function of the genes and
corresponding proteins. Therefore, databases with genes and corresponding
amino acid sequences grow much more rapidly than those of respective
experimental gene and protein labels or annotations. Besides, gaining
knowledge about the mapping between amino acid sequence and protein function
can help to engineer proteins for dedicated purposes too [Alley et al., 2019;
Yang et al., 2019; Hie and Yang, 2022; Ferruz et al., 2022; Madani et al.,
2023].

Machine learning approaches to protein function prediction can include
inferring enzymatic function [Dalkiran et al., 2018; Li et al., 2018; Zou et
al., 2019; Yu et al., 2023], Gene Ontology (GO) terms [Kulmanov et al., 2017;
Kulmanov and Hoehndorf, 2019, 2022; You et al., 2018b, a, 2021; Strodthoff et
al., 2020; Littmann et al., 2021], protein-protein/-drug interaction, remote
homology, stability, sub-cellular location, and other properties [Bepler and
Berger, 2021; Rao et al., 2019]. For structure prediction, the objective is to
infer how the amino acid sequence folds into the secondary [Zhang et al.,
2018; Rives et al., 2021] and tertiary protein structure [Torrisi et al.,
2020; AlQuraishi, 2021; Jumper et al., 2021; Weissenow et al., 2022]. Several
of the prediction tasks can be approached as well by transferring labels from
similar sequences obtained via multiple sequence alignment (MSA) [Buchfink et
al., 2014; Gong et al., 2016]. Protein prediction models are compared by the
scientific community in systematic performance benchmarks, e.g., for function
annotation [CAFA, Radivojac et al., 2013; Jiang et al., 2016; Zhou et al.,
2019], for structure prediction [CASP, Kryshtafovych et al., 2019, 2021], or
for several semi-supervised tasks [Rao et al., 2019; Fenoy et al., 2022].
Machine learning methods are continuing to win ground in comparison to MSA-
techniques with respect to performance, have a short inference time, and can
process sequences from the so-called ‘dark proteome’ too, where alignments are
not possible [Perdigão et al., 2015; Rao et al., 2019; Lin et al., 2023].

###  1.2 Protein language modeling and transfer learning

Relations to NLP. Amino acid sequences share some similarities with the
sequences of letters and words occurring in written language, in particular
with respect to the complex interrelationships between distant elements, which
are arranged in one-dimensional chains. Thus, recent progress in research on
natural language processing (NLP) employing language modeling in a transfer
learning scheme [Howard and Ruder, 2018] has driven forward protein function
prediction too [e.g., Strodthoff et al., 2020].

Self-supervised pretraining. Typically, a language model is first pretrained
on large numbers of unlabeled sequences in an unsupervised fashion, e.g., by
learning to predict masked tokens (cloze task) or the respective next token in
the sequences (which is why this unsupervised approach is also dubbed self-
supervised learning). In this way, the model learns useful representations of
the sequence statistics (i.e. language). These statistics possibly arise
because the amino acid chains need to be stable under physiological conditions
and are subject to evolutionary pressure. The learned representations can be
transferred to separate downstream tasks, where the pretrained model can be
further finetuned in a supervised fashion on labeled data, which are usually
available in smaller amounts, considering that sequence labelling by
experimenters is costly and lengthy.

Model architectures. Transformer models [Vaswani et al., 2017] making use of
the attention mechanism [Niu et al., 2021], such as bidirectional encoder
representations from transformers [BERT, Devlin et al., 2018] are currently
prevailing architectures in NLP. Transformers have been recently applied to
the study of amino acid sequences too, pushing the state of the art in the
field of proteomics as well [Rao et al., 2019, 2021; Nambiar et al., 2020;
Bepler and Berger, 2021; Rives et al., 2021; Littmann et al., 2021; Elnaggar
et al., 2022; Unsal et al., 2022; Brandes et al., 2022; Fenoy et al., 2022;
Olenyi et al., 2023]. Recurrent neural networks (RNNs) using long short term
memory (LSTM) cells are another model architecture that is particularly suited
to process sequential data. RNNs have been successfully employed to protein
[Strodthoff et al., 2020] and peptide [Vielhaben et al., 2020] property
prediction as well, within the scheme of language modeling combined with
transfer learning, as sketched out above.

###  1.3 Explainable machine learning

Need for explainability. Transformers and other modern deep learning models
are notorious for having often millions and sometimes billions of trainable
parameters, and it can be very difficult to interpret the decision making
logic or strategy of such complex models. The research field of explainable
machine learning [Lundberg and Lee, 2017; Montavon et al., 2018; Tjoa and
Guan, 2020; Samek et al., 2021; Arrieta et al., 2020; Covert et al., 2021]
aims at developing methods that enable humans to better interpret – or to a
limited degree: understand – such ‘opaque’, complex models. In certain cases,
it was demonstrated that the methods can even help to uncover flaws and
unintended biases of the models, such as being mislead by spurious
correlations in the data [Lapuschkin et al., 2019].

Attribution methods, such as integrated gradients (IG) [Sundararajan et al.,
2017], layerwise-relevance propagation [Bach et al., 2015; Binder et al.,
2016] or gradient-weighted class activation mapping [Selvaraju et al., 2017],
make it possible to identify those features in the input space that the model
apparently focuses on, because these features turn out to be particular
relevant for the final classification decision of the model. Further examples
of model explainability methods include probing classifiers [Belinkov, 2022],
testing with concept activation vectors [Kim et al., 2018], and studying the
attention mechanism [Serrano and Smith, 2019; Bai et al., 2021; Niu et al.,
2021; Jain and Wallace, 2019]. Explainability methods have been employed in
NLP too [Arras et al., 2019; Pascual et al., 2021; Manning et al., 2020;
Chefer et al., 2021]. Moreover, researchers have started to explore using
explainability methods in the area of protein function prediction [Upmeier zu
Belzen et al., 2019; Vig et al., 2021; Taujale et al., 2021; Vu et al., 2023;
Zhou et al., 2023; Hou et al., 2023].

###  1.4 Contributions of the paper

Figure 1: Illustration of the experimental design. Top: From the amino acid
sequence, the finetuned transformer model infers the applicable Gene Ontology
(GO) terms (represented as multi-label class membership vector). (The depicted
exemplary ‘catalase-3’ should be labeled with the GO terms ‘catalase activity’
as ‘molecular function’, ‘response to hydrogen peroxide’ as ‘biological
process’, ‘cytoplasm’ as ‘cellular component’ etc.; about 5K of about 45K GO
terms were considered.) Center: Relevance indicative for a selected GO term
was attributed to the amino acids per protein and correlated with
corresponding annotations per amino acid. This correlation between relevance
attributions and annotations was then statistically assessed across the test
data set proteins. The analysis was conducted for the embedding layer and
‘inside’ of the model, for each head in each layer, and was repeated for
different GO terms (see Section 2.1). Bottom: Specific amino acids of a
protein are annotated in sequence databases like UniProt, because they serve
as binding or active sites or are located in the cell membrane etc. Active
sites can, e.g., be found at the histidine (‘H’ at position 65) and asparagine
(‘N’ at position 138) of ‘catalase-3’ (structure created by AlphaFold [Jumper
et al., 2021] under the CC-BY 4.0 licence).

Goal of the study. Building upon this previous research on the interpretation
of protein classification models, we aimed at exploring how explainability
methods can further help to gain insights into the inner workings of the now
often huge neural networks, and proceeded as follows.

Specific contributions. First, we finetuned pretrained transformers on
selected prediction tasks and could push or reach the state-of-the-art. Then,
we quantified the relevance of each amino acid of a protein for the function
prediction model. Subsequently, we investigated if these relevant sequence
regions match expectations informed by knowledge from biology or chemistry, by
correlating the relevance attributions with annotations from sequence
databases (see Figure 1). For instance, we addressed the question if a
classification model that is able to infer if a protein is situated in the
cell membrane does indeed focus systematically on transmembrane regions or
not. We conducted this analysis on the embedding level and ‘inside’ of the
model with a novel adaptation of IG. In this way, we identified transformer
heads with a statistically significant correspondence of the attribution maps
with ground truth annotations, across many proteins and thus going beyond
anecdotes of few selected cases.

##  2 System and methods

###  2.1 Revealing insights into function prediction models

The prediction tasks of inferring Gene Ontology (GO) terms and Enzyme
Commission (EC) numbers, that the proteins are labeled with, from their amino
acid sequence are detailed in Appendix B. This supplementary material also
explains the finetuning of the transformers ‘ProtBert-BFD’ and ‘ProtT5-XL-
UniRef50’ [Elnaggar et al., 2022] on the GO and EC tasks, and contains
statements about data availability and composition.

Overall approach. We investigated whether specific positions or areas on the
amino acid sequence that had been annotated in sequence data bases are
particularly relevant for the classification decision of the model (see Figure
1). Annotations included UniProtKB/Swiss-Prot ‘active’ and ‘binding sites’,
‘transmembrane regions’, ‘short sequence motifs’, and PROSITE patterns related
to a GO term and its children terms in the ontology. Definitions of the
aforementioned UniProt annotations (per amino acid) and matching GO terms
(class labels of proteins) are compiled in Table A.1111Tables/figures with
prefix letters are shown in the supplement. material.. First, we attributed
relevance indicative for a given class (either a selected GO term or EC
number) to each amino acid of a protein. Then, we correlated the relevance
heat map obtained for the amino acid chain of a protein with corresponding
binary sequence annotations. To study the information representation within
the model, the explainability analysis was conducted at the embedding layer
and repeated ‘inside’ of the model, separately for its different heads and
layers, using a novel method building upon IG, described below in Section 3.

Experimental setup. For the experimental evaluation, we focus on the
pretrained ProtBert model that was finetuned either to the multi-label GO-
classification on the GO ‘2016’ data set, or to the multi-class EC number
classification on the ‘EC50 level L1’ data set. This decision is based on the
computational benefit of a substantially smaller memory footprint of BERT in
comparison to T5. We consider the comparatively narrow EC task in addition to
the much more comprehensive GO prediction, because the test split of the EC
data set contains a larger number of samples that are both labeled per protein
and annotated per amino acid, which is beneficial for the conducted
explainability analysis.

##  3 Algorithm

Integrated gradients [Sundararajan et al., 2017] represents a model-agnostic
attribution method, which can be characterized as unique attribution method
satisfying a set of four axioms (Invariance, Sensitivity, Linearity, and
Completeness). In this formalism, the attribution for feature $i$ is defined
via the line integral (along a path, parameterized as $\gamma(t)$ with
$t\in[0,1]$, between some chosen baseline $\gamma(0)=x^{\prime}$ and the
sample to be explained $\gamma(1)=x$),

$\text{IG}^{\gamma}_{i}=\int_{0}^{1}\text{d}\alpha\frac{\partial
F(\gamma(\alpha))}{\partial\gamma_{i}}\frac{\text{d}\gamma_{i}}{\text{d}\alpha}\,,$
(1)

where $F$ is the function we aim to explain. Choosing $\gamma$ as straight
line connecting $x^{\prime}$ and $x$ makes IG the unique method satisfying the
four axioms from above and an additional symmetry axiom. This path is the
typical choice in applications applied directly to the input layer for
computer vision or to the embedding layer for NLP. The approach can be
generalized to arbitrary layers if one replaces $x$ and $x^{\prime}$ by the
hidden feature representation of the network up to this layer (referred to as
‘layer IG’ in the popular ‘Captum’ library [Kokhlikyan et al., 2020]).

Head-specific attribution maps. To obtain attributions for individual heads,
we have to target the output of the multi-head self-attention (MHSA) block of
a particular layer; see Figure 2 for a visualization of the transformer
architecture. Properly separating the attributions of the individual heads
from the attribution contribution obtained from the skip connection
necessitates to target directly the output of the MHSA. Now, one cannot just
simply choose an integration path that connects baseline and sample as encoded
by the MHSA block because the input for the skip connection has to be varied
consistently. To keep an identical path in all cases, we fix the integration
path as a straight line in the embedding layer, which then gets encoded into
a, in general, curvilinear path seen as input for some intermediate layer.
Choosing not a straight path only leads to the violation of the symmetry
axiom, which is not of paramount practical importance in this application; see
[Ward et al., 2020; Kapishnikov et al., 2021] for other applications with IG
applied to general paths. For every sample, this application of IG yields a
relevance map of shape $\text{seq}\times n_{\text{model}}$, where the first
$n_{\text{model}}/n_{\text{heads}}$ entries in the last dimension correspond
to the first head, followed by the second head etc. By summing over
$n_{\text{model}}/n_{\text{heads}}$ entries in the last dimension, we can
reduce the relevance map to a $\text{seq}\times n_{\text{heads}}$ attribution
map, i.e., one relevance sequence per head.

Correlation coefficients and statistical significance. Each sequence of
relevance attributions can then be correlated with sequence annotations to
find out if the model focuses on the annotated amino acids. Coefficients of
point biserial correlation [Kornbrot, 2005], which is equivalent to Pearson
correlation, were calculated between the continuous relevance values and the
corresponding binary annotations per amino acid. This correlation analysis was
conducted separately for each head in each transformer layer. The resulting
correlation coefficients were then assembled into a $n_{\text{layer}}\times
n_{\text{head}}$ matrix per protein, which entered the subsequent statistical
analysis across proteins. Summary statistics over all proteins (which belong
to the respective GO or EC class, and, which are part of the respective test
data set split) were obtained by computing t-tests across the correlation
coefficients. The resulting p-values were corrected for the multiple tests per
condition (16 heads times 30 layers equals 480 hypothesis tests) by
controlling the false discovery rate [Benjamini and Hochberg, 1995].

Summed attribution maps. In parallel to the correlation analysis, we
furthermore sum the aforementioned attribution map along the sequence
dimension, and obtain $n_{\text{heads}}$ entries that specify the relevance
distribution onto the different heads. We can carry out the same procedure for
every transformer layer and combine all results into an
$n_{\text{layer}}\times n_{\text{head}}$ relevance map of summed attributions.
This map makes it possible to identify heads with a positive relevance with
respect to the selected class. One map was obtained per protein. Heads with a
significantly positive relevance were singled out by calculating a summary
statistic across proteins with the Wilcoxon signed-rank test.

Finally, the two parallel analysis tracks were combined by identifying
transformer heads that feature both a significantly positive (A) relevance-
annotation-correlation and (B) relevance (this overlay is displayed in the
figures by masking A with B).

Figure 2: Visualization of the explainability method based on IG that can
attribute relevance to sequence tokens (here: amino acids) separately for each
head and layer of the transformer [adapted from Vaswani et al., 2017].

##  4 Implementation

Appendix C shows implementation details as supplementary material.

##  5 Results and Discussion

###  5.1 Predictive performance

The performance results for the ProtT5 and ProtBert transformers finetuned to
the GO and EC protein function tasks are presented in Appendix D as
supplementary material. In summary, we show that finetuning pretrained large
transformer models leads to competitive results, in particular in the most
relevant comparison in the single-model category, often on par with MSA-
approaches. The approach of finetuning the entire model including the encoder
shows its particular strength in the ‘CAFA3’ benchmark.

###  5.2 Explainability analysis: embedding layer

Corr(relev., annot.) $>0$ | Relevance $>0$ | Overlay
---|---|---
|  |
Figure 3: Inside ProtBert; GO ‘membrane’ (GO:0016020). Left: The relevance
attribution (along the sequences) indicative for the GO term ‘membrane’ was
correlated with UniProt annotations as ‘transmembrane regions’, for each
transformer head and layer. Biserial correlation coefficients (r), obtained
for each attribution-annotation-pair, were aggregated in population statistics
with t-tests. The resulting p-values of the t-tests were adjusted with the
Benjamini/Hochberg method for the multiple hypothesis tests conducted in order
to limit the false discovery rate. A significance threshold was applied
(family-wise error rate of 0.05). The negative logarithm of the corrected and
thresholded p-values is displayed. All colored pixels indicate statistically
significant results. Center: ProtBert heads with a sig. positive relevance
(sum along the sequence; indicative for the GO term ‘membrane’) were singled
out with the Wilcoxon signed-rank test. The matrix plots show the negative
logarithm of the resulting p-values (adjusted with Benjamini/Hochberg and a
threshold). Right: ProtBert heads with a sig. positive attribution-annotation-
correlation (p-values from t-tests plotted) that are also characterized by a
sig. positive relevance (the latter overlaid as mask). Only results for
UniProt ‘transmembrane regions’ are shown, omitting the results for
‘active/binding sites’, ‘motifs’, and PROSITE patterns, which did not feature
heads with both a sig. positive relevance and attribution-annotation-
correlation.

Research question. Starting with embedding layer attribution maps, as the most
widely considered type of attribution, we investigate whether there are
significant correlations between attribution maps and sequence annotations
from external sources (see Section 2.1). We aim to answer this question in a
statistical fashion going beyond anecdotal evidence based on single examples,
which can sometimes be encountered in the literature.

GO prediction: GO ‘membrane’ attributions correlate in particular with UniProt
‘transmembrane regions’. Figure 4 shows the results of the explainability
analysis for the embedding layer of ProtBert finetuned to GO classification.
The relevance of each amino acid indicative for selected GO terms was computed
with IG, and then correlated with UniProt and PROSITE sequence annotations.
Subsequently, it was tested whether the correlation coefficients across all
annotated proteins from the test set were significantly positive (see Section
2.1). A significant correlation was observed when relevance attributions
indicative for the GO label ‘membrane’ were correlated with UniProt
‘transmembrane regions’ ($p\ll 0.05$) and with ‘active sites’ ($p<0.05$).
Correlation was not observed in the GO ‘catalytic activity’ and ‘binding’
cases.

Membrane | Catalytic activity | Binding
---|---|---
|  |
Figure 4: Attribution maps for the embedding layer of ProtBert finetuned to GO
term classification were correlated with sequence annotations. Left: Relevance
attributions indicative for the GO label ‘membrane’ correlate significantly
($*$) with UniProt annotations as ‘transmembrane regions’ ($p\ll 0.05$) and
‘active sites’ ($p<0.05$). Center and right: Attribution-annotation-
correlation was not observed for GO ‘catalytic activity’ and ‘binding’.
Numbers of test split samples both labeled with the GO term and annotated per
amino acid are listed below the x-axis.

EC prediction: attributions correlate significantly with several types of
sequence annotations. Figure 5 shows the results of the explainability
analysis for the embedding layer of ProtBert finetuned to EC number
classification (‘EC50 level L1’ data set; i.e., the differentiation between
the six main enzyme classes). Relevance per amino acid for each of the six EC
classes was correlated with the UniProt annotations as ‘active sites’,
‘binding sites’, ‘transmembrane regions’, and ‘short sequence motifs’. It can
be observed that the relevance attributions correlated significantly
($p<0.05$) with ‘active site’ annotations for all six EC classes, with
‘binding site’ annotations for four EC classes, and with ‘transmembrane
regions’ and ‘short sequence motifs’ for two EC classes.

Figure 5: Attribution maps calculated for the embedding layer of ProtBert
finetuned to ‘EC50 level L1’ classification were correlated with UniProt
sequence annotations. Left: Relevance attributions correlated significantly
($p<0.05$) with ‘active sites’ for all six EC classes, with ‘binding sites’
for four EC classes, and with ‘transmembrane regions’ and ‘short sequence
motifs’ for two EC classes each. Right: Numbers of annotated samples in the
test split per annotation type and EC class.

Discussion. Attribution maps obtained for the embedding layer correlated with
UniProt annotations on the amino acid level, in particular, in the EC case,
but also for the GO term ‘membrane’. To summarize, across two tasks, we
provide first quantitative evidence for the meaningfulness and specificity of
attribution maps beyond anecdotal evidence. Note that the EC case has the
benefit of often several hundred annotated samples contained in the test split
(except for ‘transmembrane regions’ and ‘motifs’; see right panel of Figure
5). In comparison, the GO case provides fewer samples in the test split of the
data set that were also annotated on the amino acid level (see numbers in
brackets below the x-axis in Figure 4).

###  5.3 Explainability analysis: Peeking inside the transformer

Research question. Given the encouraging results presented in Section 5.2, we
aim to go one step further and try to answer the more specific question if
there are specialized heads inside of the model architecture for specific
prediction tasks, using our IG variant that calculates relevance on the amino
acid level per transformer head and layer (see Section 3).

GO-prediction: membrane. Figure 3 shows the results of the explainability
analysis inspecting the latent representations inside of the ProtBert model
focusing on the selected class of the GO term ‘membrane’ (GO:0016020).
Relevance attributions indicative for GO ‘membrane’ per amino acid were
correlated with the UniProt annotations as ‘transmembrane regions’ separately
for each transformer head and layer (matrix plot pixels in Figure 3). In
parallel, ProtBert heads were singled out with a significantly positive
relevance (sum along the sequence) indicative for ‘membrane’ (see also Section
2.1 and Section 3). Both parallel analysis streams were combined by
identifying ProtBert heads with both a significantly positive attribution-
annotation-correlation and relevance. Several ProtBert heads in different
layers feature a significantly positive correlation of relevance attributions
per amino acid with the UniProt annotations as ‘transmembrane regions’, going
along with a significantly positive relevance for the GO class ‘membrane’. In
contrast, correlation of relevance attributions with UniProt ‘active’ or
‘binding sites’ or ‘motifs’ or PROSITE patterns accompanied by a positive
relevance was not observed (hence these cases were not included in Figure 3).

GO prediction: catalytic activity. Figure D.1 (in Appendix D as supplementary
material) shows the results of the explainability analysis for the case where
the GO term ‘catalytic activity’ was selected (GO:0003824). Different ProtBert
heads stand out characterized by a positive relevance accompanied by a
positive correlation of attributions with PROSITE patterns and with UniProt
‘active sites’ and ‘transmembrane regions’ (but neither with ‘binding sites’
nor ‘motifs’).

GO-prediction: binding. Figure D.2 (in Appendix D) repeats the explainability
analysis inside ProtBert for the GO term ‘binding’ (GO:0005488). For several
transformer heads and layers, a positive relevance went along with a
correlation of relevance attributions with corresponding PROSITE patterns, and
with UniProt ‘transmembrane regions’ (but neither with UniProt ‘active’ nor
‘binding sites’ nor ‘motifs’).

EC-prediction. Subsequently, we conducted the explainability analysis for the
case where ProtBert had been finetuned to EC number classification on EC50
level L1. Here, the model had learned to differentiate between the six main
enzyme classes. Figure D.3 (in Appendix D) identifies ProtBert heads
characterized both by a positive relevance (sum along the sequence) with
respect to the EC class, and by a positive attribution-annotation-correlation
(on the amino acid level). The analysis was conducted separately for UniProt
annotations as ‘active’/‘binding sites’, ‘transmembrane regions’, and
‘motifs’. (The absence of identified heads for EC4, EC5, and EC6 in the
‘transmembrane regions’ rows and for EC1 and EC5 in the ‘motif’ rows of Figure
D.3 goes along with the availability of relatively few ‘transmembrane’ and
‘motif’ annotations for these EC classes; see histogram in Figure 5.)

Discussion. In summary, we propose a constructive method suited to identify
heads inside of the transformer architecture that are specialized for specific
protein function or property prediction tasks. The proposed method comprises a
novel adaptation of the XAI method of IG combined with a subsequent
statistical analysis. We first attributed relevance to the single amino acids
per protein (per GO term or EC class), separately for each transformer head
and layer. Then, we inspected the correlation between relevance attributions
and annotations, in a statistical analysis across the annotated proteins from
the test split of the respective data set. Apparently, different transformer
heads are sensitive to different annotated and thus biologically and,
respectively, chemically ‘meaningful’ sites, regions or patterns on the amino
acid sequence. The methods is complementary to probing approaches [Vig et al.,
2021].

###  5.4 Uncovering collective dynamics

Finally, we studied collective dynamics potentially emerging among the
transformer heads (ProtBert, EC50, level L1) by a visualization of the
originally high-dimensional, summed attribution maps in two dimensions, taking
their similarities into account. For this purpose, the attribution maps that
were summed along the amino acid sequence and represented as
$n_{\text{layer}}\times n_{\text{head}}$ matrices (see Section 3) were
flattened, resulting in one vector per protein. The dimensionality of these
vectors was then reduced with principal component analysis to 50 dimensions,
and subsequently to two dimensions with t-distributed stochastic neighbor
embedding [t-SNE; van der Maaten and Hinton, 2008], using the default t-SNE
parameters. The resulting 2D points were visualized as scatter plot and
colored according to the corresponding six main enzyme classes (Figure 6).

The points form distinctive clusters matching the EC labels. Apparently, a
structure emerges in the attribution maps, that seems to indicate class-
specific collective dynamics among several ProtBert heads. It is important to
stress that the attribution map underlying the clustering no longer contains
any reference to specific positions in the sequence but relies on the
relevance distribution on the different heads through all layers of the model.
The emergence of class-specific structures therefore indicates that there are
specific combinations of heads that are relevant for a specific classification
decision.

Figure 6: PCA and t-SNE visualization of summed attribution maps (ProtBert,
EC50, L1).

##  6 Conclusion

This work provides additional evidence for the effectiveness of the currently
predominant paradigm in deep-learning-based protein analysis through the
finetuning of large language models. For different protein function prediction
tasks, this approach leads to best-performing models according to single-model
performance. The performance level is in many cases on par with MSA-
approaches. The proposed models can even be effectively combined with the
latter through the formation of ensembles.

Considering the ever increasing model complexity, XAI has started to gain
traction in the field of protein analysis too [Upmeier zu Belzen et al., 2019;
Vig et al., 2021; Taujale et al., 2021; Vu et al., 2023; Zhou et al., 2023;
Hou et al., 2023], but quantitative evidence for its applicability beyond
single examples was lacking up to now. We provide statistical evidence for the
alignment of attribution maps with corresponding sequence annotations, both on
the embedding level as well as for specific heads inside of the model
architecture, which led to the identification of specialized heads for
specific protein function prediction tasks. Emerging class-specific structures
suggest that these specialized transformer heads act jointly to decide
together in specific combinations.

XAI promises to tap into the presumably substantial knowledge contained in
large models pretrained on massive data sets of amino and/or nucleic acid
sequences [Steinegger et al., 2019; Ji et al., 2021]. Therefore, we expect
that XAI will play an increasingly important role in the future of
bioinformatics research. We see potential applications of XAI methods for
model validation and for scientific discovery (e.g., of novel discriminative
sequence patterns or motifs that have not been identified by experiments or
MSA so far). Identifying specialized heads might also help to prune or distill
overly large models, making them smaller and more efficient.

## Funding

This work was supported by the Bundesministerium für Bildung und Forschung
through the BIFOLD - Berlin Institute for the Foundations of Learning and Data
[grant numbers 01IS18025A, 01IS18037A].

## References

  * Alley et al. [2019] Alley, E. C. et al. (2019).  Unified rational protein engineering with sequence-based deep representation learning.  Nat. Methods, 16(12), 1315–22.
  * AlQuraishi [2021] AlQuraishi, M. (2021).  Machine learning in protein structure prediction.  Curr. Opin. Chem. Biol., 65, 1–8.
  * Arras et al. [2019] Arras, L. et al. (2019).  Evaluating Recurrent Neural Network Explanations.  In Proc. 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP. Assoc. Comput. Linguist.
  * Arrieta et al. [2020] Arrieta, A. B. et al. (2020).  Explainable Artificial Intelligence (XAI): Concepts, taxonomies, opportunities and challenges toward responsible AI.  Information fusion, 58, 82–115.
  * Ashburner et al. [2000] Ashburner, M. et al. (2000).  Gene Ontology: tool for the unification of biology.  Nat. genet., 25(1), 25–29.
  * Bach et al. [2015] Bach, S. et al. (2015).  On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation.  PLoS ONE, 10(7).
  * Bai et al. [2021] Bai, B. et al. (2021).  Why Attentions May Not Be Interpretable? In Proc. 27th ACM SIGKDD, KDD ’21, page 25–34, NY, USA. Assoc. Comput. Mach.
  * Belinkov [2022] Belinkov, Y. (2022).  Probing classifiers: Promises, shortcomings, and advances.  Comput. Linguist., 48(1), 207–219.
  * Benjamini and Hochberg [1995] Benjamini, Y. and Hochberg, Y. (1995).  Controlling the false discovery rate: a practical and powerful approach to multiple testing.  J. R. Stat. Soc.: series B (Methodol.), 57(1), 289–300.
  * Bepler and Berger [2021] Bepler, T. and Berger, B. (2021).  Learning the protein language: Evolution, structure, and function.  Cell Systems, 12(6), 654–669.e3.
  * Bernhofer et al. [2021] Bernhofer, M. et al. (2021).  PredictProtein - Predicting Protein Structure and Function for 29 Years.  Nucleic Acids Res., 49(W1), W535–W540.
  * Binder et al. [2016] Binder, A. et al. (2016).  Layer-wise relevance propagation for deep neural network architectures.  In Information Science and Applications (ICISA) 2016, pages 913–922. Springer.
  * Brandes et al. [2022] Brandes, N. et al. (2022).  ProteinBERT: A universal deep-learning model of protein sequence and function.  Bioinformatics, 38(8), 2102–10.
  * Buchfink et al. [2014] Buchfink, B. et al. (2014).  Fast and sensitive protein alignment using DIAMOND.  Nat. Methods, 12(1), 59–60.
  * Chefer et al. [2021] Chefer, H. et al. (2021).  Transformer interpretability beyond attention visualization.  In 2021 IEEE/CVF CVPR, pages 782–791.
  * Clark and Radivojac [2013] Clark, W. T. and Radivojac, P. (2013).  Information-theoretic evaluation of predicted ontological annotations.  Bioinformatics, 29(13), i53–i61.
  * Consortium [2020a] Consortium, G. O. (2020a).  The Gene Ontology resource: enriching a GOld mine.  Nucleic Acids Res., 49(D1), D325–D334.
  * Consortium [2020b] Consortium, U. (2020b).  UniProt: the universal protein knowledgebase in 2021.  Nucleic Acids Res., 49(D1), D480–D489.
  * Covert et al. [2021] Covert, I. et al. (2021).  Explaining by Removing: A Unified Framework for Model Explanation.  J. Mach. Learn. Res., 22(209), 1–90.
  * Dalkiran et al. [2018] Dalkiran, A. et al. (2018).  ECPred: a tool for the prediction of the enzymatic functions of protein sequences based on the EC nomenclature.  BMC Bioinform., 19(1), 334.
  * Devlin et al. [2018] Devlin, J. et al. (2018).  BERT: Pre-training of deep bidirectional transformers for language understanding.  arXiv:1810.04805.
  * Elnaggar et al. [2022] Elnaggar, A. et al. (2022).  ProtTrans: Toward Understanding the Language of Life Through Self-Supervised Learning.  IEEE Trans. Pattern Anal. Mach. Intell., 44(10), 7112–27.
  * Fenoy et al. [2022] Fenoy, E. et al. (2022).  Transfer learning in proteins: evaluating novel protein learned representations for bioinformatics tasks.  Briefings in Bioinformatics, 23(4), bbac232.
  * Ferruz et al. [2022] Ferruz, N. et al. (2022).  ProtGPT2 is a deep unsupervised language model for protein design.  Nat. Commun., 13(1), 4348.
  * Gong et al. [2016] Gong, Q. et al. (2016).  GoFDR: a sequence alignment based method for predicting protein functions.  Methods, 93, 3–14.
  * Hie and Yang [2022] Hie, B. L. and Yang, K. K. (2022).  Adaptive machine learning for protein engineering.  Curr. Opin. Struct. Biol., 72, 145–152.
  * Hou et al. [2023] Hou, Z. et al. (2023).  Learning the protein language of proteome-wide protein-protein binding sites via explainable ensemble deep learning.  Commun. Biol., 6(1), 73.
  * Howard and Ruder [2018] Howard, J. and Ruder, S. (2018).  Universal Language Model Fine-tuning for Text Classification.  In Proc. 56th Annu. Meet. Assoc. Comput. Linguist. (Vol. 1: Long Papers), pages 328–339, Melbourne, Australia. Assoc. Comput. Linguist.
  * Jain and Wallace [2019] Jain, S. and Wallace, B. C. (2019).  Attention is not explanation.  In Proc. 2019 Conf. N. Amer. Chapter Assoc. Comput. Linguist.: Hum. Lang. Tech., Vol. 1”, pages 3543–56, Minneapolis, MN. Assoc. Comput. Linguist.
  * Ji et al. [2021] Ji, Y. et al. (2021).  DNABERT: pre-trained Bidirectional Encoder Representations from Transformers model for DNA-language in genome.  Bioinformatics, 37(15), 2112–20.
  * Jiang et al. [2016] Jiang, Y. et al. (2016).  An expanded eval. of protein function prediction methods shows an improvement in accuracy.  Genome biol., 17(1), 1–19.
  * Jumper et al. [2021] Jumper, J. et al. (2021).  Highly accurate protein structure prediction with AlphaFold.  Nature, 596(7873), 583–589.
  * Kapishnikov et al. [2021] Kapishnikov, A. et al. (2021).  Guided integrated gradients: An adaptive path method for removing noise.  In 2021 CVPR, pages 5048–56.
  * Kim et al. [2018] Kim, B. et al. (2018).  Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (TCAV).  In ICML, pages 2668–77. PMLR.
  * Kingma and Ba [2015] Kingma, D. P. and Ba, J. (2015).  Adam: A method for stochastic optimization.  In 3rd ICLR, San Diego.
  * Kokhlikyan et al. [2020] Kokhlikyan, N. et al. (2020).  Captum: A unified and generic model interpretability library for PyTorch.  arXiv:2009.07896.
  * Kornbrot [2005] Kornbrot, D. (2005).  Point Biserial Correlation.  John Wiley & Sons.
  * Kryshtafovych et al. [2019] Kryshtafovych, A. et al. (2019).  Critical assessment of methods of protein structure prediction (CASP) – Round XIII.  Proteins: Structure, Function, and Bioinformatics, 87(12), 1011–20.
  * Kryshtafovych et al. [2021] Kryshtafovych, A. et al. (2021).  Critical assessment of methods of protein structure prediction (CASP) – Round XIV.  Proteins: Structure, Function, and Bioinformatics, 89(12), 1607–17.
  * Kulmanov and Hoehndorf [2019] Kulmanov, M. and Hoehndorf, R. (2019).  DeepGOPlus: improved protein function prediction from sequence.  Bioinformatics.
  * Kulmanov and Hoehndorf [2022] Kulmanov, M. and Hoehndorf, R. (2022).  DeepGOZero: improving protein function prediction from sequence and zero-shot learning based on ontology axioms.  Bioinformatics, 38(Supp. 1), i238–i245.
  * Kulmanov et al. [2017] Kulmanov, M. et al. (2017).  DeepGO: predicting protein functions from sequence and interactions using a deep ontology-aware classifier.  Bioinformatics, 34(4), 660–668.
  * Lapuschkin et al. [2019] Lapuschkin, S. et al. (2019).  Unmasking Clever Hans predictors and assessing what machines really learn.  Nat. Commun., 10, 1096.
  * Li et al. [2018] Li, Y. et al. (2018).  DEEPre: sequence-based enzyme EC number prediction by deep learning.  Bioinformatics, 34(5), 760–769.
  * Lin et al. [2023] Lin, Z. et al. (2023).  Evolutionary-scale prediction of atomic-level protein structure with a language model.  Science, 379(6637), 1123–30.
  * Littmann et al. [2021] Littmann, M. et al. (2021).  Embeddings from deep learning transfer GO annotations beyond homology.  Sci. rep., 11(1), 1–14.
  * Lundberg and Lee [2017] Lundberg, S. M. and Lee, S.-I. (2017).  A Unified Approach to Interpreting Model Predictions.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Adv. NeurIPS, volume 30. Curran Assoc., Inc.
  * Madani et al. [2023] Madani, A. et al. (2023).  Large language models generate functional protein sequences across diverse families.  Nat. Biotechnol., pages 1–8.
  * Manning et al. [2020] Manning, C. D. et al. (2020).  Emergent linguistic structure in artificial neural networks trained by self-supervision.  PNAS, 117(48), 30046–54.
  * Mardis [2017] Mardis, E. R. (2017).  DNA sequencing technologies: 2006–2016.  Nat. Protoc., 12(2), 213–218.
  * McDonald et al. [2008] McDonald, A. G. et al. (2008).  ExplorEnz: the primary source of the IUBMB enzyme list.  Nucleic Acids Res., 37(Suppl. 1), D593–D597.
  * Montavon et al. [2018] Montavon, G. et al. (2018).  Methods for interpreting and understanding deep neural networks.  Digit. Signal Process, 73, 1–15.
  * Nambiar et al. [2020] Nambiar, A. et al. (2020).  Transforming the Language of Life: Transformer Neural Networks for Protein Prediction Tasks.  In Proc. 11th ACM Int. Conf. Bioinformatics, Comput. Biology and Health Informatics, BCB ’20, New York, NY, USA. Assoc. Comput. Mach.
  * Niu et al. [2021] Niu, Z. et al. (2021).  A review on the attention mechanism of deep learning.  Neurocomputing, 452, 48–62.
  * Olenyi et al. [2023] Olenyi, T. et al. (2023).  LambdaPP: Fast and accessible protein-specific phenotype predictions.  Protein Science, 32(1), e4524.
  * Pascual et al. [2021] Pascual, D. et al. (2021).  Telling BERT’s Full Story: from Local Attention to Global Aggregation.  In Proc. 16th Conf. Eur. Ch. Assoc. Comput. Linguist.: Main Vol., pages 105–124, Online. Assoc. Comput. Linguist.
  * Paszke et al. [2019] Paszke, A. et al. (2019).  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Adv. NeurIPS 32, pages 8024–35. Curran Assoc., Inc.
  * Perdigão et al. [2015] Perdigão, N. et al. (2015).  Unexpected features of the dark proteome.  PNAS, 112(52), 15898–903.
  * Radivojac et al. [2013] Radivojac, P. et al. (2013).  A large-scale evaluation of computational protein function prediction.  Nat. Methods, 10(3), 221–227.
  * Raffel et al. [2020] Raffel, C. et al. (2020).  Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.  J. Mach. Learn. Res., 21(140), 1–67.
  * Rao et al. [2019] Rao, R. et al. (2019).  Evaluating Protein Transfer Learning with TAPE.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Adv. Neural Inf. Process. Syst., volume 32.
  * Rao et al. [2021] Rao, R. M. et al. (2021).  MSA Transformer.  In M. Meila and T. Zhang, editors, Proc. 38th ICML, volume 139 of Proc. Mach. Learn. Res., pages 8844–56. PMLR.
  * Reimers and Gurevych [2019] Reimers, N. and Gurevych, I. (2019).  Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.  In Proc. 2019 EMNLP. Assoc. Comput. Linguist.
  * Rives et al. [2021] Rives, A. et al. (2021).  Biol. structure and function emerge from scaling unsupervised learning to 250 million protein sequences.  PNAS, 118(15).
  * Samek et al. [2021] Samek, W. et al. (2021).  Explaining Deep Neural Networks and Beyond: A Review of Methods and Applications.  Proc. IEEE, 109(3), 247–278.
  * Seabold and Perktold [2010] Seabold, S. and Perktold, J. (2010).  Statsmodels: Econometric and statistical modeling with Python.  In Proc. 9th Python in Science Conf., volume 57, pages 10–25080. Austin, Texas.
  * Selvaraju et al. [2017] Selvaraju, R. R. et al. (2017).  Grad-CAM: Visual Explanations From Deep Networks via Gradient-Based Localization.  In 2017 IEEE ICCV, pages 618–626.
  * Serrano and Smith [2019] Serrano, S. and Smith, N. A. (2019).  Is Attention Interpretable? In Proc. 57th Annu. Meet. Assoc. Comput. Linguist., pages 2931–51, Florence, Italy. Assoc. Comput. Linguist.
  * Shendure et al. [2017] Shendure, J. et al. (2017).  DNA sequencing at 40: past, present and future.  Nature, 550(7676), 345–353.
  * Sigrist et al. [2012] Sigrist, C. J. et al. (2012).  New and continuing developments at PROSITE.  Nucleic Acids Res., 41(D1), D344–D347.
  * Steinegger and Söding [2018] Steinegger, M. and Söding, J. (2018).  Clustering huge protein sequence sets in linear time.  Nat. Commun., 9(1), 1–8.
  * Steinegger et al. [2019] Steinegger, M. et al. (2019).  Protein-level assembly increases protein sequence recovery from metagenomic samples manyfold.  Nat. Methods, 16(7), 603–606.
  * Strodthoff et al. [2020] Strodthoff, N. et al. (2020).  UDSMProt: universal deep sequence models for protein classification.  Bioinformatics, 36(8), 2401–9.
  * Sundararajan et al. [2017] Sundararajan, M. et al. (2017).  Axiomatic Attribution for Deep Networks.  In D. Precup and Y. W. Teh, editors, Proc. 34th ICML, volume 70 of Proc. Mach. Learn. Res., pages 3319–28. PMLR.
  * Suzek et al. [2014] Suzek, B. E. et al. (2014).  UniRef clusters: a comprehensive and scalable alternative for improving sequence similarity searches.  Bioinformatics, 31(6), 926–932.
  * Taujale et al. [2021] Taujale, R. et al. (2021).  Mapping the glycosyltransferase fold landscape using interpretable deep learning.  Nat. Commun., 12(1), 5656.
  * Tjoa and Guan [2020] Tjoa, E. and Guan, C. (2020).  A survey on explainable artificial intelligence (XAI): Toward medical XAI.  EEE Trans. Neural Netw. Learn. Syst., 32(11), 4793–4813.
  * Torrisi et al. [2020] Torrisi, M. et al. (2020).  Deep learning methods in protein structure prediction.  Comput. Struct. Biotechnol. J., 18, 1301–10.
  * Unsal et al. [2022] Unsal, S. et al. (2022).  Learning functional properties of proteins with language models.  Nat. Mach. Intell., 4(3), 227–245.
  * Upmeier zu Belzen et al. [2019] Upmeier zu Belzen, J. et al. (2019).  Leveraging implicit knowledge in neural networks for functional dissection and engineering of proteins.  Nat. Mach. Intell., 1(5), 225–235.
  * van der Maaten and Hinton [2008] van der Maaten, L. and Hinton, G. (2008).  Visualizing Data using t-SNE.  J. Mach. Learn. Res., 9(86), 2579–2605.
  * Vaswani et al. [2017] Vaswani, A. et al. (2017).  Attention is All You Need.  In Proc. 31st NIPS, page 6000–6010, NY, USA. Curran Assoc., Inc.
  * Vielhaben et al. [2020] Vielhaben, J. et al. (2020).  USMPep: universal sequence models for major histocompatibility complex binding affinity prediction.  BMC Bioinform., 21(1), 1–16.
  * Vig et al. [2021] Vig, J. et al. (2021).  BERTology Meets Biology: Interpreting Attention in Protein Language Models.  In ICLR 2021.
  * Virtanen et al. [2020] Virtanen, P. et al. (2020).  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  Nat. Methods, 17, 261–272.
  * Vu et al. [2023] Vu, M. H. et al. (2023).  Advancing protein language models with linguistics: a roadmap for improved interpretability.  arXiv:2207.00982.
  * Ward et al. [2020] Ward, G. et al. (2020).  An exploration of the influence of path choice in game-theoretic attribution algorithms.  arXiv:2007.04169.
  * Webb et al. [1992] Webb, E. C. et al. (1992).  Enzyme nomenclature 1992. Recommendations of the Nomenclature Committee of the IUBMB on the Nomenclature and Classification of Enzymes. Number Ed. 6. Academic Press.
  * Weissenow et al. [2022] Weissenow, K. et al. (2022).  Protein language-model embeddings for fast, accurate, and alignment-free protein structure prediction.  Structure, 30(8), 1169–1177.e4.
  * Yang et al. [2019] Yang, K. K. et al. (2019).  Machine-learning-guided directed evolution for protein engineering.  Nat. Methods, 16(8), 687–694.
  * You et al. [2018a] You, R. et al. (2018a).  DeepText2GO: Improving large-scale protein function prediction with deep semantic text representation.  Methods, 145, 82–90.
  * You et al. [2018b] You, R. et al. (2018b).  GOLabeler: improving sequence-based large-scale protein function prediction by learning to rank.  Bioinformatics, 34(14), 2465–73.
  * You et al. [2021] You, R. et al. (2021).  DeepGraphGO: graph neural network for large-scale, multispecies protein function prediction.  Bioinformatics, 37(Suppl. 1), i262–i271.
  * Yu et al. [2023] Yu, T. et al. (2023).  Enzyme function prediction using contrastive learning.  Science, 379(6639), 1358–63.
  * Zhang et al. [2018] Zhang, B. et al. (2018).  Prediction of 8-state protein secondary structures by a novel deep learning architecture.  BMC Bioinform., 19(1), 1–13.
  * Zhou et al. [2019] Zhou, N. et al. (2019).  The CAFA challenge reports improved protein function prediction and new functional annotations for hundreds of genes through experimental screens.  Genome biol., 20(1), 1–23.
  * Zhou et al. [2023] Zhou, Z. et al. (2023).  Phosformer: an explainable transformer model for protein kinase-specific phosphoryl. predictions.  Bioinformatics, 39(2).
  * Zou et al. [2019] Zou, Z. et al. (2019).  mlDEEPre: Multi-Functional Enzyme Function Pred. With Hierarchical Multi-Label Deep Learning.  Front. Genet., 9.

##  Appendix A Supplementary material: Introduction

Table A.1: Definitions of UniProt annotations and of corresponding GO terms quoted from the UniProt and EMBL-EBI websites. UniProt annotation |  GO term
---|---
(per amino acid) |  (per protein)
Active site: ‘used for enzymes and indicates the residues directly involved in catalysis’ |  Catalytic activity (GO:0003824): ‘Catalysis of a biochemical reaction […] catalysts are naturally occurring macromolecular substances known as enzymes [that] possess specific binding sites for substrates.’
Binding site: ‘interaction between protein residues and a chemical entity’ |  Binding (GO:0005488): ‘selective, non-covalent, often stoichiometric, interaction of a molecule with one or more specific sites on another molecule.’
Transmembrane region: ‘extent of a membrane-spanning region of the protein’ |  Membrane (GO:0016020): ‘A lipid bilayer along with all the proteins and protein complexes embedded in it an attached to it.’
Short sequence motif: ‘a short (usually not more than 20 amino acids) conserved sequence motif of biological significance. Specific sequence motifs usually mediate a common function, such as protein-binding or targeting to a particular subcellular location, in a variety of proteins.’ |  –

##  Appendix B Supplementary material: System and methods

###  I Protein function prediction tasks and data sets

GO prediction: introduction. We first focus on the protein function prediction
task of inferring Gene Ontology (GO) terms, that the proteins are labeled
with, from their amino acid sequence. This particularly comprehensive task
allows characterizing each protein with respect to several properties, taking
its molecular function, location in the cell, and involvement in different
biological processes into account. GO term prediction has been approached
using a range of deep learning methods in the past years [e.g., Kulmanov et
al., 2017; Kulmanov and Hoehndorf, 2019, 2022; You et al., 2018b, a, 2021;
Strodthoff et al., 2020; Littmann et al., 2021]. The GO consortium [Ashburner
et al., 2000; Consortium, 2020a] aims at representing biological knowledge in
the form of an ontology, i.e. as a graph comprising terms/classes as nodes and
the relationships between them as edges. Actually, GO consists of three
separate ontologies for molecular function (MFO), cellular component (CCO) and
biological process (BPO). The protein database UniProtKB/Swiss-Prot
[Consortium, 2020b] lists the amino acid sequences of proteins, and labels
these proteins with selected, corresponding GO terms. In addition, the
database contains selected annotations on the amino acid level, such as amino
acids that serve as active or binding sites (see Table A.1), which will later
play a role in this work.

GO prediction: setup. Considering that each protein can typically be labeled
with several of the numerous GO terms, the task was phrased here as a multi-
label (‘protein-centric’) prediction problem, where the model was trained to
produce a label vector of about 5000 dimensions (5220 or 5101; see below) for
each protein. The mapping from GO terms to the intricate topology of the
ontology had been simplified by Kulmanov and Hoehndorf [2019], who made the
data available in their accompanying data repository, as follows. A ‘flat’
label vector represents the original graph structure of the ontology. Labels
were propagated from children terms to parents through the hierarchy of the
ontology (in contrast to the original GO, where labels are assigned to
ontology leaves as far as possible). Only a subset of about 5k of the about
45k GO terms in total was considered, because the GO terms had been filtered
to those that occur at least 50 times in the respective data set. Like in
Kulmanov and Hoehndorf [2019], predictions are propagated towards the root
node of the respective ontology tree. Evaluation is always performed with
respect to the set of all GO terms, where the model can only predict the most
frequently occurring GO terms.

GO prediction: data sets. We used both the ‘CAFA3’ and the ‘2016’ data sets by
Kulmanov and Hoehndorf [2019] with 66841 training and 3328 testing samples
with 5220 classes (‘CAFA3’) and, respectively, with 65028 training and 1788
testing samples with 5101 classes (‘2016’) [see Section 2.1 with Table 1 in
Kulmanov and Hoehndorf, 2019]. While the ‘CAFA3’ data set is a widely used
benchmark [Zhou et al., 2019], the time-based ‘2016’ splits of the underlying
UniProtKB/Swiss-Prot database serve for the comparison with the results by You
et al. [2018b, a]; Kulmanov and Hoehndorf [2019]; Strodthoff et al. [2020]
(‘2016’ because entries from Jan.-Oct. 2016 serve for testing). The predictive
performance of the protein function prediction models was assessed separately
for the molecular function, cellular component and biological process
ontologies of GO, with the $F_{\text{max}}$ and $S_{\text{min}}$ metrics
[Clark and Radivojac, 2013] used in ‘CAFA3’ [Zhou et al., 2019], and with the
area under the precision-recall (AUPR) curve [comparable to You et al., 2018a;
Kulmanov and Hoehndorf, 2019; You et al., 2018b; Strodthoff et al., 2020].

EC prediction. During the course of the present study, we narrow down the
prediction task to the inference of the enzymatic function of a protein from
its amino acid sequence. Enzymes are proteins that catalyze chemical
reactions, which can be grouped into seven main classes and further sub-
classes, as categorized by the Enzyme Commission (EC) numbers, a nomenclature
developed by the International Union of Biochemistry and Molecular Biology
[Webb et al., 1992; McDonald et al., 2008]. EC prediction has been approached
with machine learning, e.g., by Dalkiran et al. [2018]; Li et al. [2018]; Zou
et al. [2019]; Strodthoff et al. [2020]; Yu et al. [2023].

The studied prediction task comprised the binary classification of whether a
protein serves as enzyme or not (EC level ‘L0’), the classification into the
main classes of enzymatic reactions (EC level ‘L1’), and the differentiation
among different sub-classes of enzymatic reactions (EC level ‘L2’). We
assessed the predictive performance on the ‘EC40’ (58018 samples) and ‘EC50’
(104940 samples) data sets by Strodthoff et al. [2020] (see their
Supplementary Section S1), using accuracy as metric. The numbers in ‘EC40’ and
‘EC50’ refer to the similarity threshold between train and test splits of 40%
similarity (more difficult task) and, respectively, 50% (easier task). The
data set comprises the traditional six main EC classes (oxidoreductases,
transferases, hydrolases, lyases, isomerases, ligases) leaving apart the new
seventh EC class of translocases.

Data availability and composition of the training set. GO data by Kulmanov and
Hoehndorf [2019] from their data repository and EC data by Strodthoff et al.
[2020] (see their Supplementary Section S1) were preprocessed as detailed in
the code repository by Strodthoff et al. [2020], resulting in separate data
sets for GO ‘2016’ (a.k.a. ‘temporalsplit’), for GO ‘CAFA3’, and for EC40 and
EC50 on levels L0, L1, and L2. Dedicated validation (a.k.a. development)
splits were available for the EC and GO ‘CAFA3’ tasks. For the time-based GO
‘2016’ split, a 10 % subset was randomly sampled from the training data and
used as validation split [comparable to Kulmanov and Hoehndorf, 2019]. As
described below, we use the validation set for model selection and report the
test set score.

###  II Finetuning pretrained transformers

Considered transformer models. We finetuned the transformer model ‘ProtBert-
BFD’ [Elnaggar et al., 2022] to the GO term and EC number prediction tasks
detailed earlier in Section I of Appendix B. ‘ProtBert-BFD’ is a variant of
BERT [Devlin et al., 2018], with 420 million parameters distributed among 16
heads in 30 layers, that had been pretrained in a self-supervised way with a
masked token prediction task on the massive metagenomic protein sequence
collection titled ‘Big Fantastic Database’ (BFD), which comprises more than
two billion sequences [Jumper et al., 2021; Steinegger et al., 2019;
Steinegger and Söding, 2018].

In addition, we finetuned also the more complex and memory intensive
‘ProtT5-XL-UniRef50’ model [Elnaggar et al., 2022] on the GO and EC tasks,
using the last hidden state of the encoder. ‘ProtT5-XL-UniRef50’ is based on
the ‘Text-To-Text Transfer Transformer’ [T5; Raffel et al., 2020] with about
three billion parameters (1.2 billion for the encoder used here) and had been
pretrained on 45 million sequences from UniRef50 [Suzek et al., 2014]. We
abbreviate ‘ProtBert-BFD’ as ‘ProtBert’ and ‘ProtT5-XL-UniRef50’ as ‘ProtT5’
in the following text.

Classification head. The final classification head comprised a hidden layer
and then a linear layer as connection to the dimensionality of the respective
label vector. A hidden layer with $n=50$ neurons was chosen for the simpler EC
task. For the more demanding GO task, the size of the hidden layer was
increased to $n=2\times 5220$ (GO ‘CAFA3’) and, respectively, to $n=2\times
5101$ (GO ‘2016’), such that the hidden layer contained twice as many neurons
in comparison to the dimensionality of the respective label vector. Rectified
linear units served as activation function, and layer normalization and a
dropout layer with a dropout rate of ten percent served for regularization.

Finetuning procedure. For finetuning, the input sequence was cropped to a
maximum sequence length of 1000 amino acids per protein, if necessary. The
encoder features created by the transformer were pooled by concatenating the
classification token, the maximum and the mean of the features along the
sequence, and the sum of the features divided by the square root of the
sequence length [Reimers and Gurevych, 2019]. Binary cross entropy combined
with a sigmoid layer served as loss function for the multi-label GO task, and
cross entropy loss combined with softmax for the multi-class EC task.
Gradients were accumulated over 64 batches of size one each. Adam [Kingma and
Ba, 2015] was selected as optimizer. The learning rate was set to $5\times
10^{-6}$ for the encoder, which was frozen for the first epoch, and to
$3\times 10^{-5}$ for the classification head. For the multi-label GO task,
the rate was additionally scheduled with a linear increase over 500 warm-up
steps in the beginning and then a decrease to zero following a half-cosine
function calculated for a total number of 20000 training steps.

Model selection. After each training epoch, the loss (GO) or accuracy (EC) on
the validation split was monitored and the model with the minimum loss (GO) or
maximum accuracy (EC) was saved. Training was repeated until no further
improvement was observed on the validation split. The best model was typically
found after around 2-13 epochs in the case of ProtBert and after around 2-7
epochs in the case of ProtT5.

##  Appendix C Supplementary material: Implementation

The pretrained models of ProtBert and, respectively, ProtT5 [Elnaggar et al.,
2022], were connected to an additional pooling layer and classification head
(see Appendix B, II), and finetuned on GO data provided by [Kulmanov and
Hoehndorf, 2019] and, respectively, EC data by Strodthoff et al. [2020] (see
their Supplementary Section S1), using PyTorch [Paszke et al., 2019] with
PyTorch Lightning as wrapper. The adaptation of the integrated gradients
algorithm acting separately on each transformer head in every layer (see
Section 3) was implemented in Python building on the layer integrated
gradients method of Captum [Kokhlikyan et al., 2020]. Annotations were
obtained from the databases UniProt [Consortium, 2020b] and PROSITE [Sigrist
et al., 2012] in combination with GO [Ashburner et al., 2000; Consortium,
2020a]. Statistical calculations were performed using SciPy [Virtanen et al.,
2020] and statsmodels [Seabold and Perktold, 2010]. Performance on the GO task
was assessed based on an adaptation by [Strodthoff et al., 2020] from
[Kulmanov and Hoehndorf, 2019]. Source code for model training, model
interpretation, statistics and visualization can be accessed at
https://github.com/markuswenzel/xai-proteins.

##  Appendix D Supplementary material: Results and Discussion

Research question. Here, we aim to answer the question whether finetuning of
large pretrained transformer models leads to a competitive performance in
protein function prediction.

GO prediction on the CAFA3 data set. Table D.1 shows the performance results
for GO term prediction by finetuning and testing ProtBert and ProtT5 on the
‘CAFA3’ splits. Performance was assessed with the protein-centric
$F_{\text{max}}$ metric, separately for the molecular function (MFO),
biological process (BPO), and cellular component ontologies (CCO) that
comprise the GO. For comparison with the state-of-the-art, we added results
reported in the literature [Olenyi et al., 2023; Bernhofer et al., 2021;
Littmann et al., 2021; Kulmanov and Hoehndorf, 2019; You et al., 2018b] to
Table D.1. For model evaluation, we distinguish between single-models and
ensemble models. In the former category, the finetuned ProtT5 turns out to be
the model that outperforms both prior deep-learning-based approaches based on
transformers (goPredSim) and convolutional neural networks (DeepGOCNN) as well
as the very strong MSA-based baseline DiamondScore, on all ontologies MFO, BPO
and CCO.

Turning to ensemble model, we find that the ensemble of a finetuned ProtBert
(followed by ProtT5) and DiamondScore [MSA-based; Kulmanov and Hoehndorf,
2019] produces the best predictions on BPO and CCO [scores were combined with
a weighted sum; comparable to Kulmanov and Hoehndorf, 2019, see their Section
2.3]. The ensemble method ‘GOLabeler’ [You et al., 2018b] leads on MFO.
GOLabeler had won the original CAFA3 challenge on all three ontologies BPO,
CCO, and MFO [see Fig. 3 of Zhou et al., 2019, where the $F_{\text{max}}$
results are reported].

Table D.1: GO term prediction by finetuning and testing ProtT5 and ProtBert on the ‘CAFA3’ splits, next to state-of-the-art results reported in the literature. Predictive performance was evaluated with the protein-centric $F_{\text{max}}$ metric. The molecular function, biological process and cellular component ontologies of GO are abbreviated as MFO, BPO and CCO. Best single-model results are marked in bold face. Best overall results are underlined. Arrows ($\rightarrow$) mark results achieved in this work: ProtT5 provides the best overall single-model performance in the MFO, BPO and CCO categories. Ensembling with the MSA-based DiamondScore predictions is highly effective. | Method |  MFO, $F_{\text{max}}$ (($\uparrow$) |  BPO, $F_{\text{max}}$ ($\uparrow$) |  CCO, $F_{\text{max}}$ ($\uparrow$)
---|---|---|---|---
Single-models |  $\xrightarrow{}$ ProtT5-XL-UniRef50, finetuned | 0.523 | 0.442 | 0.641
$\xrightarrow{}$ ProtBert-BFD, finetuned | 0.503 | 0.441 | 0.630
goPredSim, ProtT5 embeddings [Olenyi et al., 2023] |  0.52$\pm$3% |  0.38$\pm$2% |  0.59$\pm$2%
goPredSim, SeqVec embeddings [Bernhofer et al., 2021] |  0.52$\pm$2% |  0.37$\pm$2% |  0.58$\pm$2%
goPredSim, SeqVec embeddings [Littmann et al., 2021] |  0.50$\pm$3% |  0.37$\pm$2% |  0.57$\pm$2%
DeepGOCNN [Kulmanov and Hoehndorf, 2019] | 0.420 | 0.378 | 0.607
DiamondScore [Kulmanov and Hoehndorf, 2019] | 0.509 | 0.427 | 0.557
Ensembles |  $\xrightarrow{}$ ProtT5-XL-UniRef50, finetuned + DiamondScore | 0.565 | 0.473 | 0.635
$\xrightarrow{}$ ProtBert-BFD, finetuned + DiamondScore | 0.580 | 0.480 | 0.640
DeepGOPlus [Kulmanov and Hoehndorf, 2019] | 0.544 | 0.469 | 0.623
GOLabeler [You et al., 2018b], [Fig. 3 of Zhou et al., 2019] | 0.62 | 0.40 | 0.61

GO prediction on the 2016 data set. Table D.2 shows the GO term prediction
results of the finetuned ProtT5 and ProtBert models on the more recent ‘2016’
data set Kulmanov and Hoehndorf [2019]; You et al. [2018b], as well as the
literature results [compiled by Kulmanov and Hoehndorf, 2019; Strodthoff et
al., 2020, which report AUPR and $S_{\text{min}}$ results in addition to
$F_{\text{max}}$]. Again, we distinguish between single and ensemble models in
the following discussion.

Concerning single-models, different methods obtain the best results, depending
on the considered metric and ontology. ProtT5 is the best single-model as
measured by AUPR on all three ontologies MFO, BPO, CCO. ProtT5 also obtained
the largest $F_{\text{max}}$ on BPO and CCO. DiamondScore, an MSA-based
method, is the best single-model with respect to the $S_{\text{min}}$ metrics
on the ontologies MFO and BPO, and achieves the best $F_{\text{max}}$ on MFO.
DeepGOPlus has the lowest and thus best $S_{\text{min}}$ for CCO.

Turning to ensemble models, we find that the best overall results are obtained
again by different ensemble methods. The finetuned ProtT5 combined with
DiamondScore leads with respect to the AUPR metric on the ontologies BPO and
CCO, and with respect to $F_{\text{max}}$ on BPO. Both ProtT5 and DeepGOPlus
have the best $F_{\text{max}}$ on CCO. DeepText2GO is best on MFO according to
$F_{\text{max}}$ and AUPR, and on CCO according to $S_{\text{min}}$. GOLabeler
features the smallest and thus best $S_{\text{min}}$ on MFO and BPO.

Table D.2: GO term prediction by finetuning and testing ProtT5 and ProtBert on
the ‘2016’ data set splits Kulmanov and Hoehndorf [2019]; You et al. [2018b],
next to the literature results collected by Kulmanov and Hoehndorf [2019] and
Strodthoff et al. [2020]. Best single-model results are marked in bold face
(highest $F_{\text{max}}$ and AUPR; lowest $S_{\text{min}}$). Best overall
results are underlined. Arrows ($\rightarrow$) point towards models finetuned
in this paper, with competitive results.

|  |  $F_{\text{max}}$ ($\uparrow$) |  $S_{\text{min}}$ ($\downarrow$) |  AUPR ($\uparrow$)
---|---|---|---|---
| Method | MFO | BPO | CCO | MFO | BPO | CCO | MFO | BPO | CCO
Single-models |  $\rightarrow$ ProtT5-XL-UniRef50, finetuned | 0.526 | 0.441 | 0.685 | 10.099 | 35.229 | 8.030 | 0.487 | 0.395 | 0.730
$\rightarrow$ ProtBert-BFD, finetuned | 0.503 | 0.424 | 0.677 | 10.421 | 35.401 | 8.198 | 0.460 | 0.369 | 0.714
UDSMProt [Strodthoff et al., 2020] | 0.481 | 0.411 | 0.682 | 10.505 | 36.147 | 8.244 | 0.472 | 0.356 | 0.704
DeepGO [Kulmanov et al., 2017] | 0.449 | 0.398 | 0.667 | 10.722 | 35.085 | 7.861 | 0.409 | 0.328 | 0.696
DeepGOCNN [Kulmanov and Hoehndorf, 2019] | 0.409 | 0.383 | 0.663 | 11.296 | 36.451 | 8.642 | 0.350 | 0.316 | 0.688
DiamondScore [Kulmanov and Hoehndorf, 2019] | 0.548 | 0.439 | 0.621 | 8.736 | 34.060 | 7.997 | 0.362 | 0.240 | 0.363
Ensembles |  $\rightarrow$ ProtT5-XL-UniRef50, finetuned + $\Diamond$ | 0.591 | 0.482 | 0.699 | 8.583 | 33.752 | 7.467 | 0.571 | 0.431 | 0.736
$\rightarrow$ ProtBert-BFD, finetuned + $\Diamond$ | 0.583 | 0.470 | 0.693 | 8.687 | 33.560 | 7.597 | 0.558 | 0.423 | 0.725
UDSMProt [Strodthoff et al., 2020] \+ $\Diamond$ | 0.582 | 0.475 | 0.697 | 8.787 | 33.615 | 7.618 | 0.548 | 0.422 | 0.728
DeepGOPlus [Kulmanov and Hoehndorf, 2019] | 0.585 | 0.474 | 0.699 | 8.824 | 33.576 | 7.693 | 0.536 | 0.407 | 0.726
DeepText2GO [You et al., 2018a] | 0.627 | 0.441 | 0.694 | 5.240 | 17.713 | 4.531 | 0.605 | 0.336 | 0.729
GOLabeler [You et al., 2018b] | 0.580 | 0.370 | 0.687 | 5.077 | 15.177 | 5.518 | 0.546 | 0.225 | 0.700

$\Diamond$: DiamondScore

EC prediction. Table D.3 shows how accurately ProtBert and ProtT5 predicted
the EC number of a given protein from its amino acid sequence, separately for
the three EC levels (L0, L1, L2) and two train/test splits (EC40, EC50). For
comparison, the table lists the results of ‘UDSMProt’ too [Strodthoff et al.,
2020]. ProtBert and ProtT5 predict the EC number more accurately, with a large
gap, in comparison to UDSMProt.

Table D.3: The EC number of each protein was predicted from the amino acid sequence. Different EC levels (L0: enzyme or not, L1: six main enzyme classes, L2: enzyme sub-classes) and train/test splits (EC40/EC50) were assessed separately, with accuracy as performance metric, and compared to the ‘UDSMProt’ results [cf. Strodthoff et al., 2020, Table 1; incl. baseline of a convolutional neural network (CNN) operating on MSA features]. Arrows ($\rightarrow$) mark results established here. | EC40 | EC50
---|---|---
Model | L0 | L1 | L2 | L0 | L1 | L2
$\rightarrow$ ProtT5-XL-UniRef50 | 0.966 | 0.968 | 0.954 | 0.985 | 0.993 | 0.985
$\rightarrow$ ProtBert-BFD | 0.960 | 0.958 | 0.945 | 0.981 | 0.990 | 0.982
UDSMProt (‘Fwd+bwd; pretr.’) | 0.91 | 0.87 | 0.84 | 0.96 | 0.97 | 0.95
CNN (‘Baseline, seq’) | 0.84 | 0.61 | 0.47 | 0.92 | 0.80 | 0.79

Discussion. Finetuning ProtT5 and ProtBert, as close follower, on the tasks of
GO term as well as EC number prediction results in a strong predictive
performance that is competitive with, and in several cases, better than state-
of-the-art methods from the literature (depending on the metric and ontology
under consideration, in the GO case). This outcome highlights again the
benefits of transferring pretrained universal language models to downstream
tasks in the proteomics field [cf., e.g., Strodthoff et al., 2020; Rao et al.,
2019; Elnaggar et al., 2022].

ProtT5 consistently ranks better than ProtBert both on ‘CAFA3’ and on the more
recent, harder ‘2016’ benchmark (only the ProtBert-DiamondScore ensemble forms
an exception on ‘CAFA3’ with respect to $F_{\text{max}}$ and on ‘2016’ with
respect to $S_{\text{min}}$ in the BPO case). DeepGOCNN and DiamondScore are
competing single-models with results available on both benchmarks. ProtT5 and
ProtBert outperform DeepGOCNN in both benchmarks on all metrics. ProtT5
outperforms DiamondScore, which is based on MSAs, in both benchmarks in many
cases, except with respect to $F_{\text{max}}$ for MFO in ‘2016’ and with
respect to $S_{\text{min}}$. The approach of finetuning the entire model
including the encoder shows its strength in the ‘CAFA3’ benchmark for BPO and
CCO, where the finetuned ProtT5 and ProtBert outperform the ‘goPredSim’ model
which is based on nearest-neighbor-lookup using features extracted from ProtT5
embeddings [Olenyi et al., 2023]. The observation that this nearest-neighbor-
lookup is competitive in the MFO ontology might be related to the relatively
dense annotation of the MFO. In summary, we showed that finetuning pretrained
large transformer models leads to competitive results for protein function
prediction tasks, in particular in the most relevant comparison in the single-
model category.

|  Corr(relevance, annotation) $>0$ |  Relevance $>0$ | Overlay
---|---|---|---
PROSITE |  |  |
Active sites |  |  |
Transmem. reg. |  |  |
Figure D.1: Inside ProtBert; focus on GO ‘catalytic activity’ (GO:0003824). Left panels: Results of the correlation analysis between relevance attributions indicative for the GO term ‘catalytic activity’ with PROSITE patterns and UniProt annotations as ‘active sites’ and ‘transmembrane regions’ per head and layer (negative logarithm of corrected p-values of t-tests over correlation coefficients; significance thresholds were overlaid as masks). Center panels: (Negative logarithm of corrected p-values of) Wilcoxon tests inspecting whether the relevance (sum along the sequence) was significantly positive. Right panels: Heads with both a significantly positive attribution-annotation-correlation and a significantly positive relevance (overlay of left and center panels). Only results for corresponding PROSITE patterns (top), UniProt ‘active sites’ (center) and ‘transmembrane regions’ (bottom) are shown. ‘Binding sites’ and ‘motifs’ are ommited, since a significantly positive attribution-annotation-correlation and combined with a positive relevance was not observed in these cases.  |  Corr(relevance, annotation) $>0$ |  Relevance $>0$ | Overlay
---|---|---|---
PROSITE |  |  |
Transmem. reg. |  |  |
Figure D.2: Inside ProtBert; focus on GO ‘binding’ (GO:0005488). Left: Correlation of relevance for ‘binding’ with corresponding PROSITE patterns (top) or UniProt ‘transmembrane regions’ (bottom) per head and layer. Center: Positive relevance. Right: Heads with both a significantly positive attribution-annotation-correlation and a significantly positive relevance. Results for ‘active/binding sites’ and ‘motifs’ are not shown, given the absence of significantly positive attribution-annotation-correlations in these cases.  | EC1 | EC2 | EC3 | EC4 | EC5 | EC6
---|---|---|---|---|---|---
Active sites |  |  |  |  |  |
Binding sites |  |  |  |  |  |
Transmem. reg. |  |  |  |  |  |
Motifs |  |  |  |  |  |
Figure D.3: Inside ProtBert finetuned to EC number classification. Rows:
Results of the correlation analysis are presented separately for the different
types of annotations (UniProt ‘active sites’, ‘binding sites’, ‘transmembrane
regions’, ‘short sequence motifs’). Columns: Results for the six main EC
classes (at level L1). Pixels in the matrix plots represent the 16 transformer
heads and 30 layers. Each pixel shows the negative logarithm of the p-value
resulting from a t-test across the coefficients of correlation between
relevance attributions and annotation ($p<0.05$; after correction for multiple
comparisons and thresholding) overlaid with the mask of those heads/layers
that featured a significantly positive relevance too.