Transformers trained on proteins can learn to attend to Euclidean distance
Authors: Isaac Ellmen
Constantin Schneider
Matthew I. J. Raybould
Charlotte Deane
Venue: Accepted by TMLR ()

Abstract: While conventional Transformers generally operate on sequence data, they can be used in conjunction with structure models, typically SE(3)-invariant or equivariant graph neural networks (GNNs), for 3D applications such as protein structure modelling. These hybrids typically involve either (1) preprocessing/tokenizing structural features as input for Transformers or (2) taking Transformer embeddings and processing them within a structural representation. However, there is evidence that Transformers can learn to process structural information on their own, such as the AlphaFold3 structural diffusion model. In this work we show that Transformers can function independently as structure models when passed linear embeddings of coordinates. We first provide a theoretical explanation for how Transformers can learn to filter attention as a 3D Gaussian with learned variance. We then validate this theory using both simulated 3D points and in the context of masked token prediction for proteins. Finally, we show that pre-training protein Transformer encoders with structure improves performance on multiple downstream tasks, yielding competitive performance with custom structural models. Together, this work provides a basis for using standard Transformers as hybrid structure-language models. The code is available at: https://github.com/oxpig/attending-to-distance.

Published in Transactions on Machine Learning Research (07/2025)
Transformers trained on proteins can learn to attend to Euclidean distance
Isaac Ellmen
Deportment of Statistics University of Ouford
Constantin Schncider Eascientia
Matthew I.J. Raybould
Department of Statistics University of Oxford
Charlotte M. Deane Department of Statistics University of Onford
deane or, ac. ac.uk
Reviewed OpenReview: https: https:///peerreviau
openreview. net/ forum? id= mU59bDygqu
Abstract
While conventional Transformers generally operate on sequence data, they can be used in conjunction with structure models, typically SE(3)-invariant or equivariant graph neural networks (GNNs), for 3D applications such as protein structure modelling. These hybrids typically involve either (1) preproccssing/ htap://ssssinnooooeiiinngng structural features input for Trans- formers or (2) taking Transformer embeddings and processing them within structural representation. However, there is evidence that Transformers can learn to process struc- tural information on their own, such as the AlphaFood3 structural diffusion model. In this work we show that Transformers can function independently as structure models when passed linear embeddings of coordinates. We first provide a theoretical explanation for how Transformers can learn to filter attention as Gaussian with learned variance. We then validate this theory using both simulated 3D points and in the context of masked token prediction for proteins. Finally, show that pre-training protein Transformer encoders with structure improves performance on multiple downstream tasks, yielding competitive performance with custom structural models. Together this work provides a basis for us- ing standard Transformers as hybrid strueture-language models. The code is available at: -oodssitnnee-
1 Introduction
1.1 Background
Transformers typically operate on sequential data, however many applications of Transformers benefit from ability to learn geometric reasoning. For instance, ESM-2 LLn al. 2023 demonstrates that in order to effectively predict masked lokens in protein sequences, the model has learned some ability to predict protein structures. Other tasks such as image processing even natural language processing may benefit from an internal representation of objects 3D space. However, it is unclear how Transformers can learn to use 3D representations to perform spatial reasoning. To this end, custom structural Transformers have been created
Now Xyme

---

Published in Transactions on Machine Learning Research (07/2025)
which model data graphs and represent distance between nodes edge features in order to perform
SF(3)-invariant attention (Ingraham et al. 2019 Fuchs et al. 2020: Liao & Smidt 2023 Liao et al. 2023
SE(3) invariance means that all functions of coordinates reduce to functions of relative distance. That is, for every function f of two coordinates 21 and 22, there exists an equivalent function which depends only on their relative distance: [ = g( oz)) Since relative distance Euclidean space is defined as | = = + y² + +2, it may be casicr to learn functions of the square of the relative distance - a lincar combination of functions of the individual coordinates.
In this manuscript, we investigate how conventional Transformers can learn to approximate functions of the squared distance between points, thereby learning an approximately (E(3)-invariant mcasure of distance. In short, show that "out of the box" Transformers can act as 3D structural models. Our main contributions are (1) to provide theorctical explanation for how standard Transformers can learn to mcasure distance and perform structural reasoning, (2) to show that Transformers indeed learn Gaussian functions of distance and investigate efficient data augmentation methods which can be used to learn SE(3), and (3) to train a protein masked token prediction model with coordinates and show that finetuning it for function prediction yiclds a model which outperforms structural GNNs.
1.2 Prior work
There have been prior approaches to merge Transformers with (((3)-(in/qquiivvaiia models, espocially for computational chemistry and 3D point clouds. Some methods add attention blocks to SE(3) GNNs to create SE(3) (Ev3)-innariant GNN Transformers (Fuchs et al. 2020 Liao & Smidt 2023). These have shown good results a number of tasks, however tend to be memory-intenste edges, which grow as n² for fully-connected graphs. As a result, the graph connectedness of the GNNs is typically limited to k-nearest neighbours. In contrast, memory-efficient attention implementations such as FlashAttention (Dao et al. 2022
particularly because attention is performed
Dao 2023 have enabled linear- memory standard Transformers.
Previous works have demonstrated that sequence-only protein Transformers can learn attention maps which correlate with physical contacts (Lin et al., 2023 Vig al.) 2020). However, these works do not formally model structure and so are limited by which contacts be predicted from purcly sequential patterns. To overcome this, Transformers are often paired with structural models, for instance met such SaProt (Su et al. 2023), ProSST (Li et al.) 2024), ESM-1F (Hsu et al.] (022), and ESM3 (Hayes et al., 2024) use custom graph-based modules to create structural tokens which are fod into standard Transformers. In contrast, our work shows that standard Transformers are natively capable of using coordinates to model structure and measure distance.
Similarly, AlphaFold2 (Jumper et al. 2021 and ESMFold (Lin et al. 2023 use Transformers to preprocess protein sequences for structure prediction. Again, these preprocessed representations have been shown to cor- relate with structural contacts. Alld2 makes this explicit during training by minimizing a distrogram loss which encourages the EvoFormer to learn structural contacts. However, it is unclear if these representa- tions are learning to explicitly embed coordinates in 3D and both models still require an SE())-oquvvarrian GNN structure module to actually produce 3D structures.
In building AlphalFodd3, DeepMind replaced AlphaFold2's (E(()-equivariant structure module with linearly embedded coordinates fed into diffusion transformer (Abramson et al. 2024 AlphaFold3's structure module uses inner product attention with pair bias learned from the pair representation. At present, this still requires quadratic memory, however does indicate that nearly-standard Transformers with lincarly embedded coordinates can learn on structure. A similar approach was used by Proteina for protein structure generation [Geffner et al. 2024). Here, we explore how such linearly embedded coordinates can be used by standard Transformer attention modules to measure the Euclidean distance between tokens, and, in contrast to prior work, show that no modifications are necessary for the standard Transformer architecture learn to perform structural reasoning.
2

---

Published in Transactions on Machine Learning Research (07 (07/2025)
2 Theory
Here we introduce our theory for how Transformers can learn to measure Euclidean distance. Throughout this section, we assume that all coordinates are small. This can be achieved by scaling the inputs and outputs manually or via learned weights in the linear input/output maps. For more detail, see Appendix A.1
2.1 Gaussian spatial attention
Consider a pre-norm Transformer (Xiong et al. 2020) operating on sequence of embedded positions, E(r;)) For simplicity, assume that the key and query embeddings for all heads are trivial (Q - K Ia). Such mappings can easily be learned as long as the head dimension is at least d. Note that this assumption does not affect the non-positional component of the query embedding matrices since, at a minimum, the model can learn a block diagonal matrix with the standard query/key matrices followed by Ia over the embedded coordinate dimensions. Then, in the first. layer of the Transformer, the attention matrix for all heads will be:
Ai. SM( LN(B(G))
(1)
Vd
Where SM denotes the softmax function and LN denotes LayerNorm [aa et al. 2016
Our objective is to determine an embedding, E, such that the attention As a monotone decreasing function of the Fuclidean distance between oi and zi.- particular, we choose such that, for some a, b IR:
LE(E(E)) LN(E(E)) = -a - æz1² +
(2)
Then,
Aij SM LN(E(G))-LN(((E))
Vd
b
SMM
vd 2 - 2 2 +
(3)
Vi
SM
10 - 17
Vd
In particular, the unnormalized attention paid to zj by x; is:
A Anonormin R J/ 1a -2)1
Which is a Gaussian of the relative distance between the two points, iee, Transformers can learn to ap- proximate a 3D Gaussian to gale attention values and selectively attend to points nearby in 3D space. By learning individual LayerNorm gains or Q/K mappings, each head can tune the variance this Gaussian filter, which allows each head determine the appropriate spatial resolution for that type of information.
Below, we provide embeddings which satisfy Equation 2 and lead to Gaussian attention filters.
3

---

Published in Transactions on Machine Learning Research (07 (075)
2.2 Spatial positional embeddings
For simplicity, consider the case of 1 spatial dimension and provide 4-dimensional embeddings, which satisfy Equation
The 3 (or n) dimensional case is similar. Consider the embeddings:
Etrig (1) (sin(x): sin(()), cos(1), cos(x))
Etin (Ir) (2t, -2, 111
(5)
x² x2
Equed (1) (2i,- -2, 2 1)
Here, Eirig is similar to the standard sinusoidal positional encoding for linear sequences (Vawwan et al. 2017 Then, Ein and Equed can be thought as the first and second order approximations Errig respectively. can be shown (sce Appendix A.2 that for all three emboddings:
LN(E(e:)) LN(B(1))) ss -2/2; -xyl +
(6)
Subject to |x; - æyy being small for Exrig and subject to |a,] and |رy| being small for Ein and Equod- This surprising result for Ejan stems from the non-linean non-linearity of LayerNorm causing the constant terms to be locally quadratic. For more details, see Appendix [A.222] An important consequence of this result. that simple linear embeddings of positions can still lead to an approximately 3D Gaussian filter of relative distance for attention.
Additionally, the approximation Equad is better than that of Egin; however the requirement to encode x² explicitly makes it harder for individual attention heads to rescale the positions appropriately, since heads have to learn the linear and quadratic scaling lerms separately. However, in A.2 we show that ReGLU and SwiGLU activation functions (Shazeer] 2020) are capable of learning exactly quadratic functions of their input, which may allow some modern Transformers to learn positional embeddings of the form Equad after the appropriate rescaling. This may be useful benefit of GLU activation functions for models such
AlphaFold3 (Abramson et al. 2024
3 Experiments
3.1
Simulated points
To test our theory how Transformers learn to measure distance, we designed a Transformer encoder which is truncated such that the output is the unnormalized attention matrix for single head. A diagram of this model is shown in Figure [ We computed the loss the l difference between the output matrix and the matrix Aij e -( This corresponds to the prenormalized softmax of the negative square of the relative distance between points, as predicted by our theory. The data consisted of 10,000 "structures", each with five 3-dimensional points with coordinates randomly selected between 0 and 200. Unless otherwise indicated, the Transformer encoder for all experiments has three layers (is truncated at the third layer), an embodding dimension of 256, a foodforward dimension of 1,024, 8 heads, pre-normalization and ReLU activation. The models were trained with a batch size of 16 using the Adam optimizer with peak learning rate of 4 x 10-4 which is reached after 4,000 warmup steps, and then is quadratically decayed.
200
As in AlphaFold3 (Abramson et al. 2024), we transform the input structures before they are passed through the model. Whenever a structure is loaded, its points are recentred, randomly rolated, and rescaled by a factor of 1. This has two benefits. First, recentering and rescaling the points ensures that all coordinates stay relatively small, even before the embodding layer has learned an appropriate mapping. Sccond, recentering and randomly rotating gives the model resilience to translations and rotations which encourages it learn distance measure which is truly (En))iinariiant.
4

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
Partial Attention
1.0
0.4
0.5
0.5
0.5
0.4
1.0
0.4
0.8
0.3
Feed
Forward
0.5
0.4
1.0
0.8
0.7
0.5
0.8
0.8
1.0
0.6
Layer Norm
0.5
0.3
0.7
0.6
1.0
Prediction
(N 1
Output
Multi-Head Atlention
erp(QK" max QKT)))
1 - xjj
Linear
Linear
Layer Norm
Layer Norm
Linear
t
Coords
Embedding
Figure 1: Overview of the simulated experiment model. Coordinates are passed through Transformer encoder which is truncated such that the output is the unnormalized attention for single head. Loss is computed as the difference between the attention for each pair and Gaussian of the relative distance between those points.
3.1.1 Transformers can attend | 2
In Section show that Transformers are theoretically capable of learning Gaussian functions of distance. We experimented with learning -- for different powers p, ranging from 0.5 to 4, in increments of 0.5. Figure 2a shows the relationship between and the validation loss. As expected, Transformers can learn to reproduce the - attention matrix most accurately which shows that a Gaussian is the most natural way for Transformers to learn to filter attention spatially.
3.1.2 Transformers need embedding dimensions learn distance in R"
Next, we explored how large spatial embedding must be to learn good approximation of distance for IR.. For the case of 1 spatial dimension, we provide theoretical 4-dimensional embedding. We trained models to predict Euclidean distance in R2t for n E 1,2,3,4 4 using head dimensions from to 8. Figure 2b shows the relationship between validation loss and head dimension. We found that the model only needs n + 2 head dimensions to learn a good approximation of distance for all spatial dimensions. This means that learn a good approximation of distance in R3 a model must reserve least 5 embedding dimensions for each head. This is surprisingly compact requirement which should be casily accommodated even in small Transformers.
3.1.3 SE(3) transformations improve learned SE(3) (((3))iivrriance
We investigated whether Transformers will learn to overfit training data low data regime and if this can be prevented. This could also correspond to a scenario where there is only strong structural signal in small number of training examples. We reduced the number of training points to 100 and measured
5

---

Published in Transactions on Machine Learning Research (07 (025)
Loss depending on for exp(
Loss depending on head dimension
0.35
Spatial dimension 1
Spatial dimension 2
0.30
Spatial dimension
Validation loss
Validation loss
0.25
Spatial dimension 4
0.004
0 0.003
0.002
0.001
0.5
1.0
1.5
2.0
3.0
3.5
4.0
2
Exponent (p)
Head dimension
(a) Loss vs exponent
(b) Loss vs head dimension
Figure 2: (a) Validation loss as function of the exponent The loss is lowest for 2 which corresponds to learning a Gaussian function of distance. (b) Validation loss function of head dimension for different spatial dimensions. Models need a head dimension n + to accurately measure distance in IR..
Rolling average loss epoch
distance from prediction truth random rotation
Raw validation
Validation loss
0.08
Raw train
SE(3) divergence
Augmented validation
0.06
Augmented train
i1 distance
0
Loss
0.02
0.00
0100
2000
4000
6000
8000
10000
Raw
Augmented
Epoch
Loss epoch
(b) Loss and SE(3) divergence
Figure 3: (a) Training and validation loss per epoch in a low data setting. The rolling average of the last 500 epochs shown. Raw coordinates show large divergence between training and validation loss whereas randomly rotated coordinates show near perfect aligument, resulting in substantially lower validation loss. (b) The final validation loss is very close to the average predicted distance between randomly rotated structures, indicating that the validation loss is minimized by learning more SE(3)-invariant measure of distance.
the training and validation loss. To test the importance of data augmentation, we trained models with and without the random rotations (Figure 3a}. The raw coordinates clearly demonstrate overfitting while the randomly rotated coordinates show near perfect alignment between training and validation loss. Importantly, this form of data augmentation does not require creating new data points, only rotating the training data each epoch. For more details see Appendix A.3
We measured the average & distance between predictions of randomly rotated structures for the models trained with and without random rotations, as a mcasure of SE(3) divergence (Figure 3b). In both cases the SE(3) divergence was almost the same the validation loss, indicating that randomly rotating training structures reduces overfitting by encouraging models to learn SF(33-invariin measure of distance.

---

Published in Transactions on Machine Learning Research (07/ (07/2025)
3.2 Proteins
Proteins are a natural fit for structural Transformers because they are composed of linear sequences embedded in 3D space. As such, their propertics depend on both sequential and structural features. We considered two tasks in protein modelling: predicting masked tokens as a pretraining objective and predicting protein function conditioned on embeddings generated by pretrained models. For all protein experiments we used the GO PDB dataset from DeepFRI GGlioorijevćć et al. 2021 which comprises 336K protein chains.
3.2.1 Pretraining a structural protein language model
To test a Transformer's ability to learn useful structural patterns in proteins, trained ESM/BERTI style (Rives et al., 2021 Devlin et al. 2019) model to complete masked token prediction. We trained two models: one with coordinates ("coords model') and the other without ('non-coords model'). The version with coordinates added a linear embedding of the coordinates to the token embedding in the same way as the simulated experiments. Both models were very similar to the smallest publicly released ESM1 model, consisting of a 6-layer Transformer encoder with a hidden dimension of 768, 12 attention heads, a feedforward dimension of 2048, and GcLU activation (Hendrycks Gimpel 2023). p. As in ESM, we omitted dropout. To prevent the model from focusing too much on linear positional information, we used Sinusoidal Positional Encodings rather than Rotary Positional Encodings (Su et al. 2022). As is common in masked token prediction, we masked 15% of tokens. Of the masked tokens, 80% were replaced with a [MASK] token, 10% were replaced with a random amino acid, and 10% were left unchanged. We clustered the data by 50% sequence identity using MMSeqs2 (Steinegger & Söding, (017] and randomly held out 1% of the clusters to use as a validation set. We trained each model for 100 epochs with a fixed batch size of 24, resulting in approximately 150K updates. We used the Adam optimizer with 4,000 warmup steps to a peak lcarning rale of 2.3x 10-1, followed by inverse square decay. Each time structure was loaded, its coordinates were recentred, randomly rotated, and rescaled.
As shown in Figure 4a adding coordinates substantially improved the model, leading to a final training perplexity of 6.5 with coordinates vs 11.9 without. The final training loss for the version without coordinates (after 100 epochs) was surpassed by the version with coordinates after 8 epochs. Additionally, the final validation loss was surpassed after only epochs which may indicate that the structural features learned early in training are more robust to dissimilarity in sequence space.
We also investigated the difference in sequence recovery rates between the two models. The total sequence recovery rate was 23% for the non-coords model compared to ~38% for the coords model. Figure 4b shows breakdown of the recovery rates per amino acid type. The recovery rate for the coords model was greater than or equal to that of the non-coords model for all amino acid types. The difference was particularly stark for glycine and proline, which may be related to their distinct backbone conformational preferences Ho &
Brasseur 2005 Beck al.] 2008).
3.2.2 Pretrained models learn to measure distance
In the model trained in the previous paragraph, there are three inputs to the model for each token: amino acid type, sequential position, and 3D coordinates. To test if the pretrained model with coordinates was learning to measure distance as predicted, we plotted the average attention paid by each pair of tokens across all heads in a layer as a function of distance. We also plotted the average attention paid to each token as a function of relative sequence distance. To isolate the effect of cach fcature, we fixed all amino acids lo alanine. We also fixed the linear sequence index to constant value for all tokens while measuring 3D dependence and fixed the 3D position to 0) while measuring linear dependence. For the distance measurements, we rounded each pairwise distance to the nearest Angstrom and computed the average for each distance value.
Figure 5 shows the plots for 3D and linear positional dependence for all layers for both models as well as the amplitude and standard deviation of the fit Gaussian for each layer. As expected, the model without coordinates shows no correlation with Euclidean distance. Conversely, the model with coordinates shows a strong dependence in the early layers which is wellaaproximmtted by Gaussian. Both models progressively
7

---

Published in Transactions on Machine Learning Research (07/ (025)
coords Itrani
No coords (train)
0.6
Cooros
With coents (val)
14.
coords (val)
Train perplexity
Recovery rate
0.4
0.2
0.1
20
40
60
80
100
WQMMFYTS C H K
Epoch
Amino acid
(a) Pretraining loss
(b) Sequence recovery breakdown
Figure 4: (a) Masked token prediction loss per epoch. Adding coordinates to the model hugely improves training and validation accuracy, as the model quickly learns to use coordinates to measure distance and compute structural features. (b) Sequence recovery rate for each amino acid type. The height of each bar represents the recovery rate for the coords model and the orange portion is the recovery rate for the non- coords model (the coords model had at least as high a recovery rate as the non-coords model for all amino acid types). The recovery rales for glycine and proline were substantialyy higher in the coords model, which may correspond to identification of beta turns.
widen their field of view as information passes through the layers. This may correspond to system of rea- soning where the model collects information about the local environment before processing this information alongside global patterns. We provide a similar plot without isolating each of the factors in Appendix A.4
3.2.3 Protein function prediction
Finally, we tested whether the pretrained protein model embeddings could improve accuracy on downst task. We trained models to prodict protein molecular function Gene Ontology labels (Ashburner et al. 2000). Protein function prediction has been studied extensively, and has been shown to benefit. from both language and structural features (Gligorijević et al. 2021). We trained models to predict protein function based on the mcan token embedding output by the pretrained models, using the same data splits as DeepFRI (Gligorijević et al. 2021 b. We compare our results to the DeepFRI and DeepCNN versions which were trained on PDB sequences. DeepFRT provided useful comparison since achieved state of the art performance on protein function prediction before models were trained using ESM embeddings. Thus we could pretrain our model the same data DeepFRI and evaluate the specific contributions of model architectures. DeepFRI is ensemble of Graph Convolutional Networks (Kipf & Welling, 2022) with different, propagation rules, in Dehmamy et al. (2019). The graph structure is defined by connecting the k-nearest neighbours for each amino acid in the protein structure. DeepFRI uses a pretrained LSTM model |Graves] 2014) to generate language model embeddings for node features. DeepCNN a Convolutional Neural Network (LeCun et al.) meant to replicate DeepGO (Kulmanov et al. 2018 but retrained on the same sequences as DeepFRT (and our model).
We train two models based on the pretrained models from the previous section. Our first model is based on the simple MIP model from Kulmanov et al. (2024), which is a simple 2 layer MLP block of 1,024 dimensions with a residual connection. The input to the model is a learned linear embedding of the mean token embedding output by the pretrained models.
Our second model is finetuned version of the pretrained masked token prediction models. As BERT Devlin al. 2219] the output a learned linear projection of the final start token embedding. Each model (coords/nn no coords) is finetuned for 20 epochs with constant learning rate of 3x10-5
8

---

Published in Transactions on Machine Learning Research (07 (075)
1e0 attonic
Layer
T wnteniee
Layer 3
Mer sterier
Layer
1 artemiie
T anterrtie
Layer
(residus))
Discanee (resdues)
(resiuae))
(residues)
(residues)
(nisidus))
(a) Linear position, with coordinates
Layer
Layer 2
Layer 3
Layer
Layer
Layer
THEEOT attertoo
WACCUD atterton
attertoon
Poeeaage atteod
(b) position, with coordinates
sttrrknn
Layer
Layer
Layer 3
Layer
Layer
Weers 1 enterntoon
Layer
(resid.as)
(c) Linear position, no coordinates
Layer
Layer 2
Neragge attertion
Layer 3
Aveeege atteettn
Layer
Layer
Avercy afteention
Layer
Distance (4)
Distance
Distance (A)
Distance (A)
Distance
(d) 3D position, coordinates
With coords,
70
With coords, 3D
0.14
With coords, linear
With coords, linear
Attention std deviation
Attention amplitude
6.12
No coords, linear
No coords, linear
0.10
0.08
46
0.06
20
10-
0.00
Layer index
Layer index
(e) Gaussian amplitudes
(f) Gaussian standard deviations
Figure Average attention paid per layer to linear and 3D positional information (a-d). Fit Gaussian functions are shown in blue for each plot. The model which trained with coordinates learns filter heavily by linear and 3D positional information whereas the model trained without coordinates only filters by linear information in early layers. Amplitudes (e) and standard deviations (f) of the fit Gaussians show that both models pay attention to local features early and then gradually widen their fields of view.
The results are shown in Table
Additional results for predicting biological process and cellular component are provided in Appendix A.5. Our sequence-only MLP model compared competitively with DeepCNN and our sequence-structure MLP model compared competitively with DeepFRI. The MLP models are very simple, taking less than minutes train on single GPU, since they are O(1) in sequence length after obtaining the mean sequence embeddings. This indicates that the pretrained models have learned rich embeddings of sequence and structure.
9

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Additionally, the finetuned structure model achieves substantially better AUPRC (0.566 0.446) and max F1 score (0.575 0.460) than DeepFRI. The gap between our structural and non-structural models is much wider than that of DeepFRI, indicating that our model derives a greater benefit from structural information, despite the fact that DoepFRI uses established models for structural processing.
Table 1: GO molecular function prediction results.
(Gain from structure)
(Gain from structure)
Pretraining (# seqs)
Method
Structure
AUPRC
Max Fl
0.427
0.438
DeepFRI (~10M)
DeepFRI*
0.446
0.019
0.460
0.022
None
DeepCNN*
X
0.363
0.385
0.361
0.377
MLP
/
0.460
0.099
0.465
0.088
Ours (~35K)
0.381
0.421
Finetuned
0.566
0.185
0.575
0.154
x From Gligorijević al.
2021
3.2.4 Zero-shot property prediction
We also tested whether the imbued structural information could provide benefit in zero-shot protein property prediction the ProteinGym benchmark (Notin et al. 2023). To normalize against sequence pretraining benefit, we trained a new model with the same architecture as the ESM2-8M model, conditioned on the last layer representation of the pretrained FSM22-MM model. We trained two versions (one with coordinates and one without) on the AlphaFold-DB SwissProt subset of 500k structures (Varadi ct al. 2024h. We compare our results against other sequence/strucluue small models (<20M parameters) Table 2 In contrast to other structural models, our model receives only the Ca coordinates. Fine-grained backbone processing seems to be critical for good performance stability prediction, but our model is otherwise competitive with other methods.
Table 2: Small model zero-shot property prediction on ProteinGym
Spearman
Model inputs
Model
All
Ca Coords
C, N, O Coords
Stability
Sequence
Properties
ProteinMPNN DDuparas al.
2022
0.257
0.565
Mulan
Frolova el al. 2025
0.317
0.369
/
MIF
Yang et al. 2023)
0.383
0.522
/
ESM-2
(8M)* (Lin et al. 2023
0.253
0.399
X
X
Ours (no coords)
0.257
0.350
X
Ours (coords)
0.316
0.426
/
As reproduced by
3.2.5 Awareness of global structure
The ability of a model to reconstruct 3D contact maps is thought to be indicative of its internal understanding of global protein structure (Rao el al. 2019). This is often measured as the model's precision at determining contacts over its K most confident contact predictions (P@K). It has been shown that larger pretrained protein language models have an improved PâL (for proteins of length L), indicating that they have learned more accurate representation of the underlying structure (Lin et al. (023). Similarly, SaProt showed that adding discrete representation of the ground truth structure improved their model's P&L relative
10

---

Published in Transactions on Machine Learning Research (07/2025)
to equivalent ESM2 model (Su et al. 2023 substantially below 100%, which means that much of the global structure was lost in the discretization.
However, the values across all distance scales were still
We trained a contact prediction head on top of our model (Table 3 and found that the accuracy across all distances was nearly 100%. This shows that, unlike discrete tokenizations of structure, adding coordinates preserves global protein structure. Attending to local structure is an essential prerequisite being able aggregate information locally in 3D. Therefore, it may be that a hybrid approach of tineggraine local structural processing followed by a Transformer with coordinates could allow model to reason about structure at multiple distance scales.
Table 3: Contact prediction P@L. Short, medium, and long range contacts are those ween relative positions 6-11, 12-23 and 24+, respectively.
Short
Medium
Long
Model
PâL
PQL./5
P
P@I./5
PåL
P@LL/
ESM-2 (650M)* * Lin al.
2023
44.87
50.18
45.21
53.90
35.33
52.11
SaProt (650M)* *
[Su et al. Ours
2023
57.11
63.71
53.43
66.45
48.14
74.32
95.81
97.61
95.73
98.04
96.98
99.58
*As reported in Su et al. 2023
4 Conclusions
In this work we show that standard Transformers are capable of performing structural reasoning by learning approximately SE(3)-invariant distance filter attention. We predict that even linearly embedded positions can produce Gaussian attention filters of distance and validate this prediction using experiments on simulated points and proteins. The protein model naturally learns to use the 3D coordinates to measure distance which substantially improves its ability to predict masked tokens. The structural information also materially improves the model's ability to inform function prediction, providing even greater benefit than existing custom- built structural models.
We show that Transformers can learn to measure distance and operate hybrid structure/langure models. In contrast to many conventional structure models which are based on GNNs, Transformers do not explicitly model edges. This admits memory-elliciet implementations such as FlashAttention (Dao el al. 2022 Dao, 2023) which allow for fast, fully-connecte updates in linear memory. Most structure models store distance in edges which use quadratic memory for fully-connected graphs. Practically, this means that Transformors can perform structural reasoning on more highly connected structures, which may allow them to "see" more while making decisions.
As shown in Scction 3.2.2 the pretrained protein model trained with coordinates showed strong positional dependence in attention in early layers followed by weak positional dependence in the last few layers. It is possible that this corresponds to the model identifying structural features such as secondary structure and local physics, before encoding these and performing long-range sequential processing, This corresponds with the contemporary trend of preprocessing structural information to create structural tokens for Transformers compared to the more traditional approach of using language model embeddings as input to structural GNNs.
The information required for effective protein structure reasoning depends on the scale of the application. For tasks such as variant prediction and inverse folding, it is essential to have fine-graines view of local structural details such dihedral angles. However, methods such SaProt (Su et al. 2023) which embed local structure into tokens do not retain consistent model of global structure. This type of information is essential for aggregating information in 3D and is required for long range structural processing tasks like protein structure/ compplxx prediction and protein structure generation. In this study, we show that Transformers are adept at embedding model of structure which allows tokens to attend to other local tokens in 3D. Other recent studies such as AlphaFold3 (Abramson et al., (024) and Proteina (Geffner et al., 2024h have demonstrated the utility of this approach on protein complex prediction and structure generation.
11

---

This includes code for Published in Transactions on Machine Learning Research (07 (07/2025)
However, Transformers may lack the ability to (1) represent SE(3)-cquivariant features such side chain atoms and (2) measure specific distance offsets, e-g-, using radial basis encoding. These limitations make Transformers unsuitable for tasks which require a fine-grained atomistic view such as inverse folding. The combination GNNs for local structure processing and Transformers for global structural aggregation may allow for improved structure processing across different distance scales future models.
Finally, while proteins are natural fit for structural Transformers due their combined sequential and spatial data, there are many other possible applications of this type of model. Some of these include tasks with explicit 3D information such small molecules and 3D objects. However, there are also tasks where learning an approximate relationship between entities in Euclidean space could help with reasoning, such as vision Transformers Dosovviskky et al. 2021 or even large language models.
5 Reproducibility
Proofs for claims made in Section 2 are available in the Appendix. Code to replicate the experiments is available on GitHub at: https:/ bectps://gttutt //gtttgtt.../..g./// www.coooooooooo models, data wwwuload/processnngg model training, and evaluation.
Acknowledgments
TE was supported by funding from the Engineering and Physical Sciences Research council [FP /S/24093///1 and Exscientia/ Aussientia/RRecusiion Recursion. MIR acknowledges research funding from CEPI.
12

---

Published in Transactions on Machine Learning Research (07/ (025)
References
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Evans, Chia- Chun Hung, Michael O'Neill, David Reiman, Kathryn Tunyasuvuakooo,, Zachary Wu, Akvilé Zemgulyté, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander I. Cowen-Rivers, Andrew Cowie, Michael Figurnov, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Polapenko, Pascal Savy, Sukhdeepp Singh, Adrian Stecula, Ashok Thillaiuundaram, Catherine Tong, Sergei Yakneen, Ellen D. Zhong, Michal Zielinski, Augustin Zidck, Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M. Jumper. Accurale structure prediction of biomolecular interactions with AlphaFold 3. Nalure, May 2024. ISSN 0028-0836, 1476-4687. doi: 10.1088/$411588-024-0774 841586-024-074877
URL https: http://wwwwat.t..eee/e///t// nature cwwc articess
Michael Ashburner, Catherine A. Ball, Judith A. Blake, David Botstein, Heather Butler, J. Michael Cherry, Allan P. Davis, Kara Dolinski, Selina S. Dwight, Janan T. Eppig, Midori A. Harris, David P. Hill, Laurie Tssel-Tarver, Andrew Kasarskis, Suzanna Lewis, John C. Matese, Joel F. Richardson, Martin Ringwald, Gerald M. Rubin, and Gavin Sherlock. Gene Ontology: tool for the unification of biology. Nature Genetics, 25(1):25 29, May 2000. ISSN 1546-1718. doi: 10.1038/75550 URL https: / http://wwwwwatee www.:natttoooooo wow.amticccs..../////..oo ng0500 _5
Publisher: Nature Publishing Group.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization, July 2016. URL http:
//arxiv..o wwwlwworgg/eoe wttg/bb/////eeeeee..eeeeee 06450 NMXIv:1607.06450 [cs, stat].
David A. C. Beck, Darwin O. V. Alonso, Daigo Inoyama, and Valerie Daggett. The intrinsic confor- mational propensities of the 20 naturally occurring amino acids and reflection of these propensities in proteins. Proceedings of the National Academy of Sciences, 105(34): 12264, August 2008. doi:
10.1073 (0.1073/pnas.dr /ppnas000632770000
URL https://www pnas. ottp/////fl1111.1173////// 1ww3:////......../a.oa/am.oaaamm 10760000700
Pub-
lisher: Proceedings of the National Academy of Sciences.
Tri Dao. PlashAttentionn-2: Faster Atlention with Better Parallelism and Work Partitioning, July 2023. URL
httpaaaawww..rg/b////////// hwtgg/bb//3330 08691
08691 arXiv:2307.0860 .08691 [cs].
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory- Efficient Exact Attention with 10-Awareness, June 2022. URL MiXiv:2205.14133 [cs].
14135.
J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. J. de Haas, N. Bethel, P. J. Y. Leung, T. F. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. K. Bera, N. P. King, and D. Baker. Robust deep lcarning-basel protein sequence design using PrOteinMPNN. Science, 378 378((676)))) 6615):49 56, October 2022. doi: 10.1126/scienne hste:/ecennce.d21111 URL science.
https://www.scieecc.o org/ 11..100.111///1 http://1..12///ccienee add2111 Publisher: American Association for the Advancement of Science.
Nima Dehmamy, Albert-Laszlo Barabasi, and Rose Yu. Understanding the Representation Power of Graph Neural Networks in Learning Graph Topology. In Advonces in Newral Information Processing Systems, volume 32. Curran Associates, Tnc., 2019. URL. htttps//paapers caapperr..iisscoom hea2/79966c122412999811110222220 Aww.reetbbtttttttteaaae
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pro-training of Deep Bidirec- tional Transformers for Language Understanding, May 2019. URL httt:://arxxv.org//bba//// 04805
arXiv:1810.08805 [cs].
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Golly, Jakob Uszkoreit, and Neil Houlsby, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, June 2021. URL
http://arx:: wwtp://agbb//...1/9/99
arXiv:2010.11 1929 (cs).
13

---

Published in Transactions on Machine Learning Research (07/2025)
Daria Frolova, Marina Pak, Anna Litvin, Ilya Sharov, Dmitry Ivankov, and Ivan Oseledets. MULAN: Multimodal Protein Language Model for Sequence and Structure Encoding. Bioinformalies Advances, pp. vbafl17, May 2025. ISSN 2635-0041. doi: ww.1009/bioadv/viie bioadv/vbaf111
URL https://doi.org: htt.i/ii.or//.11.1...
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SE(3)-Transformens: 3D Roto- Translation Equivariant Attention Networks, November 2020. URL BRtLhhttt////xiiii http://raivv...rrg/////////// 10503 arXiv:2006.1055 [cs, stat].
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scaling Flow- based Protein Structure Generative Models. October 2024. URL https://openreviin http://pperrrreeeeeetefoororriia
Vladimir Gligorijević, P. Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C. Taylor, lan M. Fisk, Hera Vlamakis, Ramnik J. Xavier, Rob Knight, Kyunghyun Cho, and Richard Bonneau. Strueture-based protein function prediction using graph convolutional networks. Nature Communications, 3161), May 2021. ISSN 2041-1723. doi: URT.
057888/44467-021--111 Nature Publishing Group.
natura.
Pub-
lisher:
Alex Graves. Generating Sequences With Recurrent Neural Networks, June 2014. URL http://arxiv.orrg bbb//300.com 0850 arXiv:1308.0SI0 [os].
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J. Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q. Tran, Jonathan Deaton, Marius Wiggert, Rohil Badkundri, Irhum Shafkat, Jun Gong, Alexan- der Derry, Raul S. Molina, Neil Thomas, Yousuf Khan, Chetan Mishra, Carolyn Kim, Liam J. Bartie, Matthew Nemeth, Patrick D. Hsu, Tom Sercu, Salvatore Candido, and Alexander Rives. Simulating 500 million years of evolution with language model, July 2024. URL http:: http:///wwwbbooii.o.o.oooooo biorxiv ietorxiv.org//ooee org/ctttttt/ 10. 10.110//2024 07. 01 600583v1
Pages: 2024.07.01 01.005886633
Section: New Results.
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), June 2023. URL http://arxiv. org/abs org/abs/1606.0e 08415 KKiv:1606.084115 [cs) version: 5.
Bosco K. Ho and Robert Brasseur. The Ramachandran plots of glycine and pre-proline. BMC Structural Biology, 5(1):14, August 2005. ISSN 1472-6807. doi: 10.1186/1472-11 (4772-6077-5-1 URL sctps://doi.org. 10. 01186/1472-687-55-1
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives. Learning inverse folding from millions of predicted structures. preprint, Systems Biology, April 2022. URL htttt:t/biooxiiii U.L.hppp//bborrt..oo///oott/ttt wtgp:/oottu//a///..0000000.0...
110..11112/222111 48.777
John Ingraham, Vikas Garg, Rogina Barzilay, and Tommi Jaakkola. Generative Models for Graph- Based Protein Design.
Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
htppscc//ppprrr_ieeeee//eee//// 33444433966554440eccaee3b333 aatp:aaeeceeeaaaaaaaaatttt..... Abstract html
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakoo,, 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 Ziclinski, 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. ///////556-02--0000 2. URL httpp://///w www.:naauue...m http://wwwr...coaaaataaelll/aeaa 03819-2 Number: 7873 Publisher: Nature Publishing Group.
14

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Thomas N. Kipf and Max Welling. Semi-Supervised Classification with Graph Convolutional Networks. July 2022. URL
Maxat Kulmanov, Mohammed Asif Khan, and Robert Hochndorf. DeepGO: predicting protein functions from sequence and interactions using deep ontology-awan aware classifier. Bioinformatics, 34(4):660 668, February 2018. ISSN 1367-4803. doi: (0000///ioiffommtiiee....
URL https:///oi. :ttts:://ol..rr/10.109393
bioinformattics/txx644
Maxat Kulmanov, Francisco J. Guzmán- Vega, Paula Duek Roggli, Lydie Lane, Stefan T. Arold, and Robert Hochndorf. Protein function prediction as approximate semantic entailment. Nature Machine Intelligence, 6(2):220 228, February 2024. ISSN 2522-5839. doi:
URI. https://www Publisher: Nature Publishing Group.
nature. nature.com/arttiii
Yann LeCun, Yoshua Bengio, and others. Convolutional networks for images, speech, and time series. Publisher: Citesoer.
Mingehen Li, Pan Tan, Xinzhu Ma, Bozitao Zhong, Huiqun Yu, Ziyi Zhou, Wanli Ouyang, Bingxin Zhou, Liang Hong, and Yang Tan. ProSST: Protein Language Modeling with Quantized Structure and Disentan-
1.. gled Attention, May 2024. URL https: bixxxo wwwexivv.orgcoooo www.tonntttttntt//////totoottttt 055110/0022222.11.118888777777 Pages: 2
Section: New Results.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs, February 2023. URL
11990 arXiv:2206.11999 [physics].
Yi-Lum Liao, Brandon Wood, Abhishek Das, and Tess Smidt. Equiformer V2: Improved Equivariant Trans- former for Scaling to Higher-Degree Representations, December 2023. URI. http:///taat http://xxvvv.o// www. org/abs/ abs/ 2306.12053 arXiv:2306. 12059 [physics].
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita, Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal- vatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379 1130, March 2023. doi: attps://ww/ http://wwwwwciinccee.e science. wttpe//ooo//oo/1///o.... www.nnncee//ceeeeeeeee ade274 Publisher: American Association for the Advancement of Science.
accenc. URL
Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, Jonathan Frazer, Mafalda Dias, Dinko Franceschi, Yarin Gal, and Debora Marks. ProteinGym: Large-Scale Benchmarks for Protein Fitness Prediction and Design. Advances in Neural Information Processing Systems, 36:64331 64379, December 2023. URL
nips. cttppperr wttp:c/pppeelllee/eeeeeeeeeeeeee files//ppper/2223/maaa AbstraactDaataee Datasets and Benchmarks html
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and
Yun S. Song. Evaluating Protein Transfer Learning with TAPE, June 2019. URL http://arxiv.orrg
abs/1906. 08230 arXiv:1906.0822 [cs].
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118115):e20162331 18, April 2021. doi: 10. 10.1073/pman 0003//pa..0016231111 2016239118. URL https::/ //ww://wwwwwww/ www...ppa..... www.:paaa..oo.......o orw.aai/r//////////////// 201.2222...222222211111
Publisher: Proceedings of the National Academy of Sciences.
Noam Shazeer. GLU Variants Improve Transformer, February 2020. URL http:://txxi: http:://axxiv.oo//bb/////
05202 arXiv:2002.05202 (cs, stat).
Martin Steinegger and Johannes Söding. MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechnology, 35(11):1066-002, November 2017. ISSN 1546-
1696. doi: 10.1038/nbt.3988. URL https://wwrt www.nature...m/rrtilllee///aaaaa Publisher: Nature Publishing Group.
15

---

Published in Transactions on Machine Learning Research (07/2025)
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced Transformer with Rotary Position Embedding, August 2022. URL http://arxiv.. org/ ab0/ 210.00944444444444 arXiv:2104.09864 [cs].
Jin Su, Chenchen Han, Yuyang Zhou, Junjie Shan, Xibin Zhou, and Fajie Yuan. SaProt: Protein Language
Modcling with Structurc-aware Vocabulary. October 2023. URL sttps://ppearerrise.eeee/ooooor
GMRm3G4N1U
Mihaly Varadi, Damian Bertoni, Paulyna Magana, Urmila Paramval, Ivanna Pidruchna, Malarvizhi Radhakr- ishnan, Maxim Tsenkov, Sreenath Nair, Milot Mirdita, Jingi Yeo, Oleg Kovalevskiy, Kathryn Tunyasuvu- nakool, Agata Laydon, Augustin Zidek, Hamish Tomlinson, Dhavanthi Hariharan, Josh Abrahamson, Tim Green, John Jumper, Ewan Birney, Martin Steinegger, Demis Hassabis, and Sameer Velankar. AlphaFold Protein Structure Database in 2024: providing structure coverage for over 214 million protein sequences. Nucleic Acids Research, (2) D375, January 2024. ISSN 0305-1048. doi: 10.1093 nar gkad1011. URL httpps://doi.
10. 1098/narrgkaa11111
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, December 2017. URL http://arxiv org/ algg/bbs17706
03762 arXiv:1706.0376% [cs].
Jesse Vig, Ali Madani, Lav R. Varshney, Caiming Xiong, Richard Socher, and Nazneen Rajani. BERTology Meets Biology: Interpreting Attention in Protein Language Models. October 2020. URL https://
spenreview.net/oooee www.://orruu///dd///ii..vv
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Ticyan Liu. On Layer Normalization in the Transformer Architecture. In Proceedings of the 37th International Conference Machine Learning, pp. 10524-1053) PMLR, November 2020. URL https::/ppoceeddigg mlr. wwtppp/ssse//eppepp.. 1t11199xxxog20bbb html, ISSN: 2640-3498.
Kevin Yang, Niccolò Zanichelli, and Hugh Yeh, Masked inverse folding with sequence transfer for protein representation learning. Protein Enginecring, Design and Selection, 36:gzad015, January 2023. ISSN 1741-0126. doi: 10. 10.1093 00008///rreii//ga0015 URL https://doi http:://h...oooo.....////o/aaaaa
16

---

Published in Transactions on Machine Learning Research (07 (07/2025)
A Appendix
A.1 Coordinates can be rescaled for better approximations
The validity of the approximations shown so far depends on the coordinates being small. Figure A1 shows how well (( ( (2)) approximates quadratic as function of с and how well the resulting exponential approximates a Gaussian. The scaling parameter с can learned by the input and output linear maps of the embedding or by the LayerNorm gain parameters. In this way, all coordinates can be rescaled such that the previous sections produce arbitrarily good approximations.
()))
2
4
capp
X
x
X
Figure A1: Linear and quadratic approximation and resulting Gaussian approximation of cLN (Eisn (=)) as function of e. Target functions are shown black and approximations are shown in blue. Increasing the scaling parameter e results in better approximation without changing the shape of the underlying Gaussian.
A.2
Embedding proofs
A.2.1
Trigonometric embeddings
Consider the embedding Eirigi
Exrig(x) = (cos(x), cis()),)in(2), - sin(x))
(A1)
Then the mean, y E(Erii (x)), is:
1el ((Erig((z))
(A2)
and the variance, af (Firig (x)), is:
17

---

Published in Transactions on Machine Learning Research (07 (07/2025)
( (sin(x)² (- sin(x))² + cos(x)² + (- cos())))
V 2 cos(z)))
(A3)
2
LN (Etrig (x)) Eerig(x) - 4 Barig (z) -0
/2Eerig (x)
(A4)
V2
(Etrigg(Z)) LN (Etrin((2)) V2" (2 cos((I)) cos(22) + (in)))))
4(cos(x1 - x22))
(A5)
~ 4 - 2( x22)
A.2.2 Layer normalization can learn approximately quadratic functions of input Consider the first-order approximation of Egrigs Eis:
Eii(x) = (1,-1, -x)
(A6)
We have
((Eii(z)) =-0
(A7)
((Ein(I)) 1 1 (12 (-1)² + +
V 2 (1 +x)) x2)
2 + x2 24
(A8)
V 2 (1+ x²
1)/ 2 )2
x²
(1 +
/2
2
18

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
1
1
V2
(1+ +
2
V2 (1 7
(A10)
1
V2(x)
2
2
(1 + 4²
)1
2
(A9)
2
2
x2
2211
2
So,
x²
x²
LN -1, -2)) V2((1
):
),
-2)
(A11)
2
2
In this way, layer normalization can be used to generate approximately quadratic functions the input. In particular,
27
LN (E)) LN (N(()))
V2'(2(1 (2(1
(11
+ 2(z10))
2
2
22
4(1
+
+ x112))
2
2
I
(A12)
4( (2 (0)
2 + 2
+
))
2
2.2 2
2
22)² +2+
-1-2
- 2 +4
A.2.3 Gated linear units provide a better approximation
Lemma A.1. ReGLU and SwiGLU can produce functions of x². In particular:
ReGLU(x) + ReGLU(-a) = SSwGGUU()) + SwiGLU x?
Proof.
ReGLU(x) + (eGLU(-z) max(0, x)x + max(0, -r)
max( (22, x)x
(A13)
|
2
Similarly,
19

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
x²
(-x)²
SwiGLU(() + SwiGLU(-x) = 1te2 [te---2)
x1(1cc")
(1+10
(1+
(1+07)
(1++0)+
(A14)
1+0
222 te
2 e² +
2
Theorem A.2. With input ã = (1, x²,,, there crists a linear embedding, E, and a linear map L such that
((L)()(E)) 2,,,,²)// where 1 S o s1+
Proof. Consider the second-order approximation of Errig Equad(z):
x²
x2 (11 - ), -x)
Eqquad (2)) = 2
(A15)
2
We have
11 (E) =
(A16)
x2
22 + (-(1 -- ))2)
x²
a (x)) V (n2 (-x)² (1
V 4 (2+ 2x
x
V2 V 1 +
4
x4
x
(A17)
+
+
V2
4
64
V2 (1 a0
1
22 (1 + ()
8
The error for the LayerNorm-only approximation of x and z² was O(n²) and O(n²)) respectively. In compar- ison, the error given x² as input for æ and z² ((n²) and O(n"). Thus, simple combinations of ReGLU or SwiGLU layers give us a better approximation of d² In practice, this may mean that æ need not be as small for reasonable approximations to hold which may allow for more stable gradients.
x and x?, which in turn gives us a better approximation
20

---

Published in Transactions on Machine Learning Research (07 (07/2025)
A.3 Details structure transformations
Unless otherwise stated, all structures are recentred and randomly rotated each time they are loaded. The recentering ensures translation invariance out of the box, bounds the magnitude of the coordinates, and ensures the start and end tokens always sit the centre of mass. The rotations ensure that the model does not overfit to particular structures, and instead learns the more generalized method of measuring distance.
To recentre, compute the centre of mass for each protein in batch and subtract that value from each coordinate of each protein. To randomly rotate, we compute random rotation per batch with scipy using R. r.ndom () and apply it to each structure in batch during dataloading. This ensures that each time a given structure is loaded will have a random orientation compared to the last time it was loaded.
A.4 Raw attention plots
showed that the attention paid positional and 3D distance well-fit by Gaussians. provide the same plots but without isolating each of the factors. Note that relative distance, position, and amino acid type are all correlated with one another, especially at linear/3D distance 0.
In Section 3.2.2
Here, in Figure A2
A.5
Extended experiments
A.5.1 Protein function prediction
Here, we report the results of two additional experiments on predicting biological process (Table A1] and molecular function (Table A2) labels, also from the DeepFRI dataset. We were unable to compare these re- sults to DeepFRI because the PDB-only results not reported. As in the molccular function experiments, the inclusion of structure improved the performance of all models. In these experiments, the performance of the MLP and finetuned Transformers were comparable, and the sequence-only MT.P outperformed the sequence-only Transformer for cellular component prediction. It is possible that this is the result of the more expressive Transformer model overfitting to sequence training data, which is then mitigated by the inclusion of structure.
Table A1: GO biological process prediction results.
(Gain from structure)
(Gain from structure)
Pretraining (# seqs)
Method
Structure
AUPRC
Max F1
0.197
0.247
MLP
/
0.235
0.038
0.286
0.039
Ours (~35K)
0.191
0.232
Finetuned
0.244
0.053
0.280
0.048
Table A2: GO cellular component prediction results.
(Gain from structure)
(Gain from structure)
Pretraining (# seqs)
Method
Structure
AUPRC
Max F1
0.281
0.335
ML.P
/ 0.306
0.025
0.350
0.015
Ours (~35K)
X 0.230
0.271
Finetuned
0.308
0.078
0.343
0.062
21

---

Published in Transactions on Machine Learning Research (07/202) 2025)
Layer
Layer
Layer
Avee etrentoo
Aeree atvertod
Net Soeeen atwessn
terte Distance (resioues)
Distance (residues)
Distnce (residues)
Distance (residues)
Distance (res
Distance (residues)
(a) Linear position, with coordinates
Avorgg ettetoo
Layer 2
Waogee ettentoo
Layer
SORCO atterbor
Layer 6
(b)
position, with coordinates
Norgge attottnn
Layer
Nedtor attenton
Layer 2
athekto
Layer
weeeeg atrenton
Layer
Layer
Layer 6
- atrent
(c) Linear position, coordinates
Heoage stterdoon
Layer
intemboo
Layer 2
Wengge ettention
Layer
Layer
Layer 5
Sarre entembor
Layer 6
Distance
(A)
(d) 3D position, coordinates
With coords, 3D
17.5
With coords, 3D
With coords, linear
With coords, linear
Attention amplitude
0.4
No coords, linear
Attention std deviation
No coords, linear
30.0
5.0-
0.1
0.8
Layer index
Layer index
Gaussian amplitudes
(f) Gaussian standard deviations
Figure A2: Average attention paid per layer to linear and positional information (a-d) for the unmodified inputs. Fit Gaussian functions are shown in blue for each plot. Amplitudes (e) and standard deviations (f) of the fit Gaussians are shown. The amplitudes are higher and the fits are worse due to the cross-correlation between relative position, 3D distance, and amino acid type.
22

---

Published in Transactions on Machine Learning Research (07/2025)
A.5.2 ProteinGym
We trained two new versions of the model (one with coords, one without) with the same architecture as the ESM2-8M model, conditioned on the final layer of the pretrained ESM2-8M model. We trained on the AlphaFold-DBB SwissProt dataset on chains up 500 residues in length. We retained 0.5% of the dataset for validation the pretraining perplexity. We trained each model for 40 epochs with fixed batch size of 16. We used the Adam optimizer with 4,000 warmup steps to a peak learning rale of 3.6x10-1 followed by inverse square decay.
We performed zero-shot prediction on the ProteinGym zero-shot DMS substitution benchmark. For each prediction, we computed the predicted logits for the wildtype sequence and scored each mutant by laking the negative sum of the mutant logits (equivalent to the normalized log probability). Extended results for all categories shown in Table A3
23

---

Published in Transactions on Machine Learning Research (07 (07/2025)
x Sequence
Model imputs C, N, O Coordinates
Ca Coordinates
Table A3: Small model zero- shot property prediction on ProleinGym
Stability
0. 0.565
0.369
0.522
0.399
0.350
0.426
Organismal Fitness
0.164
0. 0.225
0.300
0.141
0.170
0.224
Expression
0.198
0.383
0.430
0.264
0.279
0.362
Spearman
Binding
0.163
0.326
0.336
0.244
0.244
0.270
Activity
0.197
0.283
0.327
0.219
0.243
0.299
Properties
0.257
0.317
0.383
E-M-2 (8M) 0.253
0.257
0.316
All
Ours (no coords)
ProteinMPNN
Ours (coords)
Model
Mulan MIF
24

---

Published in Transactions on Machine Learning Research (07 (025)
A.6 Model parameter counts
In Table A4 we list the parameter counts for all models used in the paper. The simulated parameter counts will vary slightly as described in the details of each experiment. The pretrained models have the same number of paramcters both with and without coordinates. The finetuned modols have a slightly higher number of parameters than the pretrained because of the final linear layer which projects to the number of classes. The ML.P parameter counts are relatively low because they are conditioned on the (fixed) pretrained embeddings.
Table A4: Parameter counts for all models.
Model name
count
Simulated model
1.6M
Pretrained models
33M
Finetuned models (cc)
33M
Finetuned models (mf)
34M
Finetuned models (bp)
35M
MLP models (cc)
2.2M
MLP models (mf)
2.3M
MLP models (bp)
3.8M
Pretrained models (FSM2)
16M (8M trained)
A.7 Compute resources
All models were trained on a single NVIDTA Quadro RTX 6000 with 24GB of memory. Pretraining of all models (both with and without ESM2 conditioning) took approximately 4 days. Training the MLP models for property prediction took <1 hour and finetuning models for property prediction took <1 day.
25


Reviews:
[{"review_id": "vyVJhmcA7j", "reviewer": "Reviewer_eBF9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- As the capability of transformer-based language models continue to improve, it becomes increasingly important to incorporate knowledge embedded within them into scientific discovery workflows. Models for natural sciences are still largely domain-specific with built-in inductive biases. This paper investigates the applicability of vanilla transformers on structural biology data, making a meaningful step towards the unification of language AI models and biological data modalities. \n\n- The perspective on transformer learning Euclidean distances is well-justified with toy experiments.\n\nWeaknesses:\n\n- The theoretical results are somewhat straightforward. Actionable insights don't seem to present. Afterall, real world tasks for biological sciences are beyond estimating Euclidean distances. how does the theory conclusion of this paper impact model design?\n\n- Experimental results on protein-related tasks are relatively simple and preliminary. It is unsurprising additional positional information improves model performance, but such information isn't always available. This reviewer is not very familiar with the \"GO molecular function prediction\" task and is not sure about how to evaluate the results. Are the baselines state-of-the-art?\n\n- Masked prediction of amino acid from positional information has also been studied in \"inverse folding\". How does a vanilla transformer compare to the existing literature in that space?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The main contribution of this paper is two-fold:\n\n- a theoretical analysis of the relationship between the self-attention mechanism, positional embeddings, and Euclidean distances between the input positions. \n\n- the utility of positional information as positional embedding in a vanilla transformer for protein language modeling, demonstrated through (1) improved performance in masked protein language modeling; (2) the predictive power of a protein language model in terms of amino acid distances; (3) the utility of positional information for protein function prediction.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "It might be reasonable to include inverse folding tasks and baseline models, and see how well can vanilla transformers reason about positional information, this would be a realistic task.\n\nIn the current version I feel the findings are not significant enough due to lack of actionable insights, therefore I am marking \"Audience\" as \"No\".", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "G223vWcm7W", "reviewer": "Reviewer_dq9C", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThis paper provides a theoretical explanation for how Transformer can learn to measure distance and perform structural reasoning, and applied this finding in protein function prediction task.\n\nWeakness:\n\nThe method of this paper is simple, but the finding is interesting to me. However regarding this interesting finding, there are very few experiments done in this paper. The only meaningful one is GO molecular function prediction task. To maker the finding more impactful, the authors should consider broader structural related tasks, such as more protein function prediction tasks like fitness prediction, remote homology detection and protein stability prediction. Additionally, protein design should also be considered, like functional protein sequence design and structure design to fully adapt this finding into real-world problems.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates how Transformer can learn to prediction protein functions based on the squared distance between points. It proved in theory that Transformer can learn Gaussian functions of distance and verified in protein function prediction task that after adding structural information, the prediction accuracy is improved.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "More experiments should be done regarding the finding in this paper.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AbXaQ4I9O1", "reviewer": "Reviewer_7ohZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe paper provides clear and intuitive theoretical explanations for how transformers learn to represent Euclidean distance.\n\nBoth simulated experiments and “real-world” protein experiments support the theory and claims.\n\nStructural reasoning via transformers is understudied and underappreciated, and the work provides important insight into this problem through the mix of theory and experiment.\n\nThe simplicity of the approach and architectures have advantages both for performance and training and inference time.\n\nThe authors find that only n+2 embedding dimensions are needed per head to approximate distances in dimension n, showing that compact representations capture spatial distances. \n\nThe paper is well-written and clear.\n\nWeaknesses:\nThe theory section makes some strong simplifying assumptions (e.g., small coordinates, trivial Q/K mappings, 1 spatial dimension), which might limit the generalizability. The authors do a good job acknowledging these limitations.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper the authors investigate the capabilities of transformers to process structural information (3D coordinates) using simple linear embeddings. Theory and experiments show that this is indeed possible, and that transformers filter attention as a 3D Gaussian. Further downstream experiments for protein masked language modeling and function prediction show that this approach is competitive with or better than more traditional GNN-based approaches for modeling problems that benefit from explicit structural information. The work challenges the necessity of explicit structural modeling and shows that transformer encoders learn spatial relationships.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Can the authors expand on the discussion of the simplifying assumptions? In particular, it’s not clear how restrictive the trivial Q/K mapping is and how relaxing that assumption would affect the results.\n\nHow are the random rotations applied? Can the authors expand on this and add details to the appendix?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)