The Effectiveness and Durability of Pretraining Data Filtering in Protein Language Models
Authors: Aaron Maiwald, Roderick McNeill, Jaeyoung Lee, Matthew A. Clarke, Lennart Justen, Alejandro Acelas, Ivo Andrews, Giles Harper-Donnelly, Adam Moore, Sophie Rose, Oliver M. Crook, Felix Jackson
Organizations: Department of Statistics, University of Oxford · SecureBio · UK AI Security Institute · Independent · Broad Institute of MIT and Harvard · Department of Chemistry, Kavli Institute for Nanoscience Discovery, University of Oxford
Pretraining filtering is one of the most common safeguards for open-weight language models. Here we study its effectiveness, durability and scalability using the ESM-2 family of protein language models. We find that filtering viral proteins from the pretraining data meaningfully degrades masked language modelling, mutation effect prediction and long-range contact prediction on viral proteins, while largely preserving performance on non-viral proteins. These effects are concentrated among proteins from human-infecting viruses, and grow as we scale models from 8M to 3B parameters. As models increase in size, we find that stringent filtering becomes more effective. However, fine-tuning on sequences from the removed viral families restores the lost performance. Although the cost to recover performance increases with model scale, even for the largest models it costs about 80 USD of H100 time. This work establishes pretraining filtering as moderately effective at selectively degrading performance on viral proteins, but not durable against fine-tuning with modest computational resources and access to relevant data.
Figures & tables
Figure 1: Filtering viral proteins from the pretraining data selectively degrades performance on tasks related to viral proteins, and more so for larger models. The cost of performance recovery increases with model size. All panels show the tier-0 filtered model, where we exclude all viral proteins and their homologues from the training data, against the unfiltered model of the same size. Left: Comparison of language-modelling performance on held-out viral and non-viral sequences. Middle: Comparison of mutation effect prediction performance for viral and non-viral sequences, on ProteinGym DMS datasets. Intervals represent one standard error of paired differences across sequences (left; n=5,161 viral, 3,000 non-viral) or assays (middle; n=31 viral, 186 non-viral). Right: Estimated fine-tuning cost to recover performance on viral language modelling (solid) or viral mutation effect prediction (mean performance across 45 EVEREST assays (dashed)). Markers show the first recovering checkpoint; bars extend to the preceding checkpoint, and downward arrows mark upper bounds.
Figure 2: We exclude nested sets of viral proteins and homologues from the pretraining data. Bars show the number of UniRef50 entries each tier removes from pretraining, on a logarithmic axis. Coloured segments identify shells S1-S5: these are entries that are in one tier but not the next more narrow tier (except S5, which contains influenza A). The hatched segment marks additional non-viral homologues removed by tier 0. We train models on tiers 0 to 3 and evaluate on held-out samples from each of the five viral shells.
Figure 3: The degradation of the models’ performance on viral proteins due to filtering grows with scale, while leaving non-viral performance mostly unchanged. All panels show capability evaluations, each with the forget domain vertical and the retain domain horizontal. Each line connects models trained with the same filtering strategy at different scales. The two widest filters, tier 0 and tier 1, are drawn next to the unfiltered model.
Figure 4: The filtering effect is concentrated among proteins from human-infecting viruses. Mean cross-entropy on each held-out shell against model size, for the unfiltered model and all four filters. Lower is better. n is the number of held-out entries in that shell ( Figure 2 ). The gap between the filtered and unfiltered models is largest in the human-infecting shell and generally increases with model size. No tier-2 or tier-3 model was trained at 3B.
Figure 5: Filtering lowers rank agreement on viral assays and leaves non-viral assays unchanged. One point per ProteinGym assay at 3B: mean Spearman ρ under the tier-1 model against the unfiltered model. Points below the dashed line were degraded. Viral assays are orange; the same 186 non-viral assays are repeated in blue in every panel. The large panel pools all 31 viral assays, the small ones show one family each. Most of the pooled shift comes from Retroviridae.
Figure 6: Continued training on the withheld corpus recovers performance within the first 1.4 billion tokens at every size. Both rows use the tier-0 model. Left: the gap to the same-size unfiltered model, with the same zero line for every size. Right: the fine-tuning compute needed for full recovery, in exaFLOPs, one exaFLOP being 1018 floating-point operations. We report compute rather than tokens because a 3B model spends 4.1 times the FLOPs per token that a 650M model does. We score models every 500 steps. If a bar reaches zero, recovery finished before the first checkpoint. The right-hand scale converts compute to dollars at 34% utilisation and 4 USD per H100-hour.
Figure 7: While the retain performance stays nearly constant, removing more of the corpus leads to more degradation of forget performance for our largest models. All panels plot the unfiltered model’s cross-entropy minus the filtered model’s. The ticks on the x-axis correspond to our filtering tiers 3 (leftmost) to 0 (rightmost).
Appendix figures & tables27 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Distribution of viral families across shells. Each bar is a shell, and its segments show how that shell’s entries are distributed across viral families. This complements the tier-level view in Figure 2 .
Dataset
Removes
Sequences
Removed
% Removed
UniRef50 2021_04
—
49,874,565
—
—
full
—
49,507,711
0
0
tier0 filtered
tier 1 + homologues
47,154,140
2,363,337
4.77
tier1 filtered
all viruses
49,082,570
434,907
0.878
tier2 filtered
eukaryote-infecting families
49,386,974
130,503
0.264
tier3 filtered
human-infecting families
49,461,548
55,929
0.113
Appendix
Table 1: The pretraining corpora. Each filtered dataset is the unfiltered UniRef50 corpus with that tier’s entries removed. The tiers are nested, so a wider tier removes everything a narrower one does. The full dataset underwent an additional filter, removing 9,766 additional evaluation homologues after the tier-filtered datasets were constructed.
Figure 9: Our models have similar mutation effect prediction performance to ESM-2 and higher cross-entropy at 650M and 3B parameters. We compare each unfiltered model with the corresponding facebook/esm2 checkpoint. Left: mean cross-entropy over 3,000 non-viral held-out sequences. Right: mean Spearman correlation over 217 ProteinGym substitution assays. These held-out sequences are excluded from our training corpus. ESM-2’s training splits are unpublished, so overlap with these sequences is unknown.
Training corpus
Size
Full
Tier 0
Tier 1
Tier 2
Tier 3
8M
✓
✓
✓
✓
✓
35M
✓
✓
✓
✓
✓
150M
✓
✓
✓
✓
✓
650M
✓
✓
✓
✓
✓
3B
✓
✓
✓
—
—
Appendix
Table 2: Pretrained models. A check marks a model size-corpus pair for which a model was pretrained and evaluated. The corpora are the UniRef50 pretraining corpus and the nested tier exclusions of Figure 2 .
Pretraining
Fine-tuning
Steps
150,000
5,000
Effective batch
221 tokens
218 tokens
Total tokens
≈ 315B
≈ 1.31B
Peak learning rate
4×10−4
5×10−5
Warmup steps
2,000
100
Decay horizon
500,000 steps
5,000 steps
Appendix
Table 3: Pretraining and fine-tuning hyperparameters. Values are shared across model sizes except where noted. Effective batchsizes are powers of two: 221=2,097,152 and 218=262,144 tokens. The pretraining schedule decays linearly towards a floor of one tenth of the peak learning rate at 500,000 steps, but the run stops at 150,000 steps, where the learning rate is 2.93×10−4 . The fine-tuning schedule reaches its floor at the final step.
Figure 10: At 650M, tier-2 and tier-3 filtering increase mean cross-entropy more for removed than retained viral families. Change in cross-entropy, from the unfiltered model, for 650M models after tier 2 and 3 filtering, split by viral families removed by this filter vs those retained. Top: Mean difference in cross-entropy (filtered minus unfiltered), shown separately for removed and retained families. Each point represents the difference in mean cross-entropy for the held-out sequences for this viral family. Counts in brackets give the number of families. Bars represent sequence-weighted group means. Bottom: Unfiltered versus filtered cross-entropy for retained and removed viral families. The diagonal represents equality.
Figure 11: The largest estimated cross-entropy increases among retained viral families occur in families with few held-out sequences. Cross-entropy changes for 650M models under tier-2 and tier-3 filtering. Left: Scatter plot showing the change in mean cross-entropy for retained viral families against the number of held-out sequences, highlighting that the largest cross-entropy changes in retained families occur in families with few held-out sequences. Right: Points show the fraction of families that show degradation ( Δ>0.02 nats) separately for each tier. Degradation is more frequent among removed than retained families. Error bars show 95% Wilson intervals, and labels show degraded/total family counts.
Figure 12: Filtering effect on contact prediction in short-, medium-, and long-range residue-residue pairs for the 650M model . Each row compares P@L, P@L/2, and P@L/5 across four filtering levels, from Tier 3 to Tier 0, against the unfiltered model. The upper panels show mean precision for viral and non-viral sequences before and after filtering; the lower panels show each domain’s change (filtered minus unfiltered) and the non-viral minus viral change gap, in percentage points.
Figure 13: fine-tuning recovers contact prediction performance across short-, medium-, and long-range residue-residue pairs in the 650M model . Each row compares P@L, P@L/2, and P@L/5 for filtered checkpoints from Tier 3 to Tier 0 with the corresponding models after viral-only fine-tuning (see Section A.3 ). The upper panels show mean precision for viral and non-viral sequences before and after fine-tuning; the lower panels show each domain’s change (recovered minus filtered) and the viral minus non-viral change gap, in percentage points.
Figure 14: Viral contact-prediction precision across model sizes before filtering, after tier-1 filtering, and after viral-only fine-tuning . Rows show short-, medium-, and long-range contacts; columns show P@L, P@L/2, and P@L/5. Each panel compares the mean precision of the unfiltered, tier-1 filtered, and recovered models on viral proteins.
Figure 15: Scaling of the filtering effect in short-, medium-, and long-range pairs . Each row compares unfiltered and tier-1 filtered models from 8M to 3B using P@L, P@L/2, and P@L/5. The upper panels show mean precision for viral and non-viral sequences; the lower panels show each domain’s change (filtered minus unfiltered) and the non-viral minus viral change gap, in percentage points. The gap is larger for medium- than short-range pairs at 650M and 3B, and larger for long- than medium-range pairs at 150M–3B, across all three metrics.
Figure 16: Scaling of the recovery effect in short-, medium-, and long-range pairs . Each row compares tier-1 filtered models from 8M to 3B with their counterparts after viral-only fine-tuning using P@L, P@L/2, and P@L/5. The upper panels show mean precision for viral and non-viral sequences before and after fine-tuning; the lower panels show each domain’s change (recovered minus filtered) and the viral minus non-viral change gap, in percentage points. For P@L, the gap increases from short to medium to long range at every model size.
Figure 17: Filtering differentially degrades language modelling performance across proteins from different viral families. Each plot shows mean cross-entropy on held-out entries for the eight families with the largest tier-1 effect at 3B (among families with at least 20 entries). Each curve shows the unfiltered and tier 0 to tier 3 models for each model size up to 650M, and tier 0 and 1 for 3B.
Cross-entropy, Tier 1
ProteinGym, Tier 1
Tier 0 − Tier 1
Size
Δ (nats)
p
Δρ
p
Δ (nats)
p
8M
+0.006
4×10−35
+0.001
0.844
+0.0005
0.330
35M
+0.019
<10−99
−0.002
0.816
+0.0007
0.269
150M
+0.074
4×10−81
−0.019
0.264
+0.0121
9×10−42
650M
+0.125
<10−99
−0.068
0.003
+0.0362
<10−99
3B
+0.173
<10−99
−0.124
3×10−4
+0.0855
<10−99
Appendix
Table 4: Effects of filtering on masked language-modelling and mutation effect prediction by model size. Left: Mean paired differences between tier 1 filtered and unfiltered models in cross-entropy (nats, 5,161 held-out viral sequences). Positive values indicate that the tier 1 filtered model performs better. Middle: Change in mean Spearman correlation for mutation-effect prediction (Spearman ρ , 31 assays) for tier 1 - unfiltered. Negative values indicate that the tier 1 filtered model performs worse. Right: Difference in cross-entropy (nats) between tier 0 and tier 1 filtered models on paired sequences. p -values are calculated from two-sided paired t -tests, without adjustment for multiple comparisons.
Tier 1 vs unfiltered
Tier 0 − Tier 3
Shell
n
Δ (nats)
Share
Δ (nats)
p
1 – phages, unclassified
1462
+0.029
6.5%
+0.068
5×10−61
2 – eukaryote-infecting
1472
+0.027
6.2%
+0.038
1×10−40
3 – human-infecting
1461
+0.361
81.6%
+0.098
<10−99
4 – pathogens of concern
493
+0.070
5.3%
+0.015
2×10−5
5 – influenza A
273
+0.008
0.3%
+0.007
0.003
Appendix
Table 5: Per-shell effects of filtering at 650M Left: mean paired cross-entropy differences (tier 1 minus unfiltered, nats) and the share for each shell of the summed difference across all five viral shells. Shares are weighted by the number of held-out sequences in each shell. Right: mean paired cross-entropy differences (tier 0 minus tier 3, nats). Both comparisons pair observations by sequence. p -values are from two-sided paired t -tests, without adjustment for multiple comparisons.
Figure 18: Viral cross-entropy decreases during fine-tuning, while non-viral cross-entropy increases over the evaluated checkpoints at 650M and 3B. Mean cross-entropy (nats) plotted against fine-tuning tokens (billions) for tier-0-filtered models trained on the withheld viral corpus. Panels show language modelling performance on the five viral shells and the non-viral control. The dotted horizontal lines show the performance of the unfiltered baseline at each model size. Lower cross-entropy indicates better performance. Non-viral cross-entropy increases over the evaluated checkpoints at 650M and 3B.
Figure 19: Mean cross-entropy on the human-infecting shell recovers within 1.4B fine-tuning tokens for every evaluated model size and filtering tier. Recovery of mean cross-entropy on the five viral shells and mean Spearman correlation on EVEREST during viral-only fine-tuning. Exclusion tiers run from tier 3 (narrowest) to tier 0 (widest). Recovery is defined as matching or improving on the mean performance of the same-size unfiltered model. Vertical bars show the bracket of recovery between scored checkpoints, where upper endpoints give the first checkpoint meeting this criterion. Intervals starting at zero indicate recovery before the first scored fine-tuning checkpoint. The dotted line marks 500 fine-tuning steps (approximately 0.131 billion tokens). All available size and filter combinations recover on the human-infecting shell within this fine-tuning budget. No tier-2 or tier-3 filtered model was trained at 3B.
Figure 20: The increase in viral cross-entropy from tier-1 filtering grows with model size. Mean held-out cross-entropy (nats) for unfiltered and tier-1-filtered models from 8M to 3B parameters, evaluated on non-viral (left) and viral (right) sequences. Lower cross-entropy indicates better performance. Viral cross-entropy decreases with model size under both training conditions, with a larger reduction for the unfiltered model. The filtered-minus-unfiltered performance difference increases from 0.006 nats at 8M to 0.173 nats at 3B. Differences in performance on non-viral remain within ±0.008 nats at all five sizes. Comparisons across filtering tiers are shown in Figure 7 .
Figure 21: Tier-0 degradation is negatively correlated with unfiltered cross-entropy at 650M and 3B. Association between unfiltered cross-entropy and the change in cross-entropy after filtering across viral families. Left: Spearman correlation across the 30 families (including unclassified) with at least 20 held-out entries. When performing the statistical analysis, each family’s entries are split into two disjoint halves with the same split used across model sizes and filtering tiers. One half estimates unfiltered cross-entropy and the other estimates the paired filtered-minus-unfiltered difference, so each estimate is calculated from a separate set of held-out sequences. Bars are 95% percentile intervals from 10,000 family bootstrap resamples. Open markers indicate p≥0.05 in a two-sided test with 10,000 permutations of family pairings. Intervals and tests are conditional on the fixed sequence split, and p -values are not adjusted for multiple comparisons. Lower cross-entropy is better; therefore, a negative correlation suggests that greater degradation due to filtering occurs in families where the unfiltered performance is higher. Right: filtered versus unfiltered cross-entropy at 3B, using all held-out entries in each family. Each point is a viral family. Dashed line marks equal performance and vertical bars show the paired standard error of the difference.
Figure 22: Mean Spearman correlation on viral assays increases more from 150M to 3B in unfiltered models than in tier 1 filtered models. Mean per-assay Spearman ρ between predicted and measured mutation effects, for unfiltered and tier 1 filtered models. Left: 186 non-viral ProteinGym assays. Right: 31 viral ProteinGym assays and 45 EVEREST assays (all EVEREST assays are viral ). On viral assays, unfiltered models show larger gains in mean correlation than tier 1 filtered models from 150M onwards. On non-viral assays, the difference in mean correlation between filtering conditions is less than 0.009 in magnitude at every model size.
Figure 23: Larger filtered models show greater performance gains after the same number of fine-tuning steps. Performance recovered after a fixed number of fine-tuning steps for different model sizes. Changes are measured from each model’s filtered starting checkpoint after 500, 2,000 and 5,000 fine-tuning steps. Left: Increase in mean Spearman ρ on EVEREST for tier 1 filtered models. Right: Decrease in mean held-out viral cross-entropy (nats) for tier 0 filtered models up to 650M. The gain rises with model size at all three budgets. The right panel stops at 650M as this analysis was not performed for 3B.
Figure 24: At 650M and 3B, mean Spearman correlation exceeds the unfiltered reference at 2,500 fine-tuning steps in all four families with at least three EVEREST assays. Mean per-assay Spearman ρ between predicted and measured mutation effects for 650M and 3B tier 1 filtered models during fine-tuning on the withheld viral dataset. Panels show the four viral families with at least three EVEREST assays. Note that vertical scales differ between panels. Fine-tuning steps are shown in thousands. Dashed lines indicate the mean correlation for the same-size unfiltered model while open markers indicate the filtered model before fine-tuning. At both model sizes, the mean correlation for each displayed family is below the corresponding unfiltered reference before fine-tuning and exceeds it at the 2,500-step checkpoint.
Figure 25: Sequence perplexity and recovery without Retroviridae. Top: matched 3B unfiltered and tier-0 pretrained models, with one point per sequence, with 3,000 nonviral, 926 Retroviridae and 4,235 other viral sequences. Dashed line marks equal pseudo-perplexity. Bottom: tier-1 starting models fine-tuned without Retroviridae, cross-entropy and mean signed Spearman gaps to the same-size unfiltered model (positive means worse), with observed recovery-compute brackets. The dollar axis assumes 34% H100 utilisation and $4 per H100-hour. EVEREST averages all 45 assays. Recovery brackets span adjacent evaluations; open cost markers denote first-checkpoint upper bounds, squares at zero denote initial parity, and crosses denote unavailable results. Open trajectory markers at zero tokens identify pretrained starting models.
Figure 26: Cross-entropy after fine-tuning without Retroviridae. Top: tier-0 starting models. Bottom: tier-1 starting models. Dotted lines mark the same-size unfiltered references. The human-infecting panels use 535 non-Retroviridae sequences; other sequence strata are unchanged.
Figure 27: Tokens to recover performance without Retroviridae. Top: tier-0 starting models. Bottom: tier-1 starting models. Brackets span adjacent evaluations around first attainment of the same-size unfiltered reference, without interpolation. Open markers: first-checkpoint upper bounds; squares at zero: initial parity; NR triangles: target not reached; crosses: not run. The human-infecting panels use 535 non-Retroviridae sequences; EVEREST uses mean signed Spearman over all 45 assays.
Figure 28: Data filtering degrades fitness prediction more than stability prediction. Left to right: ProteinGym viral, EVEREST viral and ProteinGym non-viral assays. Each point shows the equally weighted mean Spearman correlation on stability assays (horizontal) and fitness assays (vertical) for one pretrained model (higher is better on both axes). Orange: viral assays, blue non-viral assays. Fitness/stability assay counts are respectively 21/6, 31/5 and 56/60 in the three panels. Fitness and stability use different assay sets; differences in their means may also reflect assay composition and difficulty.
Figure 29: Filtering effects on fitness and stability prediction per tier. Rows show ProteinGym viral, EVEREST viral and ProteinGym non-viral assays; columns show fitness (left) and stability (right). Points are equally weighted mean Spearman differences, filtered minus the same-size unfiltered reference, paired by assay. Bars are pointwise 95% percentile intervals from 20,000 paired-assay bootstrap resamples (seed 0). Negative values indicate lower predictive performance; the dashed zero line denotes reference performance. Orange: viral assays, blue non-viral assays. Fitness/stability counts are 21/6, 31/5 and 56/60 by row, and the two columns share limits within each row.
Protein language models are trained on highly imbalanced datasets, raising the question of how they represent underrepresented biological sequences. Using viral proteins as a case study across ESM model families, we identify a dominant nativeness axis in embedding space, aligned with masked reconstruction perplexity, that orders sequences from well-modeled cellular proteins through viral proteins to shuffled and random sequences. Scaling contracts this axis unevenly across viral families. Despite this, protein language model embeddings retain viral-specific signal: viral proteins remain linearly separable beyond zero-shot perplexity and shallow sequence features. Together, these results suggest that pLM representations are structured by a general notion of nativeness while preserving information specific to distinct biological groups.
Arthur Bigot, Harmon Bhasin, Core Francisco Park +2
Department of Computer Science, ETH Zurich, Zurich, Switzerland · Department of Chemistry and Chemical Biology, Harvard University, Cambridge, MA, USA · TwentyTwo, USA +2
Protein language models (pLMs) are powerful tools for protein annotation and mutation-effect prediction, but progress has largely depended on scaling model size and data, making them expensive and prone to memorization. Test-time scaling offers a complementary path, but methods developed for autoregressive LLMs do not transfer directly: masked encoders lack a generation loop, and amino acids do not verbalize in a reasoning language. Instead, we adapt recursive pretraining to masked pLMs, where a weight-tied Transformer core refines residue-aligned hidden states over an adjustable number of cycles. We instantiate this recipe as recAMPLIFY, a pLM with 155M unique parameters trained on the same 2T tokens as AMPLIFY-350M. Despite 56% fewer parameters, it improves CASP contact recovery and three out of four supervised tasks, notably raising fold accuracy by 22.5%. Because optimal depth for recAMPLIFY varies across tasks and assays, we then introduce Gated Mixture Recurrence (GMR), a framework that decodes each protein from a probability-weighted mixture of hidden states visited during recursion. GMR employs a learned residual gate at scalar, low-rank, or per-dimension resolution to control how much of each cycle's proposed update is accepted into the state. This unifies halting and gating into a single adaptive computation objective while recovering PonderNet-style halting as a special case without gates. Layered on top of recAMPLIFY, GMR further reduces expected layer depth by up to 37%, although the accuracy trade-off varies by task. Mechanistic analyses show that residual gating controls recurrence stability beyond the trained depth and that halt depth tracks biological properties.
Saurav Jha, Qianli Liao, Robert Vernon +2
Mila – Quebec AI Institute · Polytechnique Montréal · Center for Brains, Minds, and Machines +4
Protein language models are trained primarily with masked language modeling (MLM), which predicts masked amino-acid identities. Joint-embedding predictive architectures (JEPA) instead predict latent representations, but have not been applied to proteins. ProteinJEPA supplements MLM with a cosine loss for predicting the half-depth hidden states of a teacher given the unmasked sequence. On 19 tasks, with ESM2 at 35M and 150M parameters and three pretraining seeds, MLM+JEPA outperforms compute-matched and step-matched MLM-only continued training in 78 and 76 of 114 comparisons (14 losses, 22 ties). The median compute-matched gain is +0.0106 on structure- and homology-sensitive tasks versus +0.0041 elsewhere, led by SCOPe-40 retrieval and remote homology with improvements of 6.1 percentage points in Recall@1 and 2.7 points in accuracy, respectively. Gains on these tasks increase with model size from 8M to 150M. Against the off-the-shelf checkpoint, MLM+JEPA wins 81 of 114 comparisons (median +0.0068) without improving MLM loss. In random initialization the gain is smaller and replicates inconsistently across seeds (p=0.059). The same recipe improves the causal ProGen3 model, beating a compute-matched next-token-prediction control on 12 of 16 tasks. Ablations show that cosine loss beats mean squared error, while adding shallower targets removes most of the task gain. JEPA-only training collapses downstream performance: latent prediction complements MLM rather than replacing it. Code: https://anonymous.4open.science/r/protJepa-FF24
Dan Ofer, Dafna Shahaf, Michal Linial
Department of Biological Chemistry The Hebrew University of Jerusalem · School of Computer Science and Engineering The Hebrew University of Jerusalem