3D-Guided Scalable Flow Matching for Generating Volumetric Tissue Spatial Transcriptomics from Serial Histology

Measuring gene activity across a whole three-dimensional organ is prohibitively expensive, because sequencing can only read tissue one slice at a time. Microscopy images are cheap to collect in dense stacks, but existing algorithms predict genetics on each slide in isolation and lose the vertical structure entirely.

We built a method that reconstructs three-dimensional gene patterns across a tissue volume by borrowing structural clues from adjacent slices.

Why this was hard

Consecutive histology cuts suffer physical distortion, variable slice angles and rotational shear, so direct spot-to-spot vertical alignment fails. Two-dimensional predictors process each slice independently, which produces erratic fluctuation and biological drift along the vertical axis. Existing three-dimensional graph networks try to connect every spot across every slice at once, so memory consumption explodes quadratically and crashes on real volumetric stacks. Raw gene counts make it worse: they are overdispersed with large spikes of exact zeros, so a standard generative noise path spends its capacity resolving mathematical artefacts instead of anatomical boundaries.

How it works

HoloTea combines flow matching with morphologically guided cross-slice conditioning. For a query spot on section \(z\), we search a candidate pool on the neighbouring sections \(z \pm 1\) using planar coordinates, then retrieve morphological matches by cosine similarity over frozen UNI pathology embeddings. Visual similarity is blended with planar distance through a temperature \(\tau\) and a spatial factor \(\beta\) to give normalised weights \(w_{ij}\), which build an adjacent expression token \(\tilde{g}_i^{adj} = \sum_{j} w_{ij} P x_{t,j}\) through a low-rank gene projection \(P\). At test time this aggregation uses the current flow state \(x_{t,j}\) of adjacent slices, so inference never touches ground-truth expression.

Those tokens are mapped onto a continuous spatial grid and passed through a lightweight ControlNet to obtain spot control representations \(u_i = h_\omega(\tilde{g}_i^{adj}, t)\), which enter selected transformer blocks as a residual modulation scaled by a time schedule \(\alpha_l(t)\). To respect the discrete nature of transcript counts, we replace the usual Gaussian start distribution with a learned zero-inflated negative binomial prior, pretrained on training expression under a negative log-likelihood loss. We freeze that network and sample start states \(x_{0,i} = \log(1 + Y_i)\), giving an aligned initial distribution before linear endpoint flow integration begins.

Scaling to whole slides needs one more idea. We use global set attention with \(m\) learned inducing tokens: the inducing slots query all \(N\) spots on a slide in a read phase, and individual spots attend back to that compact summary in a write phase. Attention cost drops to \(O(Nmd)\), linear in the number of spots. On a 52-slice human embryo volume of 327,532 spots, inference used 3.7 GB of peak GPU memory and 114.9 GFLOPs, where prior graph baselines exceeded 80 GB and failed outright. On the ST-Data breast cancer cohort we reach a spot-level correlation of 0.796 and mean squared error of 0.364.

What it does not do

The method depends on morphological continuity between slices and fails when sections are separated by large physical gaps or severe cutting tears. Once the distance between slices exceeds the range of local cell structure, cross-slice retrieval cannot find shared histological patterns and feeds misleading context into the control modules. On very small tissue fragments with sparse spatial neighbourhoods, the global inducing tokens show diminishing returns over a standard two-dimensional baseline.

Accepted as an oral at MLCB 2026. Preprint on arXiv.