Contrastive Deep Learning for Variant Detection in Wastewater Genomic Sequencing

TL;DR

Wastewater genome sequencing faces noise, fragmentation, and mixed variants, limiting reference-based variant calling. The authors present a reference-free framework using Vector-Quantized Variational Autoencoders with EMA codebooks, masked reconstruction, and contrastive fine-tuning to learn discrete genomic patterns from k-mer reads. The approach yields high token-level reconstruction accuracy, interpretable codebooks, and significantly improved variant-discriminative embeddings, enabling scalable clustering without variant labels. This framework holds practical potential for real-time, interpretable wastewater surveillance and democratized genomic epidemiology.

Abstract

Wastewater-based genomic surveillance has emerged as a powerful tool for population-level viral monitoring, offering comprehensive insights into circulating viral variants across entire communities. However, this approach faces significant computational challenges stemming from high sequencing noise, low viral coverage, fragmented reads, and the complete absence of labeled variant annotations. Traditional reference-based variant calling pipelines struggle with novel mutations and require extensive computational resources. We present a comprehensive framework for unsupervised viral variant detection using Vector-Quantized Variational Autoencoders (VQ-VAE) that learns discrete codebooks of genomic patterns from k-mer tokenized sequences without requiring reference genomes or variant labels. Our approach extends the base VQ-VAE architecture with masked reconstruction pretraining for robustness to missing data and contrastive learning for highly discriminative embeddings. Evaluated on SARS-CoV-2 wastewater sequencing data comprising approximately 100,000 reads, our VQ-VAE achieves 99.52% mean token-level accuracy and 56.33% exact sequence match rate while maintaining 19.73% codebook utilization (101 of 512 codes active), demonstrating efficient discrete representation learning. Contrastive fine-tuning with different projection dimensions yields substantial clustering improvements: 64-dimensional embeddings achieve +35% Silhouette score improvement (0.31 to 0.42), while 128-dimensional embeddings achieve +42% improvement (0.31 to 0.44), clearly demonstrating the impact of embedding dimensionality on variant discrimination capability. Our reference-free framework provides a scalable, interpretable approach to genomic surveillance with direct applications to public health monitoring.

Figures (4)

• Figure 1: VQ-VAE Architecture Details. Encoder: token embedding (4097→128) + 2 Conv1D layers (256 hidden, kernel=3) + LN + Dropout(0.1) → $z_e$ (L×64). Quantizer: nearest-neighbor lookup in codebook (512 entries, 64-dim) with EMA updates (γ=0.95). Decoder: 2 Conv1D + LN + Linear projection to vocabulary logits. Red arrows show straight-through gradient flow.
• Figure 2: Masked reconstruction examples across corruption levels. (Top row) Original complete sequences with k-mer tokens highlighted. (Middle rows) Masked inputs with 10%, 20%, and 30% random token masking shown in gray. (Bottom rows) Model reconstructions with per-token confidence scores shown as heatmaps (darker = higher confidence). The model successfully infers most masked regions from context, with graceful degradation as masking increases. Error analysis shows the model predicts consensus sequences in ambiguous regions, valuable for handling low-quality wastewater data.
• Figure 3: Contrastive learning architecture for discriminative genomic embeddings. Starting from an input sequence, we generate two augmented views using random masking (15% probability) and token dropout (10% probability). Both views are independently processed through the frozen VQ-VAE encoder (pretrained weights fixed) followed by global mean pooling to obtain sequence-level representations. A trainable 2-layer MLP projection head with dimensions 256→128→{64,128} maps pooled features to L2-normalized contrastive embeddings. The InfoNCE loss encourages positive pairs ($v_1$, $v_2$ from same sequence) to cluster together (green solid arrows) while pushing negative pairs (embeddings from different sequences) apart (red dashed arrows). The embedding space visualization shows learned clusters corresponding to viral variants, with 128-dim projections achieving superior separation (Silhouette = 0.44) compared to 64-dim (Silhouette = 0.42). Hyperparameters: temperature $\tau = 0.5$, batch size 64, learning rate $10^{-4}$, 10 epochs.
• Figure 4: t-SNE visualization comparing embedding quality across models. (Left) Base VQ-VAE encoder embeddings show diffuse clusters with substantial overlap (Silhouette=0.31). (Middle) Contrastive-64 embeddings display tighter boundaries and reduced overlap (Silhouette=0.42). (Right) Contrastive-128 embeddings achieve best separation with well-defined cluster structures (Silhouette=0.44). Point colors indicate k-means cluster assignments (k=10). Higher-dimensional contrastive embeddings enable finer variant discrimination, critical for detecting emerging viral lineages in wastewater surveillance.