Mixture models and flow matching improve inference in biological data

Variational Mixtures and Multi-Marginal Flow Matching: Advancing Statistical Inference with Biological Applications

Machine Learning

Summary

Complex biological data can come from systems with many different patterns mixed together, making it hard to analyze. The author develops new ways to better approximate these complicated data shapes, starting from simpler methods and moving to more advanced ones called variational mixtures and flow matching. They also create a special test case to understand these methods better. Their work helps improve data analysis in fields like spatial transcriptomics, where biological samples are layered in three-dimensional space.

What this means in practice

  • For bioinformatics teams: Analyze complex spatial gene expression data by modeling multi-modal distributions across tissue layers more accurately.
  • For machine learning engineers: Improve scalable inference algorithms for probabilistic models with expensive or unknown normalization, especially using mixture approximations.

Authors

Oskar Kviman

Abstract

In this thesis I develop methods for statistical inference when the distributions arising from complex biological systems are multi-modal, geometrically structured, and sometimes only defined up to a normalizing constant. I start from variational inference and, when analytic update equations are unavailable, move to black-box variational inference. To build intuition regarding inference challenges and the proposed methodologies, I introduce a novel unnormalized target density (the CoLN distribution) and reuse it as a controlled test case in the kappa. I then trace a trajectory of increasingly expressive approximations: ensembles evaluated with the multiple importance sampling ELBO (Paper A) and variational mixtures that automate component cooperation and exploration (Paper B). Because expressivity comes at a cost, I develop efficient mixture learning ideas, including Monte Carlo objective estimators to scale mixture learning more efficiently (Paper C). As a new result in the kappa, I overturn a three decades long misconception regarding the potential performance benefits of using mixtures in variational inference. Finally, I move from variational inference to flow matching, where I address the need for specialized treatment of interpolant learning in multi-marginal settings (Paper D). By combining insights from Papers A-D, I derive in Section 5.5 a new method: multi-marginal flow matching with mixtures of variational interpolants. I connect these methodological developments to biological applications, with special emphasis on three-dimensional spatial transcriptomics, where stacked tissue slices induce multi-modal dynamics across space.