Structure targets for small protein language models

Personal research project, October 2026.

Protein language models such as ESM-2 are pretrained with masked language modelling (MLM): mask a fraction of residues and predict their identities. The resulting representations encode a fair amount of structure, but only indirectly, since the objective rewards any sequence statistic that helps recover the residue.

This study replaces the residue target with a structural one. Structure is used only as a training target, never as an input, so the trained encoder still runs on sequence alone. The question is whether that produces better structural features at a fixed sequence budget, and whether it matters if the targets are predicted through a JEPA-style predictor or regressed directly.

The interactive demo runs two of the trained models on real proteins.

Setup

All models share one network: a 7.4M-parameter transformer encoder in ESM-2's smallest configuration, randomly initialised. All are trained on the same 22,839 AlphaFold Database chains (8.9M residues) for 7,630 steps, about 486M tokens. Each step masks 20% of positions with BERT-style 80/10/10 corruption. Only the target at masked positions differs:

	Target at each masked position
C1	Original residue (MLM, the standard objective).
C2	Encoder's own embedding from an EMA copy of itself, through a predictor (JEPA). No anti-collapse regulariser.
C3	ESM-IF1 embedding of the residue, through a 2-layer predictor.
C5b	The same ESM-IF1 embedding, through a single linear layer.

ESM-IF1 is an inverse-folding model. Its encoder takes backbone coordinates only and outputs a 512-d embedding per residue describing local geometry and fold context. The targets are reduced to 256-d with PCA. Because ESM-IF1 never reads residue identities, the targets cannot leak the masked residue.

C3 is the latent-prediction setup used by I-JEPA in vision. C5b is the same targets with plain regression.

C1, C2 and C3 have three seeds each and C5b has two, for eleven runs in total.

Evaluation

Each encoder is frozen and evaluated on three tasks:

• SS3: per-residue three-class secondary structure (helix, strand, coil). This is what the demo shows.
• SS8: the eight-class version.
• Fold: zero-shot SCOPe fold retrieval. Does a domain's nearest neighbour share its fold?

The SS probes are trained on the NetSurfP-2.0 set, selected on TS115 and tested on CB513. The probe is chosen from one fixed grid of three linear and three one-hidden-layer heads, with no per-model tuning. Fold retrieval has no training step: per-residue embeddings are mean-pooled, and 1,031 query domains are ranked against a 14,146-domain gallery by cosine similarity, with no superfamily on both sides. RankMe (effective rank of the encoder output) is tracked as a collapse diagnostic.

Baselines are the public ESM-2 8M model (same scale, far more training data), one-hot residue features, and the randomly initialised encoder.

Hypotheses and pass criteria were fixed before the training runs. The paper lists the deviations from the original plan.

Results

	runs	SS3	SS8	Fold
C1, MLM	3	0.688 ± 0.003	0.542 ± 0.005	0.022 ± 0.004
C2, EMA latent	3	0.492 ± 0.006	0.358 ± 0.004	0.012 ± 0.003
C3, ESM-IF1 via predictor	3	0.612 ± 0.059	0.472 ± 0.049	0.020 ± 0.002
C5b, ESM-IF1 via linear head	2	0.709 ± 0.002	0.565 ± 0.005	0.032 ± 0.002
ESM-2 8M (public)		0.756	0.611	0.095
One-hot		0.487	0.348	0.016
Random-init encoder		0.514	0.371	0.016

SS3 and SS8 are per-residue accuracy. Fold is top-1 retrieval accuracy. Raw ESM-IF1 embeddings computed from the real structures score 0.161 on Fold.

C2 collapsed. With a jointly learned target and no regulariser, a constant embedding minimises the loss. C2's final RankMe was 37.7 against 187.3 for C3, and it scored below the random-init encoder on every task. This was the expected outcome. C2 is a positive control showing the failure mode exists in this setup.

	RankMe (encoder width 320)
C1, MLM	179.3 ± 7.2
C2, EMA latent	37.7 ± 0.9
C3, ESM-IF1 via predictor	187.3 ± 6.3
C5b, ESM-IF1 via linear head	214.4 ± 2.3

Frozen targets rule out the constant solution, and neither C3 nor C5b collapsed. The PCA-256 target bank itself has a RankMe of 229.8.

C3 did not beat MLM. This was the main hypothesis (H2), and it failed. C3 was below C1 on all three tasks. It was also unstable across seeds: SS3 ranged from 0.549 to 0.666, a standard deviation of 0.059 against 0.003 for C1.

C5b did. C5b beat both C1 and C3 on all three tasks. The only difference from C3 is a linear head in place of the predictor. A likely explanation, not tested here: the predictor can take on part of the mapping to the target space, so the encoder output, which is what the probes read, carries less of it. A linear head forces the encoder output to be linearly related to the targets.

C5b has only two seeds, but they agree closely (SS3 0.707 and 0.711), and its gap to C1 is several times the seed-to-seed standard deviation.

Limitations

ESM-2 8M beats every model here. It was trained on far more sequence data, so the comparison only shows how far this budget is from saturating the architecture. The question this study answers is narrower: at a fixed sequence budget, which target gives the best features?

C3 and C5b get supervision from a 142M-parameter structure model trained on far more data than C1 sees. Matching the token budget does not remove that asymmetry. The result is about how to spend a sequence-model budget when a pretrained structure encoder is available, not about which objective is better at equal information.

With three seeds per condition, per-task p-values are not reported. Secondary structure was scored on a fixed subset of residues because of a bug in parsing the benchmark's residue mask. The subset is about half of each protein in training and validation. On the CB513 test set it is 56k residues, against 84k under the standard mask, and some of those 56k are residues the benchmark normally excludes. All models use the same residues, so comparisons between them hold, but the absolute numbers are not comparable with published CB513 results.

Demo

The demo runs C5b and C1 in the browser. Predictions for the eight preset proteins are precomputed. A pasted sequence triggers a download of both models (ONNX, 8-bit weights, about 8 MB each), which then run client-side.

The study's SS probes were not saved, so the demo's probes were re-fitted with the study's settings on a quarter of the probe training proteins, to fit in laptop memory. On the study benchmark they score 0.708 (C5b) and 0.673 (C1), close to the study figures. For presets, the 3D view uses the deposited PDB structure and the benchmark's DSSP labels.

Full method and results are in the paper.