Amortized Variational Inference for Scalable Bayesian Tensor Factorization

Abstract

This repository contains a high-performance PyTorch implementation of an Amortized Bayesian Tensor Factorization framework. It is designed to address the critical scalability bottlenecks in spatial transcriptomics and single-cell RNA-sequencing (scRNA-seq) analysis. By replacing traditional, iterative Markov Chain Monte Carlo (MCMC) and local Variational Inference (VI) updates with an Amortized Neural Encoder, this framework scales probabilistic tensor decomposition to massive tissue datasets while strictly preserving Bayesian uncertainty quantification.

This architecture is specifically optimized for high-dimensional clinical datasets, such as those generated by the FunGen-AD consortium, enabling the discovery of latent biological pathways across complex spatial morphologies.

---

1. Scientific Background & Mathematical Gap

In spatial transcriptomics, biological tissue is naturally represented as a multi-way array, or tensor $\mathcal{X} \in \mathbb{R}^{I \times J \times K}$, where the dimensions represent Genes ($I$), Cells ($J$), and Spatial Zones ($K$).

Traditional deep learning matrix factorization yields deterministic point estimates, failing to capture the inherent noise and uncertainty of biological sequencing. Conversely, rigorous Bayesian Tensor Factorization provides excellent uncertainty estimates but relies on local parameter updates (e.g., Gibbs Sampling). The computational complexity of these updates scales linearly with the number of cells $O(N)$, rendering them computationally intractable for modern datasets containing hundreds of thousands of cells.

The Solution: We bridge this gap via Amortized Inference. Instead of optimizing local variational parameters for every independent cell, we train a global deep neural network (the Encoder) to parameterize the variational distributions. This reduces the inference of a cell's latent state from an iterative optimization problem to a single, highly parallelizable $O(1)$ forward pass.

---

2. Mathematical Formulation

2.1 The Generative Model (CANDECOMP/PARAFAC Decomposition)

We assume the observed transcriptomic tensor $\mathcal{X}$ is generated by a low-rank CP decomposition of rank $R$. The true latent biological structure is defined by three continuous factor matrices:

• $\mathbf{A} \in \mathbb{R}^{I \times R}$: Gene expression signatures.
• $\mathbf{B} \in \mathbb{R}^{J \times R}$: Cell-specific pathway activation.
• $\mathbf{C} \in \mathbb{R}^{K \times R}$: Spatial morphological dependencies.

The deterministic reconstruction of the tensor is given by the sum of $R$ rank-one tensors: $$\hat{\mathcal{X}} = \sum_{r=1}^R \mathbf{a}_r \circ \mathbf{b}_r \circ \mathbf{c}_r$$

2.2 Amortized Variational Inference

We treat the spatial ($\mathbf{C}$) and gene ($\mathbf{A}$) matrices as global parameters. The cell-specific factors ($\mathbf{B}$) are treated as local latent variables.

We introduce an inference network $f_\phi$, parameterized by weights $\phi$, which maps the observed gene expression profile of the $j$-th cell $\mathbf{x}_j$ to the parameters of a Gaussian distribution:

$$[\boldsymbol{\mu}_j, \log \boldsymbol{\sigma}^2_j] = f_\phi(\mathbf{x}_j)$$

To allow gradient flow through the stochastic sampling process, we utilize the Reparameterization Trick: $$\mathbf{b}_j = \boldsymbol{\mu}_j + \boldsymbol{\epsilon} \odot \boldsymbol{\sigma}_j \quad \text{where} \quad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$

2.3 The Optimization Objective (ELBO)

The network and global parameters are jointly optimized by maximizing the Evidence Lower Bound (ELBO). We minimize the negative ELBO, which balances data likelihood (reconstruction fidelity) with a Kullback-Leibler (KL) divergence regularizer that penalizes deviation from a standard normal prior:

$$\mathcal{L}_{ELBO} = \sum_{j=1}^J \left( \mathbb{E}_{q_\phi(\mathbf{b}_j | \mathbf{x}_j)} [\log p_\theta(\mathbf{x}_j | \mathbf{b}_j, \mathbf{A}, \mathbf{C})] - D_{KL}(q_\phi(\mathbf{b}_j | \mathbf{x}_j) || p(\mathbf{b}_j)) \right)$$

---

3. Repository Architecture

BayesTensor/
├── src/
│   ├── data_gen/
│   │   └── synthetic.py        # Generates ground-truth multi-way arrays
│   ├── models/
│   │   ├── encoder.py          # Multi-Layer Perceptron (MLP) Inference Network
│   │   └── amortized_cp.py     # CP Tensor Decoder & Reparameterization framework
│   ├── real_data.py            # Automated preprocessing pipeline for scRNA-seq
│   ├── train.py                # Simulation Study (Identifiability Proof)
│   └── train_real.py           # Biological Application Engine
├── outputs/
│   └── plots/                  # Auto-generated metrics and factor visual validation
├── requirements.txt            
└── README.md

---

4 Experiment 1: Simulation Study (Mathematical Identifiability)

To prove the architecture correctly discovers underlying topological structure, we generate a synthetic tensor with a known low-rank ground truth, inject Gaussian noise, and train the model to recover the exact spatial factors.

python -m src.train

Expected Output: Successful convergence logs and an auto-generated heatmap (saved to outputs/plots/) demonstrating alignment between the simulated Ground Truth spatial factors and the model's Amortized CP Inference factors.

4.1 Experiment 2: Clinical Data Application (Scalability)

This pipeline assesses the model's performance on highly sparse, correlated biological data. It fetches a benchmark 10x Genomics Peripheral Blood Mononuclear Cell (PBMC) dataset, executes standard bioinformatics preprocessing (normalization, log-transformation, HVG selection via scanpy), projects the cells into a spatial topology, and executes the Amortized training loop.

5. Dependencies

torch >= 2.0.0
numpy >= 1.23.0
matplotlib >= 3.7.0
scanpy >= 1.9.0
anndata >= 0.9.0
pandas >= 2.0.0

6. Conclusion

This framework demonstrates that Amortized Variational Inference provides a scalable, mathematically rigorous alternative to traditional tensor decomposition methods. By leveraging neural networks for parameter estimation, we achieve a significant reduction in computational latency while preserving the probabilistic integrity required for clinical transcriptomic discovery.

7. Citation

If you use this framework in your research, please cite:

@article{Mboya2024Amortized,
  title={Amortized Variational Inference for Scalable Bayesian Tensor Factorization},
  author={Mboya, Grold Otieno},
  year={2026}
}

License

This project is licensed under the MIT License