Contrastive Deep Learning for Variant Detection in Wastewater Genomic Sequencing

Wastewater Surveillance: Opportunities and Challenges

Wastewater-based epidemiology has emerged as a powerful tool for monitoring community-level viral prevalence, offering early warning signals for disease outbreaks and variant emergence. Unlike clinical testing, wastewater surveillance provides a non-invasive, population-wide snapshot that captures both symptomatic and asymptomatic infections. However, analyzing genomic sequences from wastewater presents unique computational challenges:

• Fragmentation: Sequencing reads are highly fragmented (100-300 bp), making de novo assembly difficult
• Noise: High sequencing noise and quality variation due to degraded viral RNA in wastewater
• Low concentration: Viral RNA constitutes a small fraction of total genetic material
• Complexity: Multiple co-circulating viral strains and variants are present simultaneously
• Reference dependence: Traditional bioinformatics pipelines require known reference genomes

Our Approach: Reference-Free Discrete Representation Learning

We propose a reference-free deep learning approach based on Vector-Quantized Variational Autoencoders (VQ-VAE) to learn discrete, compressed representations of viral genomic sequences. Our pipeline consists of three key components:

1. K-mer Tokenization: Convert raw DNA sequences into overlapping k-mer tokens (k=6), creating a vocabulary of 4,097 canonical k-mers that capture local sequence patterns
2. Discrete Representation Learning: Train a VQ-VAE with a codebook of 512 discrete latent codes (dimension 64) to compress sequences while preserving reconstruction quality
3. Enhanced Learning Objectives: Extend the base model with:
   • Masked VQ-VAE: BERT-style masked token prediction for robustness
   • Contrastive VQ-VAE: SimCLR-style contrastive learning for better clustering

This approach enables unsupervised discovery of sequence patterns, variant clustering, and robust sequence reconstruction without requiring reference genomes—making it particularly suitable for novel variant detection and real-time surveillance applications.

1. Vector-Quantized Variational Autoencoder (VQ-VAE)

The VQ-VAE architecture consists of three main components:

Encoder

The encoder maps tokenized sequences to continuous latent representations:

$$z_e = f_\theta(x)$$

where \(x \in \mathbb{Z}^{L}\) is the input token sequence, and the encoder \(f_\theta\) consists of:

• Token Embedding: \(V = 4{,}097\) vocabulary, embedding dimension \(d = 128\)
• 1D Convolutions: Two layers with kernel size 3, stride 1, hidden dimension 256
• Layer Normalization + Dropout: Dropout probability \(p = 0.1\)

Vector Quantizer

The quantizer discretizes continuous representations using a learned codebook:

$$z_q = \text{Quantize}(z_e) = e_k, \quad k = \arg\min_j \|z_e - e_j\|_2$$

where \(\{e_1, \ldots, e_K\}\) is the codebook with \(K=512\) entries, each of dimension \(D=64\). The codebook is updated using Exponential Moving Average (EMA) with decay \(\gamma = 0.95\).

Decoder

The decoder reconstructs the sequence from quantized representations:

$$\hat{x} = g_\phi(z_q)$$

The decoder architecture mirrors the encoder with 1D convolutions, layer normalization, and a final linear projection to vocabulary size.

Training Objective

The complete loss function combines reconstruction, commitment, and entropy regularization:

$$\mathcal{L} = \mathcal{L}_{\text{recon}} + \beta \mathcal{L}_{\text{commit}} - \lambda H[\mathcal{C}]$$

where:

• \(\mathcal{L}_{\text{recon}} = -\sum_{i=1}^L \log p(x_i | z_q)\) is the cross-entropy reconstruction loss
• \(\mathcal{L}_{\text{commit}} = \|z_e - \text{sg}[z_q]\|_2^2\) is the commitment loss with \(\beta = 0.1\)
• \(H[\mathcal{C}] = -\sum_{k=1}^K p_k \log p_k\) is the codebook usage entropy with \(\lambda = 0.003\)

2. Masked VQ-VAE

To improve robustness to missing or corrupted data, we extend the VQ-VAE with a masked language modeling objective inspired by BERT:

During training, we randomly mask 20% of input tokens (\(p_{\text{mask}} = 0.2\)) and train the model to reconstruct these masked positions:

$$\mathcal{L}_{\text{masked}} = -\sum_{i \in \mathcal{M}} \log p(x_i | z_q(\tilde{x}))$$

where \(\mathcal{M}\) is the set of masked positions and \(\tilde{x}\) is the masked input.

3. Contrastive VQ-VAE

To learn better cluster-separated representations, we fine-tune the encoder using contrastive learning with InfoNCE loss:

$$\mathcal{L}_{\text{contrast}} = -\mathbb{E}\left[\log \frac{\exp(\text{sim}(v_i, v_{i'})/\tau)}{\sum_{k=1}^{2B} \exp(\text{sim}(v_i, v_k)/\tau)}\right]$$

where:

• \(v_i, v_{i'}\) are L2-normalized embeddings from two augmented views of the same sequence
• \(\text{sim}(u, v) = u^\top v\) is cosine similarity
• \(\tau = 0.5\) is the temperature parameter
• \(B\) is the batch size (128)

We generate two augmented views for each sequence:

• View 1: Random masking with \(p_{\text{mask}} = 0.15\) (~22 tokens per sequence)
• View 2: Random dropout with \(p_{\text{drop}} = 0.10\) (~15 tokens per sequence)

python scripts/train.py \
  --data-path cleaned_reads.fastq \
  --output-dir experiments/1_standard_vqvae \
  --experiment-name standard_vqvae_50epochs \
  --epochs 50 \
  --batch-size 32 \
  --learning-rate 2e-4 \
  --num-codes 512 \
  --code-dim 64 \
  --n-gpu 2 \
  --gpu-ids "0,1" \
  --wandb-project vqvae-genomics \
  --save-freq 5

python scripts/train.py --config configs/experiment_configs/large_model.yaml

Project Structure

genomic_sequence_detection/
├── configs/                      # Configuration management
│   ├── default_config.yaml      # Default hyperparameters
│   └── experiment_configs/      # Experiment-specific configurations
│       └── large_model.yaml
│
├── scripts/                      # Executable scripts
│   ├── train.py                 # Main training loop with multi-GPU support
│   ├── evaluate.py              # Model evaluation and metrics
│   └── preprocess.py            # Data preprocessing pipeline
│
├── src/                         # Source code modules
│   ├── models/                  # Model architectures
│   │   ├── __init__.py
│   │   └── vqvae.py            # VQ-VAE, Encoder, Decoder, VectorQuantizer
│   ├── data/                    # Data processing
│   │   ├── __init__.py
│   │   └── tokenizer.py        # K-mer tokenization utilities
│   └── utils/                   # Utility functions
│       ├── __init__.py
│       └── wandb_init.py       # W&B experiment tracking
│
├── experiments/                 # Experiment outputs
│   ├── 1_standard_vqvae/       # Standard VQ-VAE experiments
│   ├── 2_masked_vqvae/         # Masked VQ-VAE experiments
│   ├── 3_contrastive_vqvae/    # Contrastive VQ-VAE experiments
│   └── 4_final_comparison/     # Comparative analysis
│
├── data/                        # Data files
│   ├── wastewater_seq_dataset.fastq
│   ├── cleaned_reads.fastq
│   └── virus_sequences.fastq
│
├── Trimmomatic/                 # Quality control tool
│   └── trimmomatic-0.39.jar
│
├── requirements.txt             # Python dependencies
└── README.md                    # Comprehensive documentation

Installation & Setup

System Requirements

• Python: 3.8 or higher
• CUDA: 11.3+ (for GPU training)
• GPU: 8GB+ VRAM recommended (tested on NVIDIA RTX series)
• RAM: 16GB+ system memory
• Storage: 10GB+ for data and checkpoints
• Java: 8+ (for Trimmomatic preprocessing)

Step 1: Clone Repository

git clone https://github.com/arrdel/genomic_sequence_detection.git
cd genomic_sequence_detection

Step 2: Install Dependencies

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install PyTorch (adjust for your CUDA version)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# Install other dependencies
pip install -r requirements.txt

Key Dependencies:

• torch>=2.0.0 - Deep learning framework
• biopython>=1.79 - Biological sequence processing
• numpy>=1.21.0 - Numerical operations
• wandb>=0.15.0 - Experiment tracking
• pyyaml>=6.0 - Configuration management
• tqdm>=4.65.0 - Progress bars

Step 3: Download Trimmomatic

wget http://www.usadellab.org/cms/uploads/supplementary/Trimmomatic/Trimmomatic-0.39.zip
unzip Trimmomatic-0.39.zip -d Trimmomatic/

Step 4: Setup Weights & Biases

# Login to W&B (required for experiment tracking)
wandb login

# Or disable W&B with --no-wandb flag during training

Usage Examples

1. Preprocessing

python scripts/preprocess.py \
  --input-fastq data/wastewater_seq_dataset.fastq \
  --output-fastq data/cleaned_reads.fastq \
  --k-mer 6 \
  --max-seq-length 150 \
  --run-fastqc \
  --run-trimmomatic

2. Training Standard VQ-VAE

python scripts/train.py \
  --data-path data/cleaned_reads.fastq \
  --output-dir experiments/1_standard_vqvae \
  --experiment-name my_experiment \
  --epochs 50 \
  --batch-size 32 \
  --num-codes 512 \
  --code-dim 64 \
  --n-gpu 2 \
  --gpu-ids "0,1"

3. Training with Configuration File

# Use predefined configuration
python scripts/train.py --config configs/default_config.yaml

# Create custom config in configs/experiment_configs/my_config.yaml
python scripts/train.py --config configs/experiment_configs/my_config.yaml

4. Evaluation

python scripts/evaluate.py \
  --checkpoint-path experiments/1_standard_vqvae/checkpoints/best_model.pt \
  --data-path data/cleaned_reads.fastq \
  --output-dir evaluation_results \
  --num-samples 100

Key Implementation Features

Dead Code Refresh Mechanism

To prevent codebook collapse (some codes never being used), we implement a dead code refresh mechanism that periodically replaces unused codes with high-variance encoder outputs:

# Enable dead code refresh during training
python scripts/train.py \
  --refresh-codes \
  --refresh-interval 500 \
  --refresh-min-count 5

Gradient Clipping

Gradient clipping stabilizes training and prevents exploding gradients:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Checkpoint Management

Automatic checkpoint saving with configurable frequency:

• Regular checkpoints: Every N epochs (--save-freq 5)
• Best model: Saved when validation loss improves
• Resume training: --resume experiments/.../checkpoint_epoch_20.pt

Troubleshooting

CUDA Out of Memory

• Reduce batch size: --batch-size 16
• Reduce model dimensions: --hidden-dim 128
• Use gradient accumulation (implementation in progress)
• Enable mixed precision training: --mixed-precision

Low Reconstruction Accuracy

• Train longer: --epochs 100
• Increase model capacity: --num-codes 1024
• Verify data quality (run FastQC)
• Adjust commitment weight: --commitment-cost 0.15

Low Codebook Utilization

• Increase commitment cost: --commitment-cost 0.25
• Reduce codebook size: --num-codes 256
• Enable dead code refresh: --refresh-codes
• Increase entropy regularization: --entropy-weight 0.005

Citation & License

@misc{chinda2025contrastivedeeplearningvariant,
      title={Contrastive Deep Learning for Variant Detection in Wastewater Genomic Sequencing}, 
      author={Adele Chinda and Richmond Azumah and Hemanth Demakethepalli Venkateswara},
      year={2025},
      eprint={2512.03158},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2512.03158}, 
}

License: MIT License - See LICENSE file for details
Contact: adelechinda@gsu.edu

Key Findings

• ✓ VQ-VAE achieves 99.5% reconstruction accuracy on fragmented wastewater sequences
• ✓ Discrete codebook effectively captures genomic patterns with only 19.7% utilization
• ✓ Entropy regularization prevents codebook collapse while maintaining diversity
• ✓ Masked learning improves robustness to missing/corrupted data by 12%
• ✓ Contrastive learning improves clustering separability by 35%
• ✓ Reference-free approach enables novel variant detection

Advantages Over Traditional Methods

• No reference genome required: Works with novel and unknown variants
• Computationally efficient: Training and inference in minutes vs. hours for alignment-based methods
• Learns meaningful representations: Discrete codes capture interpretable sequence patterns
• Robust to noise: Handles degraded RNA and sequencing artifacts effectively
• Scalable: Can process large-scale wastewater surveillance data in real-time

Limitations

• Modest exact sequence match rate (56%) leaves room for improvement
• Codebook utilization suggests potential for more compact representations
• Model currently focused on SARS-CoV-2; generalization to other pathogens needs validation
• Interpretability of learned discrete codes requires further analysis

Future Directions

1. Hierarchical VQ-VAE: Multi-scale representations to capture both local k-mer patterns and long-range dependencies
2. Phylogenetic Integration: Combine learned embeddings with traditional phylogenetic analysis
3. Multi-Pathogen Validation: Test on diverse viral datasets (influenza, RSV, norovirus) to assess generalization
4. Real-Time Deployment: Integrate into operational wastewater surveillance pipelines
5. Temporal Modeling: Incorporate time-series data to track variant emergence and evolution
6. Interpretability Analysis: Investigate biological meaning of learned codebook entries
7. Semi-Supervised Learning: Leverage small amounts of labeled variant data to improve clustering
8. Uncertainty Quantification: Add Bayesian extensions for confidence estimation

Impact

This work demonstrates that discrete representation learning can provide a scalable, reference-free approach to genomic sequence analysis. By removing the dependency on reference genomes, our method democratizes genomic surveillance, making it accessible to resource-limited settings and enabling rapid response to emerging viral threats. The learned discrete representations could serve as a foundation for:

• Early warning systems for novel variant detection
• Automated clustering and classification of viral sequences
• Robust sequence reconstruction in noisy surveillance data
• Transfer learning to other pathogen surveillance applications