Transformers trained on proteins can learn to attend to Euclidean distance

TL;DR

This paper demonstrates that standard Transformer models trained on protein sequences can independently learn to attend to Euclidean distances in 3D space, enhancing protein structure modeling without specialized structural architectures.

Contribution

It provides a theoretical framework and empirical validation that Transformers can learn structural information directly from coordinate embeddings, enabling hybrid structure-language modeling.

Findings

Transformers can learn to filter attention as a 3D Gaussian with learned variance.

Pre-training on structural data improves downstream protein modeling tasks.

Transformers outperform custom structural models when trained with structural information.

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 a downstream task, yielding better performance than custom structural models. Together, this work provides a basis for using standard Transformers as hybrid structure-language models.