BaVE is a deep generative framework for learning batch-invariant representations of microbiome count data to enable robust cross-study integration and batch-aware machine learning. By combining a ZINB likelihood with dual batch handling through conditioning and adversarial learning, explicit modeling of sequencing depth, and covariate-aware inference, BaVE produces batch-robust sample embeddings suitable for clustering, ordination, and predictive modeling across heterogeneous microbiome cohorts, while avoiding reliance on feature-level batch-corrected tables.
BaVE is built on a Zero-Inflated Negative Binomial (ZINB) generative model, enabling it to faithfully capture the sparsity and overdispersion inherent in microbiome count data.
BaVE suppresses batch effects using two complementary strategies:
Batch information is explicitly injected into both the encoder and decoder, allowing the model to account for batch during inference and generation.
A dedicated batch classifier with gradient reversal forces the latent space to discard batch-discriminative signals, promoting batch-invariant representations.
BaVE includes an explicit library-size (sequencing depth) offset in the decoder, preventing latent embeddings from being confounded by technical depth variation.
BaVE natively integrates key host variables, including age, BMI, and antibiotic usage — directly into both inference and generation, enabling biologically informed representation learning.
Rather than producing batch-corrected abundance tables, BaVE outputs batch-robust sample embeddings, making it ideal for
- Cross-study clustering,
- Ordination and visualization,
- Phenotype prediction,
- Integration across cohorts and protocols
All training and inference workflows are controlled through YAML-based configuration, ensuring complete reproducibility and transparent experimental tracking.
BaVE explicitly models: • Zero inflation and overdispersion via a ZINB likelihood • Batch effects through both conditioning and adversarial learning • Sequencing depth via an explicit size-factor offset • Host covariates (e.g., age, BMI, antibiotic usage) in both encoder and decoder
BaVE produces batch-robust sample embeddings rather than corrected abundance tables. It is therefore intended for representation learning and integration, not feature-level batch correction or differential abundance testing.
- Handles batch effects during representation learning, not after, ensuring inherently robust embeddings
- Aligns model structure with microbiome data properties, including sparsity, overdispersion, and depth variation
- Produces batch-invariant latent representations suitable for cross-study integration and machine learning
- Actively suppresses technical signals rather than merely shifting or rescaling them
- Integrates biological covariates directly into learning, preserving meaningful variation while removing confounding
- Designed for reproducibility, scalability, and downstream analytical flexibility
- integrating microbiome datasets across studies/platforms
- building predictive models robust to batch
- learning embeddings for large-scale microbiome ML
- your goal is taxon-level inference
- differential abundance is the primary outcome
- sample size is extremely small (<150)
BaVE requires two primary input files: a raw count matrix (counts.tsv) and a sample metadata table (metadata.tsv). Together, these define the microbiome composition, technical batch structure, and biological covariates necessary for batch-aware representation learning.
The file counts.tsv must be provided as a tab-separated table in which rows correspond to taxa or features and columns correspond to samples, with each cell containing a non-negative integer representing the raw sequencing count of a given taxon in a given sample. This format allows BaVE to operate directly on untransformed count data, enabling explicit modeling of sparsity, overdispersion, and sequencing depth through its ZINB generative framework.
Example:
sample1 sample2 sample3
TaxonA 10 0 5
TaxonB 0 12 1
...The file metadata.tsv must be a tab-separated table where each row corresponds to a sample and columns contain batch labels, biological covariates, and optional phenotype information. This file provides the contextual information required for BaVE’s dual batch handling and covariate-aware inference.
Example:
sample_id batch age BMI antibiotics phenotype
sample1 A 45 23.1 0 control
sample2 B 50 28.4 1 case
...bave validate --counts data/example/counts.tsv --metadata data/example/metadata.tsvbave train --config configs/default.yamlArtifacts written to:
runs/example/
model.pt
embeddings.tsv
features_kept.txt
batch_levels.txt
covariate_stats.json
train_history.json
run_card.json
bave export --config configs/default.yamlbave qc --config configs/default.yamlGenerates:
- PCA plots colored by batch and phenotype
- Depth leakage diagnostics
- Batch predictability tests
Configuration (YAML)
All runs are controlled via YAML files.
Key sections:
preprocessing:
min_total_count_per_feature: 20
min_prevalence_fraction: 0.02
size_factor_method: median_ratio
model:
latent_dim: 16
hidden_dim: 256
batch_embedding_dim: 8
adversary_hidden_dim: 128
training:
epochs: 250
batch_size: 128
kl_weight: 1.0
adv_weight: 1.0
kl_warmup_epochs: 25
adv_warmup_epochs: 25Each experiment should preserve its config file for reproducibility.
BaVE produces:
Tab-separated latent representation:
z1 z2 z3 ...
sample1 0.15 -1.23 0.44
sample2 -0.87 0.31 -0.55
...These embeddings are suitable for:
- clustering
- ordination (PCA/UMAP)
- machine learning
- cross-cohort integration
git clone https://github.com/KarwaiHong/bave
cd bave
pip install -e .BaVE is released under the Apache-2.0 License.
Maintained by:
Dr. Hong Kar Wai
Jeffrey Cheah School of Medicine and Health Sciences, Monash University Malaysia
Email: hong.karwai@monash.edu