Uncertainty-Aware Fine-Tuning in Genomic Language Models

Gavin Hearne; Gail Rosen

doi:10.1145/3768322.3769100

Transformer-based Genomic Language Models (GLMs), such as DNABERT and the Nucleotide Transformer, have achieved strong performance on a wide range of genomic prediction tasks. However their tendency to produce overconfident predictions, particularly on noisy or out-of-distribution (OOD) data, remains a barrier to their reliable use in practice. In genomics, where unknown species, novel variants, and heterogeneous sequencing data are common, robust uncertainty quantification (UQ) is essential. While deep ensembles are often considered a gold standard for UQ, their computational overhead makes them impractical for large GLMs and large-scale genomic datasets, where billions of sequences may be processed through models consisting of billions of parameters. This motivates the exploration of lighter-weight approaches that preserve scalability without sacrificing predictive reliability. We investigate two such strategies for uncertainty-aware fine-tuning: Monte Carlo (MC) Dropout and Epistemic Neural Networks (Epinets). Both are implemented as lightweight uncertainty-aware heads on top of frozen embeddings from the 100M-parameter multispecies Nucleotide Transformer. MC Dropout estimates predictive distributions by sampling multiple stochastic forward passes, approximating Bayesian inference. Epinet, by contrast, augments the model with a secondary epistemic head that directly models uncertainty in function space, providing estimates in a single forward pass. To compare their effectiveness, we evaluate both methods along two key axes: calibration, which measures the agreement between predicted probabilities and observed frequencies, and OOD detection, which assesses the ability to distinguish in-distribution from out-of-distribution inputs. For calibration we report mean-squared error, and for OOD detection we use the area under the ROC curve (ROC-AUC) computed from epistemic uncertainty (mutual information). Our experiments span two tasks: gene classification, representing a stable in-domain problem, and metagenomic read classification, which introduces substantial distributional shift between bacterial training data and viral or eukaryotic test sequences. In gene classification, MC Dropout achieves stronger calibration, with lower error by MSE (0.008 vs. 0.028) compared to Epinet, reflecting more reliable confidence alignment when training and testing distributions are closely matched (Table 1). Under distributional shift in metagenomic classification, Epinet clearly outperforms Dropout, maintaining near-perfect calibration (MSE ≈ 0.000) and delivers higher OOD ROC-AUC across all non-bacterial classes (Table 2), with notable gains from 0.527 to 0.643 on withheld bacterial genomes, 0.436 to 0.744 on Archaea, and 0.447 to 0.775 on viral sequences. Consequently, MC Dropout is best suited to in-domain tasks, while Epinet provides sharper separation and more reliable uncertainty under OOD conditions. This demonstrates the task-dependent tradeoffs of uncertainty methods for Genomic Language Models. MC Dropout offers well-calibrated confidence in controlled, in-domain settings, but degrades under distributional shift. Epinet, by contrast, delivers robust calibration and OOD detection across diverse metagenomic datasets, establishing itself as the more reliable choice in realistic applications where unseen species are common. By evaluating lightweight UQ approaches, we show that Epinet in particular provides an effective and scalable basis for uncertainty-aware fine-tuning of Genomic Language Models.

Uncertainty-Aware Fine-Tuning in Genomic Language Models

Files and links (1)

Abstract

Metrics

Details

Uncertainty-Aware Fine-Tuning in Genomic Language Models

Files and links (1)

Abstract

Metrics

Details

Drexel University Social media