DepMicroDiff: Multimodal Imputation in Microbiome Data

• DepMicroDiff is a multimodal imputation framework that integrates latent diffusion modeling with dependency-aware transformers for robust microbiome data reconstruction.
• It conditions imputation on patient metadata using LLM-based encoding, enhancing context by incorporating features like tissue type and pathological stage.
• Performance evaluations on TCGA datasets show improved metrics over KNN and mbVDiT, validating its effectiveness for biomarker discovery and multi-omics applications.

DepMicroDiff is a multimodal imputation framework designed for microbiome data, integrating latent-structured diffusion modeling with explicit dependency-aware attention and conditioning on patient metadata. Developed to address the pronounced sparsity and technical dropout inherent to microbial abundance measurements, DepMicroDiff introduces novel mechanisms to capture both mutual dependencies among microbial taxa and context provided by host-level features, leading to marked improvements in imputation accuracy for biomarker discovery and downstream multi-omics applications (Sadia et al., 31 Jul 2025).

1. Model Architecture and Workflow

DepMicroDiff comprises four interconnected modules within its imputation pipeline:

• VAE Pretraining Module: A variational autoencoder (VAE) encoder $E$ learns compact latent representations $\hat x_{mb}^i = E(x_{mb}^i)$ from raw, partially observed microbiome vectors $x_{mb}^i \in \mathbb{R}^d$. The decoder reconstructs the original input, producing a structured latent space suitable for capturing inter-tissue generalization and minimizing overfitting in low-data regimes.
• LLM-Based Metadata Encoder: Patient metadata—such as tissue type, pathological stage, and age—are converted to text tokens, embedded with a pretrained transformer-based LLM (e.g., BERT), and projected to match the microbial latent dimension, yielding $m_{meta}$. This embedding is fused with the masked observed latent for conditioning.
• Diffusion Model: A forward–reverse process is conducted in the VAE latent space. Gaussian noise is injected iteratively:

$$q(\hat x^{i,t} \mid \hat x^{i,t-1}) = \mathcal{N}(\hat x^{i,t};\sqrt{1-\beta_t}\,\hat x^{i,t-1},\ \beta_t I)$$

with a cosine noise schedule. At inference, denoising reconstructs the clean latent via $p_\theta(\hat x^{i,t-1}|\hat x^{i,t},c)$, where $c$ encodes both observed latents and contextual metadata.

• Dependency-Aware Transformer (DAT): The core denoising function $\epsilon_\theta(\cdot, t, c)$ is implemented with a masked transformer integrating custom attention masks based on a learned binary dependency graph and autoregressive ordering. The mask limits attention flow to plausible microbe-microbe connections determined by Granger causality and mutual information statistics, merged with an exponentially decaying autoregressive schedule.

During training, input vectors are encoded with the VAE, masked in the latent space, augmented with Gaussian noise, and denoised by DAT under contextual conditioning. The inverse diffusion chain at inference reconstructs latent representations, which are decoded by the VAE into complete microbiome profiles.

2. Diffusion Process and Training Objective

The DepMicroDiff forward process incrementally corrupts the latent vectors with Gaussian noise through

$$q(\hat x^{0:T}) = q(\hat x^0) \prod_{t=1}^T q(\hat x^t|\hat x^{t-1})$$

with each step defined as

$$q(\hat x^t|\hat x^{t-1}) = \mathcal{N}(\hat x^t; \sqrt{1-\beta_t}\,\hat x^{t-1}, \beta_t I).$$

The reverse process is learned as

$\hat x_{mb}^i = E(x_{mb}^i)$0

parameterized by the DAT.

The primary training loss is the DDPM-style noise prediction objective: $\hat x_{mb}^i = E(x_{mb}^i)$1 where $\hat x_{mb}^i = E(x_{mb}^i)$2 and $\hat x_{mb}^i = E(x_{mb}^i)$3.

3. Dependency-Aware Transformer and Conditioning

DAT processes each microbial feature as a token in the latent space, applying two critical forms of masking:

• Autoregressive Mask: Tokens are partitioned into $\hat x_{mb}^i = E(x_{mb}^i)$4 subsets $\hat x_{mb}^i = E(x_{mb}^i)$5 with exponentially decaying support; at each step, only tokens from previous subsets and conditioning tokens are visible.
• Dependency Mask: A binary graph $\hat x_{mb}^i = E(x_{mb}^i)$6 encodes meaningful taxa-taxa influences ($\hat x_{mb}^i = E(x_{mb}^i)$7 from Granger causality; $\hat x_{mb}^i = E(x_{mb}^i)$8 from mutual information), used to mask attention logits in the transformer.

Positional and diffusion-timestep embeddings are added per token. Conditioning tokens—comprising observed portions of the latent vector and the metadata embedding—are prepended at each diffusion step.

4. Data, Preprocessing, and Evaluation Metrics

DepMicroDiff is validated on The Cancer Genome Atlas (TCGA) microbiome datasets:

• STAD: 530 samples, 106 microbes, 87.6% zeros
• COAD: 561 samples, 106 microbes, 63.2% zeros
• HNSC: 587 samples, 106 microbes, 79.6% zeros
• Pretraining: Smaller READ (182) and ESCA (248) are also included.

Data is normalized by scaling and log-transforming abundances: $\hat x_{mb}^i = E(x_{mb}^i)$9 During training, entries are randomly masked to reflect empirical dropout patterns.

Performance is quantified using Pearson correlation coefficient (PCC), cosine similarity (COS), root-mean-square error (RMSE), and mean absolute error (MAE). The model consistently outperforms KNN and the previous state-of-the-art mbVDiT across these metrics in all datasets.

Metric	STAD (KNN/mbVDiT/DepMicroDiff)	COAD (KNN/mbVDiT/DepMicroDiff)	HNSC (KNN/mbVDiT/DepMicroDiff)
PCC	0.232 / 0.634 / 0.641	0.188 / 0.704 / 0.712	0.247 / 0.626 / 0.636
COS	0.430 / 0.806 / 0.812	0.241 / 0.791 / 0.789 (2nd)	0.357 / 0.802 / 0.798 (2nd)
RMSE	3.371 / 1.320 / 1.290	4.336 / 0.934 / 0.927	4.128 / 1.155 / 1.098
MAE	3.255 / 0.956 / 0.933	3.862 / 0.530 / 0.524	3.679 / 0.807 / 0.799

5. Component Ablations and Contributions

Ablation studies demonstrate the importance of each architectural element:

• VAE Pretraining: Eliminating this step reduces PCC by approximately 0.03.
• Metadata Conditioning: Omitting LLM-based metadata encoding lowers PCC by 0.02–0.04 across all cohorts.
• Dependency-Aware Masking: Replacing Granger/mutual-information masks with a causal-agnostic mask results in a 0.015 PCC decrease.

These ablations confirm that the combination of structured latent space, context encoding, and explicit dependency modeling collectively accounts for the observed performance gains.

6. Limitations and Future Areas of Investigation

Several limitations and potential avenues for enhancement are identified:

• Computational Cost: The joint diffusion and transformer stack with 1,000 diffusion steps is compute-intensive. Reducing the number of steps or applying alternatives such as DDIM samplers may address this.
• Interpretability: While Granger-causal and mutual-information priors guide mask construction, deeper investigation of biological plausibility and interpretability of learned dependencies remains an open problem.
• Temporal and Multimodal Extensions: The current formulation is sample-wise; longitudinal extensions for temporal cohort data and cross-omics imputation are plausible next steps.
• Applicability to Other Data Types: Deployment to other zero-inflated omic profiles (e.g., single-cell ATAC-seq) could test generalizability.

7. Context within Diffusion-Based and Multimodal Imputation Research

DepMicroDiff advances diffusion-based imputation by explicitly merging pairwise microbial dependencies with contextual patient metadata, contrasting with traditional approaches which neglect inter-taxa interactions or host context. The use of transformer-based dependency masking, VAE pretraining for robust latent modeling, and LLM-based patient conditioning collectively define a new paradigm for imputation in sparsely observed, multimodal biomedical datasets (Sadia et al., 31 Jul 2025). A plausible implication is that the framework could serve as a foundation for unified generative modeling of microbiome and host multi-omics data in clinical prediction workflows.