UncertainGen: Uncertainty-Aware Representations of DNA Sequences for Metagenomic Binning

TL;DR

This work tackles metagenomic binning under sequence-level uncertainty, arguing that deterministic representations struggle to separate fragments that belong to multiple genomes or have highly similar features. It proposes UncertainGen, a probabilistic embedding that represents each DNA fragment as a Gaussian with mean $\boldsymbol{\mu}$ and covariance $\mathbf{S}$, enabling an expanded latent space and a closed-form probabilistic similarity score $q$ derived from $K_{ij} = \alpha(\mathbf{S}_i + \mathbf{S}_j)$. The approach is trained with a contrastive objective and backed by theoretical results on distinguishability and expressivity, including $\ extit{epsilon}$-distinguishability and packing-number bounds. Empirically, UncertainGen delivers improvements over deterministic k-mer and LLM-based embeddings on real metagenomic datasets, producing robust, uncertainty-aware partitions with a lightweight, scalable architecture. This probabilistic framework enhances binning robustness and opens avenues for uncertainty-aware genomic analyses and future extensions beyond Gaussian representations.

Abstract

Metagenomic binning aims to cluster DNA fragments from mixed microbial samples into their respective genomes, a critical step for downstream analyses of microbial communities. Existing methods rely on deterministic representations, such as k-mer profiles or embeddings from large language models, which fail to capture the uncertainty inherent in DNA sequences arising from inter-species DNA sharing and from fragments with highly similar representations. We present the first probabilistic embedding approach, UncertainGen, for metagenomic binning, representing each DNA fragment as a probability distribution in latent space. Our approach naturally models sequence-level uncertainty, and we provide theoretical guarantees on embedding distinguishability. This probabilistic embedding framework expands the feasible latent space by introducing a data-adaptive metric, which in turn enables more flexible separation of bins/clusters. Experiments on real metagenomic datasets demonstrate the improvements over deterministic k-mer and LLM-based embeddings for the binning task by offering a scalable and lightweight solution for large-scale metagenomic analysis.

Figures (9)

• Figure 1: Illustration of the metagenomic binning process. Starting from a set of DNA sequences (a), the process ends with their two-dimensional embeddings derived from 2-mer profiles (d). In general, these embeddings allow the DNA fragments from two different species to be correctly clustered. However, the second and third DNA sequences in (a) pose an exception: although distinct, their $k$-mer representations shown in (c) are highly similar and, consequently, their embeddings are also very close (shown as the two empty circled points in (d)). Because the k-mer profiles of DNA within a species tend to be (locally) similar, the contrastive learning procedure attempts to position such fragments in both clusters but, since this is not possible, ultimately places them between them.
• Figure 2: Packing number ($\mathcal{P}_\tau^D$)
• Figure 3: Visualization of the learned embeddings and the variance distribution of the sequences.
• Figure 4: Metagenomic binning results. Cluster counts are segmented by F1-score quality ranges. The dark blue portion highlights the highest-quality bins for each model-dataset combination.
• Figure 5: Ablations studies examining the behavior of the proposed UncertainGen model.

Theorems & Definitions (15)

• Lemma 3.1
• Definition 3.2
• Lemma 3.3
• Corollary 3.3.1
• Theorem 3.4
• Lemma A.1
• proof
• Lemma A.2
• proof
• Lemma A.3