Can We Make Protein-Structure Prediction More Efficient with Linear Attention?

A protein is a chain of amino acids that folds into a three-dimensional shape, and its function depends on that shape. Predicting the shape from the sequence alone was a central problem in computational biology for decades. AlphaFold2 showed that it can be solved to near-experimental accuracy for many proteins [4], using multiple-sequence alignments, structural templates and a large attention-based network.

Those models are expensive. Part of the cost comes from attention itself: every residue attends to every other residue, so the cost grows with the square of the chain length. Long proteins and proteome-scale screening are where this hurts. Our ICTA 2025 paper [1] asked a modest version of the question: can a transformer with linear-cost attention learn useful geometric quantities from sequence and evolutionary features?

Why long sequences are expensive for attention

Standard attention [6] compares every query with every key:

For a chain of \(N\) residues, each head in each layer materialises an \(N\times N\) score matrix. At \(N = 256\) that is about 65,000 entries. At \(N = 2{,}048\) it is over four million, per head, per layer, before backpropagation stores it again. Time grows the same way. The quadratic term, not the number of parameters, is what limits how long a sequence fits on a GPU.

The Performer idea: attention without the \(N \times N\) matrix

The Performer [2] rests on one observation. The softmax kernel \(\exp(q^{\top}k)\) can be written, in expectation, as an inner product of random feature maps: \(\exp(q^{\top}k)\approx\varphi(q)^{\top}\varphi(k)\) with \(\varphi:\mathbb{R}^{d}\to\mathbb{R}^{m}\). Once the kernel factorises, matrix multiplication can be re-associated:

Computing \(\varphi(K)^{\top}V\) first gives an \(m\times d\) matrix that summarises all keys and values, and every query then reads from it. No \(N\times N\) object is ever formed, so time and memory scale as \(O(N\,m\,d)\) instead of \(O(N^2 d)\). The FAVOR+ construction uses positive, orthogonal random features, which keep the estimate unbiased and its variance low. The price is approximation error, controlled by \(m\).

Putting evolution into the input: PSSMs

A single sequence says little about which positions matter. Evolution says more. If a position is conserved across many related proteins, it is probably structurally or functionally important. A position-specific scoring matrix (PSSM) records this: for each position in the sequence, it gives 20 scores, one per amino acid, describing how often that residue appears there across an alignment of homologous sequences. PSSMs were the input that made neural secondary-structure prediction work in the late 1990s [5], and they remain a strong, cheap source of evolutionary signal compared with full alignments.

Our model represents each residue as its one-hot amino-acid identity concatenated with a standardized 20-dimensional PSSM row, projected into a 256-dimensional embedding.

The model and its two targets

The two heads are trained jointly. Distances use a mean-squared error over all residue pairs, and angles use a cosine loss so that \(-179^\circ\) and \(+179^\circ\) count as close:

It is worth being precise about what these targets are. A matrix of predicted distances and a list of predicted angles are not a 3-D structure. Turning them into coordinates needs a separate reconstruction step, as in trRosetta [8], and the result must then be compared with the true structure using metrics such as TM-score or lDDT. Our paper does neither. At best, it can show that distances and angles were learned.

What the experiment shows

What we observed. Both losses drop sharply in the first two epochs and then flatten. Validation loss fluctuates between about 2.18 and 2.32 without trending upward, so there is no sign of overfitting. The figure shows 15 epochs; the paper’s text describes 20.

How we read it. A smooth loss curve shows that optimisation was stable. It does not show that the model learned distances. The reported distance MSE is 38,820 Å², while the plotted total loss is around 2. So the plotted loss must have been computed on rescaled targets, or weighted in a way the paper does not describe. Either way, this curve cannot be read as distance accuracy improving.

This is the central question, and the most direct evidence is the distance maps. The paper shows two test proteins; here is the first.

What we observed. Three things stand out once the colour bars are read alongside the images.

• The predictions are essentially zero. The centre panel’s scale runs from about −0.06 to +0.07 Å; the second protein in the paper is similar (about 0 to 0.16 Å). Negative distances are physically impossible. As a result, the error map on the right is visually indistinguishable from the reference on the left.
• The reference map varies only with sequence separation. We sampled the left panel on the residue grid. Along each diagonal (fixed \(|i-j|\)), the colour changes by less than 1/255, so the reference distance depends on \(|i-j|\) alone. It rises at about 3.7 Å per residue, close to the 3.8 Å spacing between consecutive Cα atoms. That is the profile of a fully extended, straight chain. A folded 64-residue domain spans a few tens of ångströms and shows off-diagonal contacts between residues far apart in sequence.
• The headline error matches a zero prediction. The paper reports RMSE 197.03 Å and MAE 156.89 Å. For straight-chain targets \(d_{ij}=3.8\,|i-j|\) and a prediction of zero, a chain of 128 residues gives RMSE ≈ 199 Å and MAE ≈ 162 Å.

How we read it. In this experiment, the distance head did not learn inter-residue distances: its outputs sit near zero regardless of the target. The reference maps do not look like folded proteins either, which points to a problem in how distance targets were built from the data. We have not yet traced where in the pipeline it arises. ProteinNet stores coordinates in picometres, so units are one place to look, but a unit error alone would not remove the off-diagonal structure. The paper describes these maps as showing accurate local predictions and long-range consistency, and the maps do not support that description.

The paper reports an angle RMSE of 4.07° and an MAE of 3.28°. On their own, these numbers would be remarkable. Two things keep us from relying on them. First, given the problem with the distance targets, the angle targets also need checking before the numbers mean anything. Second, the paper does not say which dihedrals were predicted. The backbone \(\omega\) angle sits near 180° in almost every residue, so including it would make the average error look much smaller than for the informative \(\varphi\) and \(\psi\) angles. We report the numbers as published, but they should not yet be read as evidence of accurate angle prediction.

In principle, yes, and this is the part of the project I still find most promising. Figure 1 shows the size of the problem: at 2,048 residues, a softmax attention map is about eight times larger than a FAVOR+ factorisation with \(m = 256\). The gap grows linearly with length after that.

The paper also lists measured-sounding gains: 3–5× less memory for sequences above 256 residues, 2–4× faster training, and support for 1,024-residue sequences on a single GPU. It gives no baseline configuration, hardware or measurement procedure, and the model is configured for sequences up to 512 residues. We therefore treat these figures as unverified and keep to the theoretical argument. A real efficiency claim needs a softmax baseline of the same size on the same hardware, with peak memory and throughput reported as a function of sequence length.

What a sound version of the experiment would need

• Look at the targets before training on them. Plot a few true distance maps and check that they show contacts. Convert ProteinNet coordinates from picometres to ångströms. Mask missing residues and padding.
• Compare with trivial baselines. Predicting the average distance for each sequence separation \(|i-j|\) is a strong, cheap baseline. A model that does not beat it has not learned structure.
• Use the standard formulation. Predict binned distance distributions (distograms) with cross-entropy, as AlphaFold 1 and trRosetta did [7][8], rather than raw distances with MSE. Evaluate long-range contact precision.
• Specify the angles, report \(\varphi\) and \(\psi\) separately, and compare with a baseline that predicts each residue type’s most common angle.
• Measure efficiency properly, against a matched softmax transformer, across sequence lengths.
• Reconstruct coordinates before making any claim about structure, and score them with TM-score or lDDT.

What we learned

1. The efficiency case for linear attention is sound in theory. FAVOR+ removes the \(N\times N\) matrix and makes long proteins tractable. Whether the approximation preserves enough accuracy for structural targets remains to be tested.
2. A smooth loss curve is not evidence of learning. Inspecting predicted outputs against simple baselines is the check that matters, and in this experiment it shows that the distance head did not learn.
3. Distances and angles are not structures. Claims about structure prediction need reconstructed coordinates and structural metrics.
4. Negative results deserve a clear write-up. This one points directly at what to fix: the target pipeline, the baselines and the efficiency measurements.

Limitations and open questions

• Where in the data pipeline do the straight-chain reference maps come from, and do the angle targets have the same problem?
• With correct targets, can a Performer with PSSM inputs beat the per-separation baseline, and how does it compare with a softmax transformer of the same size?
• How many random features does FAVOR+ need before the approximation stops limiting accuracy on structural targets?
• Do the efficiency gains materialise in practice at the sequence lengths where they matter?

@inproceedings{roufas2026performer,
  title     = {Efficient Protein Folding with Transformer Models
               Using the Performer Architecture},
  author    = {Roufas, Nikolaos and Karamitsos, Ioannis and
               Al-Hussaeni, Khalil and Gerogiannis, Vassilis C. and
               Kanavos, Andreas},
  booktitle = {Advances in Information and Communication Technology
               (ICTA 2025)},
  series    = {Lecture Notes in Networks and Systems},
  volume    = {1833},
  pages     = {191--198},
  publisher = {Springer, Cham},
  year      = {2026},
  doi       = {10.1007/978-3-032-18162-6_20}
}