Empirical Bayes Flow Matching for Continuous Cryo-EM Heterogeneity

NeurIPS 2026

University of Haifa

EB-FM learns a prior over 3D atom ensembles from 2D cryo-EM images

Abstract

We address the problem of learning generative priors over high-dimensional latent variables from indirect, noisy observations, with a focus on continuous heterogeneity in cryo-electron microscopy (cryo-EM). In this setting, 3D atomic structures x ∈ ℝ3Na are never directly observed and must be inferred from 2D projection images y generated by a complex, non-linear forward operator. We propose Empirical Bayes Flow Matching (EB-FM), a method that learns a flow matching prior pθ(x) directly from observations using a Monte Carlo stochastic approximation Expectation-Maximization framework. EB-FM alternates between an E-step that performs approximate posterior sampling from pθ(x | y) via guided flow-based generation while maintaining persistent latent estimates through stochastic approximation, and an M-step that updates the flow model using these stabilized latent variables as training targets. This approach eliminates the need for pretraining on clean data and naturally accommodates non-linear forward models. We validate our method on controlled inverse problems using two-moons and MNIST data, and demonstrate its effectiveness on simulated cryo-EM data with known poses, showing that EB-FM can recover continuous distributions of protein conformations directly in atomic-coordinate space.

Method

EB-FM learns a flow matching prior over atomic structures from indirect observations by iteratively inverting the forward model during training, formulated as a Monte Carlo stochastic approximation Expectation-Maximization algorithm.

Diagram of EB-FM. E-step: a flow model maps a harmonic-prior sample to a denoised structure, which is rendered and compared with the observed image to guide sampling and update the persistent latent. M-step: the flow model is trained by flow matching on the latents.
Overview of EB-FM, illustrated for the cryo-EM setting. Each observation yi is associated with a persistent latent x̃i that tracks the posterior mean. In the E-step, the current flow matching prior pθ(x) generates a posterior sample x̂i via gradient-based guidance through the differentiable forward operator A, and x̃i is updated by EMA. In the M-step, the prior is trained by conditional flow matching on the persistent latents {x̃i} in place of clean data.

Sampled Conformational Trajectories

Green is the ground-truth MD simulation; blue is prior samples from EB-FM. The samples are aligned to the ground-truth molecular-dynamics trajectories for four proteins. The prior is trained only on simulated 2D cryo-EM projections, one noisy image per conformation.

Video coming soon1a62_A_traj_sample.mp4
1a62_ATranscription termination factor Rho
Video coming soon1egw_C_traj_sample.mp4
1egw_CMyocyte-specific enhancer factor 2A
Video coming soon3rjs_A_traj_sample.mp4
3rjs_ADynein light-chain protein
Video coming soon7r79_C_traj_sample.mp4
7r79_CCalcium-binding protein

Controlled Benchmarks

Two moons

We sample 2D points zi ∈ ℝ2 and lift them into a high-dimensional space xi ∈ ℝD (D = 128) via a fixed random matrix P. For each data sample, we draw a fixed random measurement matrix Ai ∈ ℝM×D and generate noisy observations yi = 𝒜(xi) + εi, where εi ∼ 𝒩(0, σ2IM). We test two forward operators: a linear model 𝒜(xi) = Aixi and a non-linear model 𝒜(xi) = sin(Aixi). To evaluate performance across varying degrees of problem difficulty, we sweep the measurement dimension M ∈ {128, 64, 32, 16, 8, 4}, representing a transition from a fully-determined system to a highly under-determined regime (M/D = 0.03).

Grid of two-moons samples for DiEM-iso, EB-FM without persistent latents and EB-FM, for linear and sine forward operators and measurement dimensions M from 128 down to 4.
Two-moons benchmark for linear and non-linear forward operators, across decreasing measurement dimensions M. Each panel shows samples from the learned prior projected back to 2D. DiEM-iso (top) is sharp at high M but its linearized CG solver collapses by M=32 under increasing uncertainty. EB-FM without persistent latents (middle) preserves the manifold further into the under-determined regime but loses structure by M=8. EB-FM with persistent latents (bottom) is the only method that recovers the two moons at M=8.

MNIST

We follow a very similar setup on MNIST: the data are images with D = 784 pixels, and each observation is a random projection of the image to M = 100 or M = 50 measurements.

Two rows of generated handwritten digits for M=100 and for M=50 measurements per image.
Unconditional samples from the EB-FM prior. 784-dimensional MNIST observed only through random Gaussian projections. Left: M=100 measurements per image. Right: M=50. The model never sees clean digits during training, yet recovers digit structure from the indirect observations.

Cryo-EM Heterogeneity in Atomic-Coordinate Space

We apply EB-FM to recovering the distribution of 3D protein conformations from noisy 2D cryo-EM projections, directly in atomic-coordinate space. The data come from molecular-dynamics (MD) simulations: each of four proteins has an MD ensemble of 10,000 conformations, and each conformation is rendered with TEM-simulator to obtain a single simulated cryo-EM image.

For 1a62_A, 1egw_C, 3rjs_A and 7r79_C, two ground-truth states in green next to the closest prior sample in cyan and their superposition, with RMSDs from 0.511 to 3.012 angstroms.
Unconditional EB-FM samples on four proteins, compared to ground-truth MD conformations. For each protein we show two ground-truth states selected to span the conformational change in the ensemble. Green: ground-truth. Cyan: the closest unconditional sample from the learned prior. Each pair is also shown superimposed with the corresponding RMSD.

BibTeX

@misc{aibinder2026empiricalbayesflowmatching,
      title={Empirical Bayes Flow Matching for Continuous Cryo-EM Heterogeneity},
      author={Daniel Aibinder and Azmi Haider and Dan Rosenbaum},
      year={2026},
      eprint={2611.05202},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2611.05202},
}