Investigating Data Pruning for Pretraining Biological Foundation Models at Scale

AI summary

Overview

• Research area: Machine learning for biology — specifically data pruning / coreset selection for pretraining biological foundation models (BioFMs) on RNA and protein sequences.
• Technical level: Advanced. The method is built on influence functions, Hessian/Fisher curvature approximations, and second-order Taylor expansions, though the paper motivates them with practical compute constraints.
• Scope: The paper proposes a post-hoc, influence-guided data pruning framework, derives a scalable subset-based self-influence estimator, and validates it by pretraining RNA-FM and ESM-C from scratch on pruned coresets of 0.2 million sequences.

What This Paper Is About

Biological foundation models are pretrained on millions to billions of sequences, which makes full pretraining prohibitively expensive and hard to reproduce for most academic labs — RNA-FM used over 23 million RNA sequences, and ESM has scaled to 2.78 billion protein sequences. The authors ask whether a small, carefully selected coreset of sequences can be used to retrain such models from scratch and still reach competitive downstream performance, without needing access to the original full training run. Their goal is a pruning method that works post-hoc, at biological scale, where training-dynamics methods and pairwise-similarity density methods are both impractical.

Key Contributions

1. A post-hoc influence-guided data pruning framework for BioFMs that does not require access to the full training process or its training dynamics.
2. A theoretical reformulation of the self-influence function based on curvature estimated over a small, randomly sampled training subset. The derivation relies on a "subset-based ERM" assumption, under which the subset curvature approximates the full-training curvature, and is stated formally as Proposition 1.
3. Two influence-guided selection strategies: Top-k Influence-guided Selection (Top I), which keeps the highest-influence examples, and Coverage-Centric Influence-guided Selection (CCI), which stratifies over the influence-score distribution so both easy and hard examples survive extreme pruning.
4. A scalable approximation using a diagonal empirical Fisher matrix, reducing the inverse-Hessian cost from O(M · d² + d³) to O(M · d), where M is the subset size and d is the number of model parameters — making influence estimation tractable for models with billions of parameters.

Main Findings

• RNA results beat random selection across the board: On RNA Type Classification, RNA Modification Prediction, and CRISPR On-Target Prediction, both Top I and CCI outperform the Random 0.2M baseline. Full RNA-FM (23M) reaches 91.93 ACC / 91.87 F1 on TypeCls and 94.98 AUC on Modif; CCI at 0.2M reaches 82.88 ACC / 83.12 F1 and 93.86 AUC.
• CCI was best on all three function/engineering tasks: CCI at 0.2M scored 82.88 ACC, 83.12 F1 (TypeCls), 93.86 AUC (Modif), and 32.90 SC / .0135 MSE (CRI-On). The paper states CCI surpasses full RNA-FM on CRISPR On-Target Prediction (SC 32.90 vs. 31.87); RNA-FM retains the lower MSE (.0118 vs. .0135).
• Top I was strongest on structure and interaction tasks: On RNA Distance Map Prediction, Top I at 0.2M reached 79.25 R² and 86.47 SC versus Random's 76.71 R² and 84.90 SC (RNA-FM full: 83.26 R², 89.21 SC). On Contact Map Prediction, Top I exceeded RNA-FM on Top-1.0L (94.36 vs. 93.93) and Top-0.5L (98.41 vs. 98.28) precision.
• Extreme pruning still works: RNA pruning was performed over the full 23M-sequence training corpus, i.e. over 99% pruning, and the conclusion notes that RNA structure prediction performance comparable to the original RNA-FM was achieved using less than 1% of the full 23-million-sequence training set.
• Small coresets beat 10× larger random subsets: Models trained on Top I or CCI coresets (0.2M) outperformed Random 2M baselines on all function and engineering prediction tasks despite using only 10% of the data volume, which the authors present as evidence of substantial redundancy in biological sequence datasets.
• Protein generalizability confirmed on ESM-C: Top I and CCI at 0.2M outperformed both Random 0.2M and Random 2M on Binary Localization (CCI 79.25 ACC vs. Random 0.2M 73.64 and Random 2M 75.76), Secondary Structure (71.48 vs. 66.18 and 67.20), and PPI Affinity (Top I MAE 2.06 vs. 2.51 and 2.39). ESM-C trained on the full 2.78B sequences remains higher at 91.63, 86.10, and MAE 1.92.
• Fine-tuning on the subset was necessary: Adapted variants (Top I, CCI) consistently outperformed their "w/o ft" counterparts across both RNA tables, supporting the role of the subset-based ERM assumption in reducing estimation error.

Methodology in Plain English

The authors start from classical influence functions, which estimate how much a single training example changes a model's behavior by combining gradients with the inverse Hessian (curvature) of the training loss. Computing that Hessian over an entire biological corpus with billions of parameters is impossible, so they make three simplifications in sequence.

First, instead of the curvature of the full training set, they use the curvature of a small, randomly sampled subset, under an assumption that the model is well trained on that subset. In practice they satisfy this with a light-weight fine-tuning pass (one epoch) on a randomly selected 0.2-million-sequence subset — the cost is described as negligible relative to the rest of the pipeline.

Second, they replace the Hessian with the empirical Fisher information matrix, which is standard for models trained with negative log-likelihood, and then use only its diagonal. This turns the inverse-Hessian computation into a linear-time operation, making scoring feasible at scale.

Third, they use the resulting per-example self-influence scores in two ways: Top I simply keeps the 0.2 million highest-scoring sequences, while CCI sorts by score, discards a fixed proportion of the hardest examples (a hard cutoff rate β = 5), partitions the remainder into 50 equal-width influence strata, and draws a budgeted sample from each stratum, redistributing unused budget if a stratum is too small. The pruning rate is α = 90.

Evaluation is deliberately extreme: only 0.2 million sequences are retained and models are pretrained from scratch for 10 epochs. For RNA this means pruning the entire 23M corpus; because the full 2.78-billion-sequence ESM-C corpus was unavailable, the authors instead assembled roughly 4.5 million protein sequences from UniRef50 and pruned over that. They compare against the original full-data model, an untrained model, and uniform random sampling at both 0.2M and 2M.

Why This Matters

Research impact. The paper argues that current biological pretraining datasets contain substantial redundancy, and that a compact, influence-selected subset can substitute for the full corpus in several RNA structure and interaction tasks. If it holds more broadly, this changes how BioFM pretraining is scoped: reproducibility stops requiring industrial-scale compute, and the community can study pretraining choices without retraining on billions of sequences. The code is released at https://github.com/victor-yifanwu/bio-coreset.

Real-world applications (drawn from the tasks evaluated in the paper):

• RNA function and modification annotation — the TypeCls and Modif tasks use datasets from Amin et al. (2019) and Duan et al. (2019).
• CRISPR guide design — CRISPR On-Target Prediction (Chuai et al. 2018), where the pruned CCI model matched or exceeded the full RNA-FM on SC and MSE.
• RNA structure characterization — secondary structure on bpRNA (Danaee et al. 2018), distance map prediction, and contact map prediction.
• Protein analysis — binary localization prediction (Almagro Armenteros et al. 2017), protein secondary structure (Klausen et al. 2019), and protein–protein interaction affinity (Moal and Fernández-Recio 2012).

Industry relevance. Reducing pretraining data by more than 99% directly lowers GPU hours, energy, and storage requirements. RNA-FM's original training used 8 A100 GPUs over 30 days, so even a partial reduction in required data has a large cost implication for any organization training or reproducing biological foundation models.

Future Directions

• Finding better coreset sizes for protein models. The paper explicitly leaves the exploration of more suitable coreset sizes for protein foundation models as future work, noting that the 0.2M- versus 2.78B-scale gap explains the residual performance shortfall against full ESM-C.
• Testing whether the redundancy findings generalize to other modalities and scales. The framework was validated on two models (RNA-FM and ESM-C), so behavior on DNA, multi-modal, or larger protein corpora is open.
• Reducing reliance on fine-tuning for the subset-based ERM assumption. The ablation shows adaptation matters, but the condition (one epoch on a random subset) is a practical heuristic rather than a guarantee.
• Choosing between Top I and CCI more systematically. The results suggest Top I is preferable for structure- and interaction-related tasks while CCI wins on function and engineering prediction, but the paper does not provide a rule for selecting between them in advance.

Target Audience

Researchers working on biological foundation models and data-efficient pretraining will get the most from this paper, particularly those who lack the compute to reproduce full-scale BioFM training. It also suits machine learning researchers interested in coreset selection and influence functions applied outside CV/NLP, and bioinformatics practitioners who want to know whether small curated sequence sets can substitute for massive corpora in RNA and protein tasks. Readers need some familiarity with empirical risk minimization and curvature-based approximations to follow the methods section.

Authors’ abstract

Biological foundation models (BioFMs), pretrained on large-scale biological sequences, have recently shown strong potential in providing meaningful representations for diverse downstream bioinformatics tasks. However, such models often rely on millions to billions of training sequences and billions of parameters, resulting in prohibitive computational costs and significant barriers to reproducibility and accessibility, particularly for academic labs. To address these challenges, we investigate the feasibility of data pruning for BioFM pretraining and propose a post-hoc influence-guided data pruning framework tailored to biological domains. Our approach introduces a subset-based self-influence formulation that enables efficient estimation of sample importance at low computational cost, and builds upon it two simple yet effective selection strategies, namely Top-k Influence (Top I) and Coverage-Centric Influence (CCI). We empirically validate our method on two representative BioFMs, RNA-FM and ESM-C. For RNA, our framework consistently outperforms random selection baselines under an extreme pruning rate of over 99 percent, demonstrating its effectiveness. Furthermore, we show the generalizability of our framework on protein-related tasks using ESM-C. In particular, our coreset even outperforms random subsets that are ten times larger in both RNA and protein settings, revealing substantial redundancy in biological sequence datasets. These findings underscore the potential of influence-guided data pruning to substantially reduce the computational cost of BioFM pretraining, paving the way for more efficient, accessible, and sustainable biological AI research.