CryoFM: A Flow-based Foundation Model for Cryo-EM Densities

## Abstract

Abstract Cryo-electron microscopy (cryo-EM) is a powerful technique in structural biology and drug discovery, enabling the study of biomolecules at high resolution. Significant advancements by structural biologists using cryo-EM have led to the production of over 38,626 protein density maps at various resolutions 1 1 1 The Electron Microscopy Data Bank (EMDB) is a public repository for cryo-EM volumes (https://www.ebi.ac.uk/emdb/). As of 18 September 2024, EMDB contains 38 , 626 38 626 38,626 38 , 626 entries. . However, cryo-EM data processing algorithms have yet to fully benefit from our knowledge of biomolecular density maps, with only a few recent models being data-driven but limited to specific tasks. In this study, we present CryoFM , a foundation model designed as a generative model, learning the distribution of high-quality density maps and generalizing effectively to downstream tasks. Built on flow matching, CryoFM is trained to accurately capture the prior distribution of biomolecular density maps. Furthermore, we introduce a flow posterior sampling method that leverages CryoFM as a flexible prior for several downstream tasks in cryo-EM and cryo-electron tomography (cryo-ET) without the need for fine-tuning, achieving state-of-the-art performance on most tasks and demonstrating its potential as a foundational model for broader applications in these fields.

## 1 Introduction

Cryo-electron microscopy is an important technique in structural biology and drug discovery, allowing the determination of high-resolution 3D structures of biomolecules that are difficult to study through conventional methods, offering insights into molecular mechanisms and aiding in fields like drug discovery (Nogales & Scheres, ). A major component of cryo-EM data processing involves reconstructing 3D structures from noisy 2D projections of particles. It can be formulated as an inverse problem (Bendory et al., ; Singer & Sigworth, ), which aims at recovering the signal ${\mathbf{x}}\in\mathbb{R}^{n}$ (i.e., a clean protein density) from the observation ${\mathbf{y}}\in\mathbb{R}^{m}$ (^2^22The signal ${\mathbf{x}}$ denotes a clean protein density whose shape is $\mathbb{R}^{D\times D\times D}$, where $D$ is the side length. However, the observation ${\mathbf{y}}$ is a degradation and can have any arbitary shape, so we use ${\mathbf{x}}\in\mathbb{R}^{n}$ for ease of notation but without loss of generality.). Under this formulation, following Bayesian statistics, the objective is to sample a density from the posterior $p({\mathbf{x}}|{\mathbf{y}})$, which can be factorized to $p({\mathbf{x}}|{\mathbf{y}})\propto p({\mathbf{y}}|{\mathbf{x}})p({\mathbf{x}})$. Thus, the prior distribution becomes essential for guiding the reconstruction process and improving the accuracy of the resulting structures.

Scheres () first introduced a Gaussian distribution as the prior $p({\mathbf{x}})$ in 3D reconstruction for cryo-EM, effectively functioning as a frequency-dependent low-pass filter, which has proven useful in many cases. Building on this, recent works have explored more sophisticated regularizers rather than explicitly defining a prior distribution (Punjani et al., ; Tegunov et al., ; Kimanius et al., ; Li et al., ; Schwab et al., ; Kimanius et al., ; Liu et al., ). While these approaches were initially developed for 3D reconstruction tasks, their methodologies are closely aligned with the methods for density map modification and post-processing (Jakobi et al., ; Ramírez-Aportela et al., ; Terwilliger et al., ; Kaur et al., ; Sanchez-Garcia et al., ; He et al., ). Both share the common objective of improving the quality of cryo-EM maps, whether during refinement or in post-processing. Among all these approaches, a line of methods has emerged that utilizes data-driven models to directly learn $p({\mathbf{x}}|{\mathbf{y}})$ from data, showing strong potential (Sanchez-Garcia et al., ; He et al., ; Kimanius et al., ). While powerful, these models are often designed for specific tasks, limiting their general applicability and versatility.

Figure: Figure 1: The overview of CryoFM. In the training stage, CryoFM learns a vector field $v_{t}({\mathbf{x}}_{t})$, whose corresponding probability flow generates the data distribution $p_{0}({\mathbf{x}}_{0})$ of high-quality protein densities. In the inference stage, given a degraded observation ${\mathbf{y}}$, a likelihood term $p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})$ is incorporated to convert the unconditional vector field $v_{t}(x_{t})$ to a conditional one $v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$, so that we can sample from the posterior distribution $p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})$. This enables signal restoration of the density map, resulting in improved resolution of the alpha helices in the shown case.
Refer to caption: x1.png

Foundation models leverage prior knowledge and provide versatility across a wide range of applications. In particular, they have revolutionized AI for protein structure prediction and design by utilizing vast datasets to tackle diverse tasks (Jumper et al., ; Baek et al., ; Lin et al., ; Ingraham et al., ; Krishna et al., ; Hayes et al., ; Abramson et al., ; Wang et al., ). However, experimental electron density maps, a key modality in structural biology, have largely been neglected in this space. Meanwhile, the cryo-EM field has not yet taken full advantage of foundation models, which integrate prior knowledge and offer flexibility across multiple applications. In this study, we present a foundation model, CryoFM, that directly models the prior distribution $p({\mathbf{x}})$ by learning the distribution of high-quality cryo-EM density maps through flow matching. Notably, flow matching (Lipman et al., ; Liu et al., ) is a well-established generative model which achieves state-of-the-art results on several tasks (Dao et al., ; Esser et al., ; Yim et al., ). The key ingredient of flow matching is a time-dependent vector field $v_{t}({\mathbf{x}}_{t})$, which can be employed to progressively denoise a noise vector into a data sample and is closely related to the prior term $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t})$. In CryoFM, when applied to downstream tasks, we draw a sample from the posterior distribution with its score (Ho et al., ; Song et al., ): $\nabla_{\mathbf{x}}\log p({\mathbf{x}}|{\mathbf{y}})=\nabla_{\mathbf{x}}\log p
({\mathbf{x}})+\nabla_{\mathbf{x}}\log p({\mathbf{y}}|{\mathbf{x}})$. In contrast to expert-designed prior, a data-driven prior is more powerful and expressive. Compared with models that directly learn the posterior, a decomposition of the prior and likelihood is more robust and versatile. As long as we can derive the likelihood term for a downstream task, we can combine it with the learned prior in a plug-and-play fashion (Zhang et al., ; Venkatakrishnan et al., ) to sample from the posterior.

CryoFM consists of two models: CryoFM-S, which captures fine local details at high resolution, and CryoFM-L, which focuses on the overall global shape of the biomolecule at medium to low resolution(^3^33Unless stated otherwise, CryoFM refers to CryoFM-S throughout this paper.). Fig. 1 illustrates the overview framework of CryoFM. In the training stage, we regress the vector field with high-quality density data downloaded from the EMDB database (wwPDB Consortium, ). When applied to downstream tasks, we derive a flow posterior sampling method to convert the vector field $v_{t}({\mathbf{x}}_{t})$ to a conditional one $v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$, which is then able to generate a posterior distribution $p({\mathbf{x}}|{\mathbf{y}})$. When applied to various downstream tasks, CryoFM achieves state-of-the-art results on most of the experiments, demonstrating the power and potential of CryoFM as a foundation model for solving more complex problems in cryo-EM and cryo-ET.

Our main contributions are as follows: (1) we present the first flow-based generative model as a foundation model that learns the distribution of high-quality cryo-EM density maps; (2) we derive a flow posterior sampling algorithm, enabling CryoFM to be effectively utilized as a prior for various downstream tasks; (3) in an unsupervised manner, we achieve better performance in several downstream tasks compared to DeepEMhancer (Sanchez-Garcia et al., ), EMReady (He et al., ) and spIsoNet (Liu et al., ); (4) we explore different model architectures and configurations to optimize the training of generative models on large biomolecular density maps.

## 2 Related Work

In this section, we will briefly review the mostly related work on cryo-EM and diffusion/flow matching for inverse problems. A more detailed discussion on these topics can be found in Section A.

Density modification and denoising in cryo-EM Density modification refers to the process of using known properties of the expected density in specific regions of a map to correct errors in observed cryo-EM density maps (Terwilliger et al., ). Deep learning has increasingly contributed to cryo-EM map denoising and modification, with methods divided into two categories: pretrained models and self-supervised approaches. Pretrained models, like DeepEMhancer (Sanchez-Garcia et al., ) and EMReady (He et al., ), learn the posterior $p({\mathbf{x}}|{\mathbf{y}})$ from data to recover high-frequency details. In contrast, self-supervised approaches, such as M (Tegunov et al., ) and spIsoNet (Liu et al., ), are trained on the dataset being processed, offering more robustness but take longer to process one dataset.

Missing wedge problem in cryo-ET Cryo-electron tomography (cryo-ET) is an imaging technique used to reconstruct 3D volumes of biological specimens from 2D images captured at various tilt angles (Lučić et al., ). The resulting 3D volumes, known as tomograms, provide detailed views of cellular structures, with smaller regions, called subtomograms, extracted to focus on specific structures like proteins (Wan & Briggs, ). A key challenge in cryo-ET is the missing wedge problem, caused by the limited range of tilt angles during data acquisition, which leaves a wedge-shaped region in Fourier space without information, leading to anisotropic resolution and artifacts. Traditional approaches use signal processing and regularization techniques to address the missing wedge (Goris et al., ; Deng et al., ; Yan et al., ; Zhai et al., ), while recent deep learning methods have shown promise in tackling this issue by leveraging data-driven models (Liu et al., ; Van Veen et al., ).

Ab initio modeling in cryo-EM Ab initio modeling in cryo-EM involves estimating the 3D structure of a protein from 2D particle images with unknown orientations (Crowther et al., ). One of the early approaches involved using 2D class averages, which are representative images created by aligning and averaging particles with similar views to improve the signal-to-noise ratio (SNR) (Ludtke et al., ; Voss et al., ). However, since the introduction of cryoSPARC, modern methods now use raw particle images directly for model estimation, leveraging stochastic gradient descent (SGD) to bypass the need for class averages (Punjani et al., ).

Diffusion/Flow Matching for Inverse Problem Inverse problems aim to restore the original sample from a degraded observation, such as in image super-resolution (Haris et al., ), inpainting (Yeh et al., ), deblurring (Kupyn et al., ), etc. Denoising diffusion probablistic models (DDPM) (Sohl-Dickstein et al., ; Ho et al., ) and flow-based models (Lipman et al., ; Liu et al., ) have demonstrated their effectiveness in solving inverse problems through unsupervised approaches (Kawar et al., ; Chung et al., ; Dao et al., ). Since a pretrained DDPM/Flow models the data distribution (prior distribution), it can help discover the posterior distribution given a degradation model (likelihood). Specifically, Chung et al. () developed diffusion posterior sampling for general inverse problems, removing the need for strong assumptions like the linearity of the degradation operator. Another line of research examines the solution to the inverse problem when the degradation operator is unknown (Chung et al., ; Kapon et al., ).

## 3 Preliminaries

### 3.1 Flow Matching

Flow matching (Lipman et al., ; Liu et al., ) defines the generative process of a data point ${\mathbf{x}}_{0}\in\mathbb{R}^{n}$ by progressively transforming a sample ${\mathbf{x}}_{1}\in\mathbb{R}^{n}$ from the noise distribution. The process is modeled in terms of an ordinary differential equation (ODE):

$$ $\displaystyle d{\mathbf{x}}_{t}=v_{\theta}(t,{\mathbf{x}}_{t})dt,$ (1) $$

where $v_{\theta}:[0,1]\times\mathbb{R}^{n}\to\mathbb{R}^{n}$ is a time-dependent vector field (^4^44For ease of notaion, sometimes we omit the $\theta$ and use $v_{t}({\mathbf{x}}_{t})$ for abbreviation in later sections.) parameterized by a neural network $\theta$. It denotes the velocity of a sample ${\mathbf{x}}_{t}\in\mathbb{R}^{n}$ at some intermediate timestep $t$. The vector field generates a probability density path $p_{t}({\mathbf{x}}_{t})$, which reshapes the noise distribution $p_{1}({\mathbf{x}}_{1})$ to the data distribution $p_{0}({\mathbf{x}}_{0})$. Chen et al. () directly solves Eq. (1) via a differentiable neural ODE solver. Lipman et al. () and Liu et al. () suggest that it is more efficient to learn a manually designed conditional vector field $u|{\mathbf{x}}_{0}:[0,1]\times\mathbb{R}^{n}\to\mathbb{R}^{n}$ given a data point ${\mathbf{x}}_{0}$. The principle of designing the vector field $u|{\mathbf{x}}_{0}$ is to simplify the generated flow trajectory. For example, the conditional flow trajectory can be a straight line between ${\mathbf{x}}_{0}$ and a noise sample ${\mathbf{x}}_{1}$:

$$ $\displaystyle{\mathbf{x}}_{t}|{\mathbf{x}}_{0}=(1-t){\mathbf{x}}_{0}+t{\mathbf {x}}_{1},$ (2) $$

where the corresponding vector field is $u(t,{\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\mathbf{x}}_{0}-{\mathbf{x}}_{1}$.
Finally, the (conditional) flow matching objective becomes:

$$ $\displaystyle\mathcal{L}=\|v_{\theta}(t,{\mathbf{x}}_{t})-u(t,{\mathbf{x}}_{t} |{\mathbf{x}}_{0})\|^{2}_{2}=\|v_{\theta}(t,{\mathbf{x}}_{t})-({\mathbf{x}}_{0 }-{\mathbf{x}}_{1})\|^{2}_{2}.$ $$

### 3.2 Diffusion Posterior Sampling

Diffusion posterior sampling (DPS) has emerged as a promising approach for solving inverse problems (Song et al., ; ; Chung et al., ). Given a measurement ${\mathbf{y}}\in\mathbb{R}^{m}$ derived from ${\mathbf{x}}\in\mathbb{R}^{n}$ with a degradation operator $\mathcal{A}:\mathbb{R}^{n}\to\mathbb{R}^{m}$, the goal is to sample ${\mathbf{x}}$ from the posterior $p({\mathbf{x}}|{\mathbf{y}})$. Leveraging the DDPM which learns the score of the prior distribution $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$, it is straightforward to plug a likelihood term, allowing us to compute the score of the posterior:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t}|{\mathbf{y}} )=\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})+\nabla_{{\mathbf{x}}_{ t}}\log p_{t}({\mathbf{y}}|{\mathbf{x}}_{t}).$ $$

However, $\log p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})$ is generally intractable since the tractable likelihood $\log p_{0}({\mathbf{y}}|{\mathbf{x}}_{0})$ is only defined in the data space (i.e., $t=0$). This term requires marginalizing over all possible ${\mathbf{x}}_{0}\in p_{0}({\mathbf{x}}_{0})$. Chung et al. () proposed to use a Laplace approximation of the likelihood term so that $p_{t}({\mathbf{y}}|{\mathbf{x}}_{t})\approx p_{0}({\mathbf{y}}|\hat{{\mathbf{x
}}}_{0})$. The conditional score can thus be approximated by:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t}|{\mathbf{y}} )\approx\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})+\nabla_{{\mathbf {x}}_{t}}\log p_{t}({\mathbf{y}}|\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})).$ $$

If we assume the observation distribution $p_{0}({\mathbf{y}}|{\mathbf{x}}_{0})$ is Gaussian, taking derivative to the log-probability will produce:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{y}}|{\mathbf{x}}_{t}) \propto\nabla_{{\mathbf{x}}_{t}}\|{\mathbf{y}}-\mathcal{A}\hat{{\mathbf{x}}}_{ 0}({\mathbf{x}}_{t})\|^{2}_{2}.$ $$

Putting everything together, we can approximate the score of the posterior distribution by:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t}|{\mathbf{y}}) \approx\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})+\lambda_{t}\cdot \nabla_{{\mathbf{x}}_{t}}\|{\mathbf{y}}-\mathcal{A}\hat{{\mathbf{x}}}_{0}({ \mathbf{x}}_{t})\|^{2}_{2},$ $$

where $\lambda_{t}$ is a hyper-parameter which controls the step of the likelihood term.

## 4 CryoFM

In this section, we present the implementation of CryoFM. First, we introduce the pretraining dataset in Section 4.1. Next, we illustrate the architecture of the neural network in Section 4.2. Finally, we propose a posterior sampling algorithm for the downstream tasks.

### 4.1 Pretraining Dataset

Our training dataset consists of deposited sharpened density maps from the EMDB (wwPDB Consortium, ), specifically those with: 1) a reported resolution better than $3.0$ Å, 2) structures resolved by single-particle cryo-EM, and 3) data entries that include half-maps, ensuring that resolution estimates are based on the “gold standard” Fourier shell correlation (FSC) (Henderson et al., ). We manually curated this subset by removing exceptionally large complexes, helical structures and problematic cases through visual inspection. We also excluded density maps with side lengths greater than $576$ Å. This curation resulted in a total of $3479$ density maps for model training and testing. The density maps in the selected subset were lowpass filtered to $1.5$ Å/voxel and $3$ Å/voxel for CryoFM-S and CryoFM-L, respectively. A total of $32$ density maps were selected as test set and excluded from training. For CryoFM-S, we applied random cropping to volumes of size $64^{3}$ along with random rotations for augmentation, whereas data for CryoFM-L were center cropped to volumes of size $128^{3}$ and augmented solely with random rotations.

### 4.2 Architecture

Figure: Figure 2: CryoFM’s architecture. The side length $D$ of the input ${\mathbf{x}}_{t}$ undergoes dimension reduction through down-sampling layers and is then expanded back to its original size. To minimize computational cost near the input and output, the model employs neighborhood attention (NA). Neighborhood attention only attend to a localized area, whereas global attention (GA) calculates attention across all positions.
Refer to caption: x2.png

A major challenge in applying the vanilla Transforemr architecture to 3D density ${\mathbf{x}}_{t}\in\mathbb{R}^{D^{3}}$ is the computational complexity, which is $\mathcal{O}(D^{6})$ (^5^55The computational complexity of Transformers is $\mathcal{O}(L^{2})$, where $L=D^{3}$ is the number of tokens.). CryoFM processes the 3D input data through a Transformer with a hierarchical architecture, based on HDiT (Crowson et al., ). As illustrated in Fig. 2, the hierarchical structure downsamples the spatial dimension at initial levels and upsamples it at final levels, implemented via PixelUnshuffle and PixelShuffle layers (Shi et al., ). Consequently, the spatial dimension at the middle level is significantly reduced, making the model more easily scalable. Furthermore, at the two ends of the hierarchical structure, CryoFM employs Neighborhood Attention (NA) (Hassani et al., ) since a localized attention layer is able to capture local dependencies while considerably reducing the number of tokens to attend.

Additionally, it is important to note that we normalize the data ${\mathbf{x}}_{0}$ before training CryoFM to avoid large changes in the variance of model’s output. We empirically sample some data from the dataset and calculate a mean value of $0.04$ and a standard deviation of $0.09$.

### 4.3 Flow Posterior Sampling

Given a vector field $v_{t}({\mathbf{x}}_{t})$ that generates the prior distribution $p_{0}({\mathbf{x}}_{0})$, we aim to convert it to a vector field $v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ that generates the posterior $p_{0}({\mathbf{x}}_{0}|{\mathbf{y}})$.
We firstly correlate the vector field in Eq. (1) with the score function $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ in that (Dao et al., ; Song et al., ):

$$ $\displaystyle v_{t}({\mathbf{x}}_{t})=f_{t}({\mathbf{x}}_{t})-\frac{g_{t}^{2}} {2}\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t}),$ $$

where $f_{t}$ is the drift term and $g_{t}$ is the diffusion coefficient in the forward-time SDE(^6^66An SDE of the standard form is $d{\mathbf{x}}_{t}=f_{t}({\mathbf{x}}_{t})dt+g_{t}d{\mathbf{w}}$, where $d{\mathbf{w}}$ is the Wiener process. The probability ODE which has the same marginals as the reverse of the SDE is $d{\mathbf{x}}_{t}=\left(f_{t}({\mathbf{x}}_{t})-\frac{g_{t}^{2}}{2}\nabla_{{ \mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})\right)dt$. Compared with Eq. (1), we can find the relationship between the vector field and the score function.).
Plugging the condition into both sides, we have:

$$ $\displaystyle v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})=f_{t}({\mathbf{x}}_{t}|{ \mathbf{y}})-\frac{g_{t}^{2}}{2}\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x }}_{t}|{\mathbf{y}})=v_{t}({\mathbf{x}}_{t})-\frac{g_{t}^{2}}{2}\nabla_{{ \mathbf{x}}_{t}}\log p_{t}({\mathbf{y}}|{\mathbf{x}}_{t}).$ $$

Thus, adding a likelihood term weighted by the diffusion coefficient generates a vector field that models the posterior distribution.
We follow Dao et al. () to set $f_{t}({\mathbf{x}}_{t})=-\frac{{\mathbf{x}}_{t}}{1-t}$ and $\frac{g_{t}^{2}}{2}=\frac{t}{1-t}$ for the flow trajectory defined in Eq. (2), so the conditional vector field is:

$$ $\displaystyle v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})=v_{t}({\mathbf{x}}_{t})- \frac{t}{1-t}\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{y}}|{\mathbf{x}}_{t}).$ $$

Similar to Section 3.2, we estimate the likelihood term by:

$$ $\displaystyle v_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\approx v_{t}({\mathbf{x}}_{ t})+\frac{t}{1-t}\cdot\lambda_{t}\nabla_{{\mathbf{x}}_{t}}\|{\mathbf{y}}- \mathcal{A}\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})\|^{2}_{2}.$ $$

In practice, we incorporate some tricks to avoid numerical instability. The detailed algorithm of flow posterior sampling is explained in Alg. 1. The algorithm can be adapted for different tasks; we refer the reader to Appendix E for the additional versions.

Figure: Algorithm 1 Flow Posterior Sampling

## 5 Experiments

Fourier Shell Correlation (FSC) is a widely used metric that compares two density maps in Fourier space (Harauz & van Heel, ), allowing for the assessment of alignment between the ground truth and the reconstructed map. In the experiments, to evaluate the quality of the density maps, we use three primary metrics: $\text{FSC}_{\text{AUC}}$, $\text{FSC}_{0.5}$, and Fail Rate (FR). Specifically, $\text{FSC}_{\text{AUC}}$ measures the overall correlation across all spatial frequencies, and $\text{FSC}_{0.5}$ represents the resolution of the reconstructed map at the standard 0.5 cutoff (Rosenthal & Henderson, ). The Fail Rate (FR) identifies cases where the method fails to run or produces a result that significantly deviates from the ground truth. For the FSC metrics, we only report the results for cases that did not fail. More detailed explanation can be found in Section C.1.

### 5.1 Spectral Noise Denoising

In cryo-EM reconstruction, spectral noise is the most commonly used noise model due to its ability to capture the varying noise characteristics across different spatial frequencies, which arise from factors such as the contrast transfer function (CTF) and detector imperfections.
We consider the introduction of noise in the Fourier domain, where the variance of the added Gaussian noise is frequency-dependent, with higher frequencies exhibiting greater noise variance. Given a density $\tilde{{\mathbf{V}}}\in\mathbb{C}^{D\times D\times D}$ in the Fourier domain, the degradation model $\mathcal{A}:\mathbb{C}^{D\times D\times D}\to\mathbb{C}^{D\times D\times D}$ is:

$$ $\displaystyle\mathcal{A}(\tilde{{\mathbf{V}}})=\tilde{{\mathbf{V}}}+{\mathbf{ \epsilon}},$ $$

where ${\mathbf{\epsilon}}\in\mathbb{R}^{D\times D\times D}$ is a noise volume, whose values on the spherical shell with the same radius $\nu$ are the same ($\nu$ denotes the index of a component in the frequency space):

$$ $\displaystyle\quad{\mathbf{\epsilon}}(\nu)\sim\mathcal{N}(0,\sigma_{\text{ noise}}^{2}(\nu)).$ (3) $$

In the experiment, we estimate $\sigma_{\text{noise}}^{2}(\nu)$ from the half maps, which are two independent observations with the same underlying signal with uncorrelated noise. See Section D.1 for details.

**Table 1: Results of the spectral noise denoising task, comparing CryoFM with DeepEMhancer (Sanchez-Garcia et al., ) and EMReady (He et al., ). The best attainable $\text{FSC}_{0.5}$ is $3.0$ Å since the voxel size is $1.5$ Å.**
|  | Estimated resolution at different level of noise |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | $3.2$ Å | $4.3$ Å | $6.1$ Å |  |  |  |  |  |  |
|  | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ |
| After degradation | - | $0.8876$ | $3.21$ | - | $0.5527$ | $4.96$ | - | $0.4969$ | $6.07$ |
| DeepEMhancer | $0.06$ | $0.92\pm 0.02$ | $\mathbf{3.0\pm 0.0}$ | $0.34$ | $0.66\pm 0.05$ | $\mathbf{4.3\pm 0.2}$ | $0.44$ | $\mathbf{0.62\pm 0.05}$ | $\mathbf{4.5\pm 0.3}$ |
| EMReady | $0.03$ | $0.73\pm 0.02$ | $3.4\pm 0.1$ | $0.16$ | $0.59\pm 0.08$ | $4.9\pm 0.9$ | $0.19$ | $0.55\pm 0.09$ | $5.3\pm 0.9$ |
| CryoFM | $\mathbf{0.00}$ | $\mathbf{0.95\pm 0.01}$ | $\mathbf{3.0\pm 0.0}$ | $\mathbf{0.00}$ | $\mathbf{0.68\pm 0.04}$ | $4.4\pm 0.2$ | $\mathbf{0.00}$ | $\mathbf{0.62\pm 0.04}$ | $4.6\pm 0.3$ |

Figure: Figure 3: Result of the spectral noise denoising task. Two density maps from EMDB were added spectral noise so that the estimated resolution is $4.3$ Å. The degraded density maps were filtered by relion_postprocess for visual clarity.
Refer to caption: x3.png

We manually introduce spectral noise, resulting in resolutions ranging from 3.2 Å to 15.0 Å. DeepEMhancer and EMReady, both pretrained models that directly learn the posterior $p({\mathbf{x}}|{\mathbf{y}})$, are used as baselines. CryoFM demonstrates superior robustness, with no failed cases across all experiments, while both baselines show higher fail rates, especially at medium resolutions. Additionally, as shown in Tab. 1 and Tab. 5 in Section D.1, CryoFM achieves higher $\text{FSC}_{\text{AUC}}$ and better $\text{FSC}_{0.5}$ in most cases, significantly improving degraded density maps and enhancing the resolution of helices and loops, as illustrated in Fig. 3. Notably, the FSC curves demonstrate that CryoFM successfully improves the signal across all frequencies, while the baselines fails to retain trustworthy signals in the low-frequency range.

### 5.2 Anisotropic Noise Denoising

Figure: Figure 4: Degradation operators for (i) the anisotropic noise and (ii) the missing wedge.
Refer to caption: x4.png

Anisotropic noise in cryo-EM occurs when noise distribution varies by direction, affecting some orientations more than others and leading to uneven reconstruction quality. As a result, certain orientations have lower signal-to-noise ratios than others. To approximate this degradation in a simplified manner, we amplify the spectral noise by a factor when the particle orientation falls within a specific range of angles, simulating the increased uncertainty in less-sampled orientations. The degradation model $\mathcal{A}:\mathbb{C}^{D\times D\times D}\to\mathbb{C}^{D\times D\times D}$ is given by:

$$ $\displaystyle\mathcal{A}(\tilde{{\mathbf{V}}})=\begin{cases}\tilde{{\mathbf{V} }}+{\mathbf{\epsilon}},&\text{if }\theta_{\min}\leq\theta\leq\theta_{\max}\\ \tilde{{\mathbf{V}}}+\alpha\cdot{\mathbf{\epsilon}},&\text{otherwise}\end{ cases},$ $$

where ${\mathbf{\epsilon}}\in\mathbb{R}^{D\times D\times D}$ is the same as defined in Eq. (3), and $\alpha>1$ is a scaler that increases the noise for specific orientation angles, reflecting the heightened uncertainty in those directions. $\theta_{\min}$ and $\theta_{\max}$ are the angles that controls the portion of the noise being amplified as illustrated in Fig. 4. In the experiment, we estimate $\sigma_{\text{noise}}^{2}(\nu)$ and $\alpha$ from the half maps. See Section D.2 for details.

**Table 2: Results of the anisotropic noise denoising task, comparing CryoFM with DeepEMhancer (Sanchez-Garcia et al., ) and spIsoNet (Liu et al., ).**
|  | $\theta\in[-45^{\circ},+45^{\circ}]$ | $\theta\in[-30^{\circ},+30^{\circ}]$ | $\theta\in[-15^{\circ},+15^{\circ}]$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ |
| After degradation | $-$ | $0.6623$ | $4.15$ | $-$ | $0.6324$ | $4.27$ | $-$ | $0.6111$ | $4.38$ |
| DeepEMhancer | $0.22$ | $0.80\pm 0.03$ | $3.2\pm 0.1$ | $0.22$ | $0.79\pm 0.05$ | $3.3\pm 0.2$ | $0.22$ | $0.77\pm 0.05$ | $\mathbf{3.4\pm 0.2}$ |
| spIsoNet | $\mathbf{0.00}$ | $0.65\pm 0.01$ | $4.15\pm 0.01$ | $\mathbf{0.00}$ | $0.62\pm 0.01$ | $4.27\pm 0.03$ | $\mathbf{0.00}$ | $0.61\pm 0.01$ | $4.37\pm 0.03$ |
| CryoFM | $\mathbf{0.00}$ | $\mathbf{0.88\pm 0.03}$ | $\mathbf{3.1\pm 0.1}$ | $\mathbf{0.00}$ | $\mathbf{0.84\pm 0.03}$ | $\mathbf{3.2\pm 0.1}$ | $\mathbf{0.00}$ | $\mathbf{0.81\pm 0.04}$ | $\mathbf{3.4\pm 0.2}$ |

Figure: Figure 5: Result of the anisotropic noise denoising task. Two density maps from EMDB were added anisotropic spectral noise that the estimated resolution is 4.38 Å. The degraded density maps were filtered by relion_postprocess for visual clarity.
Refer to caption: x5.png

We manually add varied anisotropic spectral noise to the test set density maps and use DeepEMhancer, spIsoNet, and CryoFM for signal restoration. As shown in Tab. 2, CryoFM outperforms both baselines on FSC metrics and has a zero fail rate. Fig. 5 demonstrates that CryoFM restores the corrupted signal without introducing significant bias, unlike DeepEMhancer, which introduces artifacts and alters the overall shape. While spIsoNet preserves signals across frequencies, it fails to improve correlation with the ground truth. This suggests that pretrained models that learn the posterior $p({\mathbf{x}}|{\mathbf{y}})$ like DeepEMhancer can struggle to retain reliable signals, and self-supervised methods like spIsoNet lack a data-driven foundation, limiting their effectiveness. Additionally, spIsoNet requires over 5 hours on a V100 GPU for one dataset, whereas CryoFM completes the task in about 30 minutes.

### 5.3 Missing Wedge Restoration

The missing wedge problem in cryo-ET arises from the limited tilt range of the electron microscope during data acquisition, causing incomplete Fourier space sampling and artifacts in the subtomograms. We simulate this effect by applying a wedge-shaped mask in the Fourier domain, removing data from unmeasured orientations. The degradation model $\mathcal{A}:\mathbb{C}^{D\times D\times D}\to\mathbb{C}^{D\times D\times D}$ is given by:

$$ $\displaystyle\mathcal{A}(\tilde{{\mathbf{V}}})=\begin{cases}\tilde{\bm{V}},& \text{if }\theta_{\min}\leq\theta\leq\theta_{\max}\\ 0,&\text{otherwise}\end{cases},$ $$

where $\theta_{\min}$ and $\theta_{\max}$ are typically set to $-60^{\circ}$ and $+60^{\circ}$ respectively, representing the common tilt angle limits in experimental setups as illustrated in Fig. 4. For posterior sampling, we slightly modify the flow posterior sampling algorithm to Alg. 2.

**Table 3: Results of the missing wedge restoration task.**
|  | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{\text{AUC}}\uparrow$<br>(Missing Region) |
| --- | --- | --- |
| After degradation | $0.80\pm 0.02$ | $0.0000$ |
| CryoFM | $0.92\pm 0.02$ | $0.76\pm 0.06$ |

Figure: Figure 6: Result of missing wedge restoration task. Two density maps from EMDB were masked in Fourier space to simulate the missing wedge effect in cryo-ET. The central slices through the maps were also shown.
Refer to caption: x6.png

Figure: Figure 7: Result of the ab initio modeling task. One synthetic data and two real data are shown.
Refer to caption: x7.png

We apply the missing wedge effect to 32 test set density maps, with results in Tab. 3 and an example in Fig. 6. CryoFM effectively restores the original signal, reducing the missing wedge artifacts, as shown by both FSC metrics and visual inspection. The FSC curve confirm a strong correlation with the ground truth, particularly in the low-frequency range.

### 5.4 Ab initio Modeling

Ab initio modeling in cryo-EM involves reconstructing a coarse density map from 2D particle projections, which serves as a reference for refinement. Here, we simplify the problem by focusing on generating a coarse density map from a few clean 2D projections obtained via upstream 2D classification. This approach mirrors earlier ab initio methods (Ludtke et al., ; Voss et al., ) before cryoSPARC’s introduction of SGD (Punjani et al., ). Given $K$ projections, there exist $K$ degradation operators, where the $k$-th degradation model is $\mathcal{A}^{(k)}:\mathbb{R}^{D\times D\times D}\to\mathbb{R}^{D\times D}$ given by:

$$ $\displaystyle\mathcal{A}^{(k)}({\bm{V}})=\mathcal{P}(\phi^{(k)},{\bm{V}}),~{}~ {}~{}k\in[1,2,\cdots,K],$ $$

where $\mathcal{P}$ is a projection operator in the real space, and $\phi^{(k)}\in\mathbb{SO}(3)\times\mathbb{R}^{2}$ is the pose of the 2D projections. In this task, $\phi^{(k)}$ is unknown and can not be pre-determined, so it must be iteratively searched to find the optimal value (Scheres, ; Punjani et al., ; Zhong et al., ), we modify the flow posterior sampling to Alg. 3. Moreover, since ab initio modeling focuses on capturing the global shape at low resolution, we use CryoFM-L for this task.

We apply CryoFM to one synthetic dataset and two real datasets from the Electron Microscopy Public Image Archive (EMPIAR). For the synthetic dataset, we use 2 clean projections, while for the real datasets, we use 8 or 6 selected class averages as the input. As shown in Fig. 7, CryoFM closely matches the ground truth at low frequencies on the synthetic dataset. For the real datasets, CryoFM produces ab initio models that achieve similar final resolutions after refinement as those generated by cryoSPARC, demonstrating its potential for 3D reconstruction in cryo-EM.

### 5.5 Ablation & Discussion

In this section, we perform ablation studies to analyze the impact of different design choices and hyper-parameters. We briefly present some conclusions here and refer the reader to Appendix F for details.

Figure: Figure 8: MSE loss on test set as a function of total training compute. The estimation of training compute is consistent with Peebles & Xie ().
Refer to caption: x8.png

Figure: Figure 10: Averaged ranking ($\downarrow$) across downstream tasks for different hyperparameters.
Refer to caption: x10.png

Moderate patchifying and downsampling parameters lead to efficient training with minimal performance degradation. As shown in Fig. 9, reducing both patchifying and downsampling parameters lowers the test loss, since less downsampling causes less information loss. However, the reduction becomes marginal while the training compute (Gflops) increases significantly. To achieve a balance between the performance and training cost, we set the patchifying parameter to $4$ and the downsampling parameter to $1$, respectively.

Data normalization enhances downstream performance across all tasks. Fig. 9 demonstrates that models trained with normalized data consistently improves the $\text{FSC}_{\text{AUC}}$ metrics compared to the one without normalization in three downstream tasks. The observation is in line with some DDPM-based work (Yim et al., ; Rombach et al., ), as a significant change in the variance of the data can make the model difficult to train (Karras et al., ).

More sampling steps boosts the performance in density restoration. Fig. 11 presents the averaged ranking across three downstream tasks for different time steps ($\#\mathrm{timesteps}$) and the maximum step sizes ($\lambda_{\text{max}}$). The main factor influencing performance is the number of time steps. A smaller $\lambda_{\text{max}}$ may reduce the impact of the likelihood term and thus degrade performance. We choose the best combination: $\lambda_{\text{max}}=5$ and $\#\mathrm{timesteps}=1000$.

Flow models trained in the voxel space converge much faster and better than the latent flow models.
We conducted experiments on training the model in the latent space (Rombach et al., ). Fig. 11 illustrates the unconditional sampling results from two latent flow models (LFM), where the high-frequency information contains more noise than signal. The observation concides with that of Crowson et al. (), where they found that latent diffusion models fail to generate fine details. We suspect that the VAE is under-trained due to a small amount of data, and leave for future work.

## 6 Conclusion

In this study, we present CryoFM, a flow matching-based foundation model that learns the prior distribution of high-quality cryo-EM densities. During inference, we derive the degradation operators for specific tasks, allowing protein densities to be sampled from the posterior distribution based on given observations. CryoFM demonstrates versatility by restoring protein densities across four distinct tasks without fine-tuning, showcasing the potential of deep generative models as a powerful prior in cryo-EM. While promising, CryoFM has limitations: though we use real data with synthetic noise for comparison, applying the method directly to real-world noisy densities remains challenging. Additionally, our work does not address reconstructing 3D densities from raw 2D particles, but we believe CryoFM can contribute to solving these complex tasks in future work.

## Appendix A Additional Related Work

### A.1 Density modification and denoising in cryo-EM

Traditional density modification methods often apply Fourier space weighting, utilizing frequency-dependent scaling to suppress noise and enhance signal (Jakobi et al., ; Ramírez-Aportela et al., ; Terwilliger et al., ; Kaur et al., ). These heuristic-based approaches, like Wiener-style deblurring (Ramírez-Aportela et al., ), tend to rely on local resolution estimates and filtering strategies but may struggle with intricate structural details due to limited priors. On the deep learning side, pretrained models like Blush (Kimanius et al., ) further refine density maps by denoising half-maps during iterative refinement. However, they can introduce hallucinated details in lower-resolution maps. In contrast, methods such as M (Tegunov et al., ) and spIsoNet (Liu et al., ) avoid external data, with M leveraging noise2noise (Moran et al., ) for denoising independent half-maps, and spIsoNet addressing anisotropic signal distributions through self-supervised learning. Though slower, these approaches tend to reduce the risk of hallucinations and offer better robustness, but their power is often limited by not being data-driven.

### A.2 Missing wedge problem in cryo-ET

Cryo-ET enables detailed visualization of macromolecular complexes and cellular structures in their native environments. The missing wedge problem occurs due to the physical limitation of tilt angles during data acquisition, typically restricted to $(-60^{\circ},+60^{\circ})$, resulting in missing data in Fourier space and artifacts such as elongation along the missing axes. Traditional methods that use signal processing techniques and regularization strategies (Goris et al., ; Deng et al., ; Yan et al., ; Zhai et al., ) are often based on heuristic assumptions and have limitations in fully recovering lost data. Recently, deep learning models have been applied to this problem, offering improved recovery of complex patterns in tomograms (Liu et al., ; Van Veen et al., ). However, these methods often focus on entire tomograms, making it difficult to incorporate specific prior knowledge of protein structures, limiting their effectiveness in subtomogram reconstructions.

### A.3 Ab initio modeling in cryo-EM

Earlier approaches to ab initio modeling relied on experimental techniques, such as image tilt pairs (Radermacher et al., ; Leschziner & Nogales, ), which provided indirect pose information, or negative stain (De Carlo & Harris, ), which improved SNR but at the cost of high-frequency detail. Computationally, 2D class averages were commonly used as input due to their higher SNR compared to raw particles, though they had limitations, particularly in fully sampling Fourier space due to the restricted range of particle orientations (Voss et al., ). While cryoSPARC’s SGD approach improved this by working directly with raw particles (Punjani et al., ), certain structural features or heterogeneity observed in 2D class averages may still be lost in the final 3D reconstruction, especially in challenging samples.

### A.4 Diffusion/Flow-based Model

Flow-based models (Lipman et al., ; Liu et al., ) and denoising diffusion probablistic models (DDPM) (Sohl-Dickstein et al., ; Ho et al., ) are two modern deep generative models. They exhibit many similarities in technical details. DDPM implements an iterative refinement process by learning to gradually denoise a sample from a normal distribution. It has achieved the state-of-the-art results on many generative tasks, including image generation (Rombach et al., ; Podell et al., ; Peebles & Xie, ), video generation (Blattmann et al., ), and molecule generation (Yim et al., ; Abramson et al., ; Ingraham et al., ; Watson et al., ; Wang et al., ), etc. Recently, diffusion models have been adopted in the cryo-EM field. Kreis et al. () traverses the latent space of cryoDRGN (Zhong et al., ) with a diffusion model, while Wang et al. () refines structures for model building by iteratively denoising the density. Flow-based models regress a vector field that generate a disired probability path. Their simple and efficient implementation enables fast learning. These models have shown success in various domains, including image generation (Esser et al., ) and moleculer generation (Yim et al., ; Bose et al., ). The DDPM objective can also be unified into flow-based models by converting it into a probability flow ODE (Song et al., ).

### A.5 Vision Transformers for Diffusion Models

Diffusion transformers (Peebles & Xie, ) have demonstrated significant scalability and generative capabilities in image-related tasks (Esser et al., ; Hoogeboom et al., ; Hatamizadeh et al., ; Zhou et al., ). In particular, HDiT (Crowson et al., ) leverages the inherent hierarchical nature of visual patterns in its model design. By integrating the characteristics of Diffusion Transformer (Peebles & Xie, ) and Hourglass transformers (Nawrot et al., ), and employing local attention mechanisms (Hassani et al., ), HDiT offers an efficient model structure suitable for training the diffusion model in the data space.

## Appendix B Additional Detailed Information of CryoFM

### B.1 Implementation Details

Fig. 12 illustrate the model architectures for input dimensions of $64^{3}$ and $128^{3}$. Fig. 13 shows the details of different attention structure we used in the transformer block. Tab. 4 details the specific model parameter configurations (aligned with the naming conventions used in Crowson et al. ()) and the training hyperparameters.

In all experiments, we employed the FairseqAdam (Ott et al., ) optimizer with a default learning rate of 1e-4, betas set to (0.9, 0.98), and a weight decay of 0.01. A linear warm-up strategy was applied during the first 2000 steps of training.

Figure: Figure 12: Overview of two 3D CryoFM model details. (a) CryoFM-S design for input shape $64^{3}$ and (b) CryoFM-L for shape $128^{3}$.
Refer to caption: x12.png

Figure: Figure 13: Illustration of QKV structure of neighborhood attention (NA) and global attention (GA).
Refer to caption: x13.png

**Table 4: Details of training and model setup.**
| Parameter | CryoFM-S | CryoFM-L |
| --- | --- | --- |
| Parameters | 335.18 M | 308.54 M |
| GFLOP/forward | 395.87 | 427.26 |
| Training Steps | 150k | 300k |
| Batch Size | 128 | 128 |
| Precision | bf16 | bf16 |
| Training Hardware | 8$\times$A100 | 8$\times$A100 |
| Patchifying | 4 | 4 |
| Levels (Local + Global Attention) | 1 + 1 | 2 + 1 |
| Depth | [4, 8] | [2, 2, 12] |
| Widths | [768, 1536] | [320, 640, 1280] |
| Attention Heads (Width / Head Dim) | [12, 24] | [5, 10, 20] |
| Attention Head Dim | 64 | 64 |
| Neighborhood Kernel Size | 7 | 7 |

### B.2 Likelihood Estimation

Given a data point ${\mathbf{x}}_{0}\in\mathbb{R}^{n}$, we estimate its likelihood by solving a probability flow ODE (Song et al., ; Lipman et al., ; Chen et al., ). We start with the continuity equation of a velocity field $v_{t}:\mathbb{R}^{n}\to\mathbb{R}^{n}$:

$$ $\displaystyle\frac{d}{dt}\log p_{t}({\mathbf{x}}_{t})+\nabla\cdot v_{t}({ \mathbf{x}}_{t})=0.$ $$

Integrating $t\in[0,1]$ gives:

$$ $\displaystyle\log p_{0}({\mathbf{x}}_{0})=\log p_{1}({\mathbf{x}}_{1})-\int_{0 }^{1}\nabla\cdot v_{t}({\mathbf{x}}_{t})dt.$ $$

Since $p_{1}$ is a standard normal distribution, $\log p_{1}({\mathbf{x}}_{1})$ can be calculated exactly. The second term in the right-hand side of the equation can be computed by solving an ODE:

$$ $\displaystyle\frac{d}{dt}\left[\begin{array}[]{cc}{\mathbf{x}}_{t}\\ f_{t}\end{array}\right]=\left[\begin{array}[]{cc}v_{t}({\mathbf{x}}_{t})\\ \nabla\cdot v_{t}({\mathbf{x}}_{t})\end{array}\right],$ $$

with initial conditions ${\mathbf{x}}_{0}$ and $f_{0}=0$. Since the divergence operator requires a lot of computational cost, we follow Grathwohl et al. () to use the Hutchinson trace estimator to get an unbiased estimate:

$$ $\displaystyle\frac{d}{dt}\left[\begin{array}[]{cc}{\mathbf{x}}_{t}\\ f_{t}\end{array}\right]=\left[\begin{array}[]{cc}v_{t}({\mathbf{x}}_{t})\\ {\mathbf{z}}^{T}\nabla v_{t}({\mathbf{x}}_{t}){\mathbf{z}}\end{array}\right],$ $$

where ${\mathbf{z}}\sim\mathcal{N}(0,{\bm{I}})$.
In practice, we use the Runge–Kutta method (Dormand & Prince, ) to solve the ODE to get $f_{1}$, and the log probability is $\log p_{0}({\mathbf{x}}_{0})=\log p_{1}({\mathbf{x}}_{1})-f_{1}$.

##### Normalization

A commonly employed trick for training flow models is to linearly transform a data point, so that the transformed data are centered around 0 0 and the standard deviation is not too off. An additional step is required to estimate the likelihood of the original data point, since the transformation changes the probability density. Given a volume density in the dataset ${\mathbf{x}}\in\mathbb{R}^{64\times 64\times 64}$, if we normalize it by ${\mathbf{y}}=({\mathbf{x}}-0.04)/0.09$ and compute the log-probability of $\log p({\mathbf{y}})$, then the probability of ${\mathbf{x}}$ can be further calculated by:

$$ $\displaystyle\log p({\mathbf{x}})$ $\displaystyle=\log p({\mathbf{y}})+\log\left(\frac{1}{0.09}\right)^{64^{3}}$ $\displaystyle=\log p({\mathbf{y}})+64^{3}\log\left(\frac{1}{0.09}\right)$ $\displaystyle\approx\log p({\mathbf{y}})+631,228$ $$

## Appendix C Additional information of the experiments

### C.1 Metrics

$\text{FSC}_{\text{AUC}}$  The area under the curve (AUC) of the FSC curve ($\text{FSC}_{\text{AUC}}$) provides an overall correlation across all spatial frequencies, offering a comprehensive view of the alignment between the ground truth and the restored density map (Harauz & van Heel, ).

$\text{FSC}_{0.5}$  The FSC resolution at the 0.5 cutoff ($\text{FSC}_{0.5}$) is a standard metric in cryo-EM, indicating the resolution of the reconstructed map relative to the ground truth (Rosenthal & Henderson, ).

Fail Rate (FR)  The fail rate is defined as the portion of the test cases where the method either fails to run or produces an $\text{FSC}_{0.5}$ between the result and the ground truth greater than 30 Å. This metric helps identify cases where the reconstruction process fails to achieve meaningful results, ensuring a more comprehensive evaluation of the method’s robustness.

Both $\text{FSC}_{\text{AUC}}$ and $\text{FSC}_{0.5}$ are reported as the mean and standard deviation across the 32 density maps in the test set, excluding failed cases.

### C.2 Spectral Noise Denoising

The results of spectral noise denoising at additional resolutions are presented in Tab. 5.

**Table 5: Additional results of the spectral noise denoising task, comparing CryoFM with DeepEMhancer (Sanchez-Garcia et al., ) and EMReady (He et al., ).**
|  | Estimated resolution at different level of noise |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | $8.5$ Å | $15.0$ Å |  |  |  |  |
|  | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ | FR$\downarrow$ | $\text{FSC}_{\text{AUC}}\uparrow$ | $\text{FSC}_{0.5}\downarrow$ |
| After degradation | - | $0.3275$ | $8.60$ | - | $0.2014$ | $15.98$ |
| DeepEMhancer | $0.13$ | $0.37\pm 0.02$ | $\mathbf{7.4\pm 0.3}$ | $0.03$ | $0.16\pm 0.04$ | $17.5\pm 2.5$ |
| EMReady | $\mathbf{0.00}$ | $0.37\pm 0.04$ | $7.7\pm 0.8$ | $0.03$ | $0.22\pm 0.02$ | $\mathbf{13.3\pm 1.6}$ |
| CryoFM | $\mathbf{0.00}$ | $\mathbf{0.38\pm 0.01}$ | $7.5\pm 0.3$ | $\mathbf{0.00}$ | $\mathbf{0.24\pm 0.01}$ | $15.0\pm 0.7$ |

## Appendix D Estimation of the degradation operators

### D.1 Spectral noise power estimation

Given two half maps from a cryo-EM reconstruction, ${\mathbf{y}}_{1}$ and ${\mathbf{y}}_{2}$ (observations), where the signal and noise power are the same and the signal and noise are uncorrelated, the obeservation models for ${\mathbf{y}}_{1}$ and ${\mathbf{y}}_{2}$ are:

$$ $\displaystyle\tilde{{\mathbf{y}}}_{1}=\tilde{{\mathbf{V}}}+{\mathbf{\epsilon}} _{1},\quad{\mathbf{\epsilon}}_{1}(\nu)\sim\mathcal{N}(0,\sigma_{\text{noise}}^ {2}(\nu)),$ $$

$$ $\displaystyle\tilde{{\mathbf{y}}}_{2}=\tilde{{\mathbf{V}}}+{\mathbf{\epsilon}} _{2},\quad{\mathbf{\epsilon}}_{2}(\nu)\sim\mathcal{N}(0,\sigma_{\text{noise}}^ {2}(\nu)).$ $$

The so-called “Gold-Standard” Fourier Shell Correlation (GSFSC) is the FSC between two half maps, which is defined as:

$$ $\displaystyle\text{GSFSC}(\nu)\triangleq\frac{\tilde{{\mathbf{y}}}_{1}(\nu) \cdot\tilde{{\mathbf{y}}}_{2}(\nu)}{\|\tilde{{\mathbf{y}}}_{1}(\nu)\|\cdot\| \tilde{{\mathbf{y}}}_{2}(\nu)\|},$ $$

where $\tilde{{\mathbf{y}}}(\nu)$ represents the component of $\tilde{{\mathbf{y}}}$ at frequency $\nu$ (i.e., all values on a spherical shell), and $\|\tilde{{\mathbf{y}}}(\nu)\|^{2}$ is the radial power spectrum.

From the GSFSC, the signal-to-noise ratio (SNR) can be computed as (Rosenthal & Henderson, ):

$$ $\displaystyle\text{SNR}(\nu)=\frac{\mathbb{E}[\sigma_{\text{signal}}^{2}(\nu)] }{\mathbb{E}[\sigma_{\text{noise}}^{2}(\nu)]}=\frac{\text{GSFSC}(\nu)}{1-\text {GSFSC}(\nu)}.$ $$

Since ${\mathbf{y}}_{1}$ and ${\mathbf{y}}_{2}$ share the same signal and noise at each frequency, and the signal $\tilde{{\mathbf{V}}}$ and the noise ${\mathbf{\epsilon}}$ are uncorrelated, we have:

$$ $\displaystyle\mathbb{E}[\|\tilde{{\mathbf{y}}}(\nu)\|^{2}]=\mathbb{E}[\|\tilde {{\mathbf{V}}}(\nu)+{\mathbf{\epsilon}}(\nu)\|^{2}]=\mathbb{E}[\sigma_{\text{ signal}}^{2}(\nu)]+\mathbb{E}[\sigma_{\text{noise}}^{2}(\nu)].$ $$

Thus,

$$ $\displaystyle\mathbb{E}[\|\tilde{{\mathbf{y}}}_{1}(\nu)\|^{2}]=\mathbb{E}[\| \tilde{{\mathbf{y}}}_{2}(\nu)\|^{2}]=\mathbb{E}[\sigma_{\text{signal}}^{2}(\nu )]+\mathbb{E}[\sigma_{\text{noise}}^{2}(\nu)].$ $$

Given the known SNR, we can derive:

$$ $\displaystyle\mathbb{E}[\sigma_{\text{noise}}^{2}(\nu)]=\frac{\mathbb{E}[\| \tilde{{\mathbf{y}}}_{1}(\nu)\|^{2}]}{1+\text{SNR}(\nu)}=\frac{\mathbb{E}[\| \tilde{{\mathbf{y}}}_{2}(\nu)\|^{2}]}{1+\text{SNR}(\nu)},$ $$

and

$$ $\displaystyle\mathbb{E}[\sigma_{\text{signal}}^{2}(\nu)]=\frac{\mathbb{E}[\| \tilde{{\mathbf{y}}}_{1}(\nu)\|^{2}]\cdot\text{SNR}(\nu)}{1+\text{SNR}(\nu)}= \frac{\mathbb{E}[\|\tilde{{\mathbf{y}}}_{2}(\nu)\|^{2}]\cdot\text{SNR}(\nu)}{1 +\text{SNR}(\nu)}.$ $$

### D.2 Anisotropic noise power estimation

Anisotropic noise means that the noise is not uniform across all directions; in certain directions, the noise power is stronger or weaker than in others. In the context of spectral noise, anisotropic noise implies that the variance of the noise depends on the direction in Fourier space. We denote by $\xi=(\phi,\theta)$ the anisotropic angular components on the sphere, and by $\nu$ the frequency component, corresponding to the radial component in spherical coordinates. We assume that the signal power $\sigma^{2}_{\text{signal}}(\nu)$ is isotropic, meaning $\sigma^{2}_{\text{signal}}(\xi,\nu)$ is constant for different values of $\xi$.

The noise power averaged over all directions for a given frequency $\nu$ is denoted as:

$$ $\displaystyle h(\nu)=\frac{1}{N_{\nu}}\sum_{\xi}\sigma^{2}_{\text{noise}}(\xi, \nu),$ $$

where $N_{\nu}$ represents the number of voxels in the frequency shell at $\nu$.

Our goal is to estimate the anisotropic noise power, $\sigma^{2}_{\text{noise}}(\xi,\nu)$, in spherical coordinates. We assume that for a given angle $\xi$, the noise power $\sigma^{2}_{\text{noise}}(\xi,\nu)$ can be expressed as:

$$ $\displaystyle\sigma^{2}_{\text{noise}}(\xi,\nu)=\alpha(\xi,\nu)\cdot\sigma^{2} _{\text{noise}}(\xi_{0},\nu),$ $$

where $\xi_{0}$ represents an arbitrary reference direction in spherical coordinates, serving as the baseline for the noise power $\sigma^{2}_{\text{noise}}(\xi,\nu)$, with $\alpha(\xi,\nu)$ adjusting the noise power in other directions. In practice, $\alpha(\xi,\nu)$ can be inferred from the number of particles contributing to each pose during the reconstruction process.

Substituting the anisotropic noise power expression, we can derive:

$$ $\displaystyle h(\nu)=\frac{1}{N_{\nu}}\sum_{\xi}\alpha(\xi,\nu)\cdot\sigma^{2} _{\text{noise}}(\xi_{0},\nu).$ $$

Simplifying, we obtain:

$$ $\displaystyle h(\nu)=\frac{\sigma^{2}_{\text{noise}}(\xi_{0},\nu)}{N_{\nu}} \sum_{\xi}\alpha(\xi,\nu),$ $$

where $h(\nu)$ is the average noise power derived from the the half-map GSFSC in Section D.1.

Thus, we can solve for $\sigma^{2}_{\text{noise}}(\xi_{0},\nu)$ as:

$$ $\displaystyle\sigma^{2}_{\text{noise}}(\xi_{0},\nu)=\frac{h(\nu)\cdot N_{\nu}} {\sum_{\xi}\alpha(\xi,\nu)}.$ $$

Finally, the anisotropic noise power for any angle $\xi$ is given by:

$$ $\displaystyle\sigma^{2}_{\text{noise}}(\xi,\nu)=\alpha(\xi,\nu)\cdot\frac{h( \nu)\cdot N_{\nu}}{\sum_{\xi}\alpha(\xi,\nu)}.$ $$

## Appendix E Additional Algorithms

In this section, we present algorithms for addressing different downstream tasks within the flow posterior sampling framework. For spectral noise and anisotropic noise denoising, it is natural to apply Alg. 1 in the main text since the loss function is well-defined. However, regarding missing wedge restoration and ab initio modeling, we can make slight modifications to the algorithm to better solve these problems.

### E.1 Missing Wedge Restoration

For certain tasks, such as missing wedge restoration, it is not necessary to compute the gradient with respect to the loss function between the observation ${\mathbf{y}}$ and the degraded signal $\mathcal{A}{\mathbf{x}}$, since we can directly attain the optimum in an analytical form.
The missing wedge restoration is an in-painting task, where a noiseless part of the signal is already observed. $\mathcal{A}$ is an operator extracting a part of the signal, hence the inverse of $\mathcal{A}$ just fills the observed part into the signal. In Alg. 2, we formulate it as the pseudo-inverse operator $\mathcal{A}^{\dagger}:\mathbb{R}^{m}\to\mathbb{R}^{n}$ (Kawar et al., ; ).

Figure: Algorithm 2 Flow Posterior Sampling for the Missing Wedge Problem

### E.2 Ab initio Modeling

For ab initio modeling, the $k$-th degraded volume (we will name it the projection later) is $\mathcal{A}^{(k)}({\bm{V}})=\mathcal{P}(\phi^{(k)},{\bm{V}})$, where $\phi^{(k)}\in\mathbb{SO}(3)\times\mathbb{R}^{2}$ is the pose.
The posterior sampling of ab initio modeling is more complex than the previous tasks in that the pose $\phi^{(k)}$ can not be easily pre-determined. The setting is closely related to blind diffusion posterior sampling (Chung et al., ; Kapon et al., ; Laroche et al., ; Gan et al., ), where the degradation operator is unknown. We find $\phi^{(k)}$ in the sampling process, by minimizing the discrepancy between the projection and the observation ${\bm{y}}^{(k)}$ using a coarse-to-fine strategy (Zhong et al., ). The discrepancy is defined by the negative correlation between a projection ${\mathbf{p}}$ and an observation ${\mathbf{y}}$:

$$ $\displaystyle\psi_{\mathrm{cor}}({\mathbf{p}},{\mathbf{y}})=-\frac{{\mathbf{p} }\cdot{\mathbf{y}}}{\|{\mathbf{p}}\|\cdot\|{\mathbf{y}}\|}.$ $$

When computing the likelihood, we further introduce two variants of the loss function to make the sampling process more robust, since this setting is highly ill-posed compared to others. The first one applies a stop gradient operator (denoted by “$\mathrm{sg}$”) to the norm of the projection, where the stop gradient operator is identity at forward computation time but has zero partial derivatives (Van Den Oord et al., ; Roy et al., ; Jang et al., ):

$$ $\displaystyle\psi_{\mathrm{sg}}({\mathbf{p}},{\mathbf{y}})=-\frac{{\mathbf{p}} \cdot{\mathbf{y}}}{\mathrm{sg}(\|{\mathbf{p}}\|\cdot\|{\mathbf{y}}\|)}.$ $$

This variant does not constrain the norm of the projection, as long as it has a large correlation with the observation. The second one penalizes additional noises in the projection so that it has a value of 0 0 in the background area of the observation ${\mathbf{y}}$. The background area is defined by a set of indices $\mathcal{B}(q,{\mathbf{y}})$ where the corresponding value is smaller the $q$-th percentile of ${\mathbf{y}}$.

$$ $\displaystyle\psi_{\mathrm{bg}(q)}({\mathbf{p}},{\mathbf{y}})=\sum_{i\in \mathcal{B}(q,{\mathbf{y}})}{\mathbf{p}}_{i}^{2}.$ $$

The final loss function $\mathcal{L}$ can be a combination of $\psi_{\mathrm{cor}}$, $\psi_{\mathrm{sg}}$ and $\psi_{\mathrm{bg}(q)}$. We tune the parameter on different datasets (Tab. 6) by examining the human preference for the reconstruction results. A full alorithm can be found in Alg. 3.

Figure: Algorithm 3 Flow Posterior Sampling for Ab-initio Modeling

**Table 6: Hyperparameters for ab initio modeling**
| EMPIAR-ID | stop gradient | $q$ |
| --- | --- | --- |
| $10345$ (synthetic) | no | $0.7$ |
| $10028$ | yes | $0.0$ |
| $10827$ | no | $0.7$ |

## Appendix F Additional Ablation Study

### F.1 Data Augmentation

Protein density data are three-dimensional data characterized by large size differences and lack of canonical orientation. To train the model efficiently, we need to crop the data to the same size. Additionally, to expose the model to enough data with various orientations, we also need to apply appropriate rotational augmentation to the data. In this section, we discuss different data augmentation strategies and how they affect the training and the performance on the downstream tasks.
We implemente four different augmentation strategies, as shown in Fig. 14:
center crop: Only crops a patch from the center of the original data for training;
random crop: Randomly crops patches from the original data;
rot24: Achieves data rotation by swapping axes order and flipping the data (resulting in 24 possible orientations);
random rotation: Randomly rotates data in the $\mathrm{SO}(3)$ space.

As seen in Fig. 14 (a), all strategies enabled the model’s loss to converge stably on the training set. However, Fig. 14 (b) and (c) indicate that the weakest augmentation strategy, center crop, leads to some overfitting on the test set. From the estimated likelihood of the data in the training and test set in Fig. 14 (d), model exhibits similar behavior across different strategies. From the downstream task performance shown in Fig. 14 (e), the differences among the strategies were negligible. Given that our modeling targets a more holistic $p({\mathbf{x}})$, we choose random crop combined with random rotation as our augmentation strategy.

Figure: Figure 14: Results on different data augmentation strategies.
Refer to caption: x14.png

### F.2 Patchifying & Downsampling Factor

For three-dimensional data, the number of tokens is proportional to the cube of the edge length $D$ of the data. Applying patchification to the input and down-sampling within the model significantly impacts the number of tokens the transformer block needs to process, thereby influencing the model’s efficiency. In this section, we conducted ablation experiments on the number of patchifying and down-sampling operations. As shown in Tab. 7, the default configuration we used is p4-d1, where p4 represents a patchifying factor of 4, meaning a $4\times 4\times 4$ patch in the input is flattened and used as a token vector, and d1 represents that the data is down-sampled once in the model. We compare configurations with p $=2,4,8$ and d = $0,1,2$.

In Fig. 15, we present the loss curves on the test set during training for different model variants. Given the substantial differences in the computational costs among these model variants, the horizontal axis represents the training compute in Gflops. The estimation of training compute is: model Gflops $\cdot$ batch size $\cdot$ training steps $\cdot$ 3 (Peebles & Xie, ). We observe that the models with higher computational costs achieve lower loss levels, indicating the potential for scalability. On the other hand, as shown in Tab. 7, the inference speeds of the p2 and d0 models are exceedingly slow. Therefore, we selected the p4-d1 configuration with a comprehensive consideration of model performance and efficiency in downstream applications.

**Table 7: Different model configurations.**
|  | Default | Patchifying variants | Down-sampling variants |  |  |
| --- | --- | --- | --- | --- | --- |
| Short name | p4-d1 | p2-d1 | p8-d1 | p4-d0 | p4-d2 |
| Patching | 4 | 2 | 8 | 4 | 4 |
| Depth | [4, 8] | [4, 8] | [4, 8] | [12] | [2, 2, 8] |
| Width | [768, 1536] | [768, 1536] | [768, 1536] | [1536] | [384, 768, 1536] |
| Parameters | 335.18 M | 335.10 M | 335.87 M | 377.75 M | 316.06 M |
| Thru. (FPS) | 10.23 | 1.19 | 40.61 | 2.67 | 28.52 |

Figure: Figure 15: MSE loss on test set as a function of total training compute.
Refer to caption: x15.png

### F.3 Hyperparameters for Posterior Sampling

Posterior sampling process involves continuously moving forward and correcting in the directions of prior $p({\mathbf{x}})$ and the likelihood $p({\mathbf{x}}|{\mathbf{y}})$. In this section, we present ablation experiments on the two parameters that significantly impact the effectiveness of the posterior sampling algorithm: the number of total timesteps, and $\lambda_{\text{max}}$ which controls the extent to which $x_{t}$ is modified based on the observed data.

Tab. 8 and Tab. 9 present the ablation studies on these two parameters across three downstream tasks. The results show that for a fixed $\lambda_{\text{max}}$, increasing the number of sampling timesteps improves quality, with the best performance at $1000$ timesteps. Moreover, a $\lambda_{\text{max}}$ of $5$ yields optimal or near-optimal results. Therefore, we defaulted to a $\lambda_{\text{max}}$ of $5$ and $1000$ timesteps for the posterior sampling parameters in our main results.

**Table 8: Results with different posterior sampling parameters on spectral noise and anisotropic noise denoising task.**
|  |  | Spectral Noise Denoising | Anistropic Noise Denoising |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | $3.2$ Å | $4.3$ Å | $6.1$ Å | $8.5$ Å | $15.0$ Å | $\pm 45^{\circ}$ | $\pm 30^{\circ}$ | $\pm 15^{\circ}$ |
| After degradation | $0.8874$ | $0.5526$ | $0.4968$ | $0.3278$ | $0.2000$ | $0.6626$ | $0.6326$ | $0.6113$ |  |
| $\lambda_{\text{max}}$ | # timesteps |  |  |  |  |  |  |  |  |
| $1.0$ | $50$ | $0.2511$ | $0.2573$ | $0.1811$ | $0.1736$ | $0.1349$ | $0.1994$ | $0.1638$ | $0.1391$ |
| $1.0$ | $200$ | $0.7301$ | $0.5525$ | $0.4985$ | $0.3855$ | $0.2476$ | $0.5965$ | $0.4836$ | $0.4126$ |
| $1.0$ | $1000$ | $0.9512$ | $0.6771$ | $0.6082$ | $0.3928$ | $0.2460$ | $0.8745$ | $0.8191$ | $0.7741$ |
| $5.0$ | $50$ | $0.4618$ | $0.4140$ | $0.3577$ | $0.3140$ | $0.2156$ | $0.3778$ | $0.8191$ | $0.2576$ |
| $5.0$ | $200$ | $0.8994$ | $0.6417$ | $0.5808$ | $\mathbf{0.4115}$ | $\mathbf{0.2678}$ | $0.7923$ | $0.6694$ | $0.5714$ |
| $5.0$ | $1000$ | $\mathbf{0.9545}$ | $\mathbf{0.6843}$ | $\mathbf{0.6191}$ | $0.3794$ | $0.2313$ | $\mathbf{0.8879}$ | $\mathbf{0.8453}$ | $\mathbf{0.8053}$ |
| $10.0$ | $50$ | $0.4989$ | $0.4312$ | $0.3904$ | $0.3381$ | $0.2249$ | $0.4012$ | $0.3027$ | $0.2568$ |
| $10.0$ | $200$ | $0.9025$ | $0.6397$ | $0.5798$ | $\underline{0.4103}$ | $\underline{0.2631}$ | $0.7926$ | $0.6713$ | $0.5503$ |
| $10.0$ | $1000$ | $\underline{0.9536}$ | $\underline{0.6649}$ | $\underline{0.6183}$ | $0.3752$ | $0.2314$ | $\underline{0.8780}$ | $\underline{0.8329}$ | $\underline{0.7926}$ |

**Table 9: Results with different sampling steps on missing wedge restoration task.**
|  | $\text{FSC}_{\text{AUC}}\uparrow$ |
| --- | --- |
| After degradation | $0.7988$ |
| Timesteps = $50$ | $0.9295$ |
| Timesteps = $200$ | $0.9295$ |
| Timesteps = $1000$ | $\mathbf{0.9310}$ |

### F.4 Latent Flow Models

In this section, we discuss the choice of the input space of the flow-matching model: should we learn in the voxel space or the latent space? Most of the early works (Goodfellow et al., ; Van Den Oord et al., ; Brock, ; Dhariwal & Nichol, ) for generative modeling learned the data distribution directly on the original image space. Rombach et al. () pointed out that learning diffusion models on the latent space can generate high-quality images with much less computational cost. The key is to learn an variational auto-encoder (VAE) (Kingma, ) which provides a low-dimensional representation that is equivalent to the original data space (Rombach et al., ; Podell et al., ). Then, a Diffusion Transformer (DiT) (Peebles & Xie, ) is trained on the latent codes of the VAE.

We have trained several latent flow models (LFMs). An LFM-$d$ model is comprised of a VAE-$d$ and a DiT, where $d$ denotes the downsampling factor of the VAE. A larger downsampling factor $d$ leads to greater information compression and increased information loss in the VAE, while it may ease the training of the DiT with smaller size of latent codes.

Fig. 11 shows the results of unconditional sampling from two latent flow models (LFM): LFM-$4$ and LFM-$8$, where the high-frequency information contains more noise, and secondary structures such as $\alpha$-helicies can not be observed. From the preliminary result, we concluded that: For protein density modeling, the flow models trained in the voxel space converge much faster and better than the latent flow models.

Our observation concides with that of Crowson et al. (), where they found that latent diffusion models may fail to generate fine details. We attribute the failure of latent flow models to the weak representational power of the learned variational auto-encoder (VAE) (Kingma, ). The protein density dataset for training the VAE is much smaller than the large-scale dataset in image synthesis tasks (Rombach et al., ; Podell et al., ). The low-dimensional space of the VAE is poorly regularized and can be sensitive to perturbations. A small error in the latent space may lead to large peceptual loss in the voxel space, especially for cryo-EM tasks which is highly sensitive to fine details.

#### F.4.1 Architecture of LFM

It is somewhat difficult to conduct a rigorous comparison between latent-space models and voxel-space models as their architectures may differ in several aspects. For the completeness of this work, we will present the results of these models. However, for the sake of simplicity in presentation, we only show two representative settings, and the conclusion is generally consistent across various settings (such as parameter size, model architecture, batch size, etc.).

The procedure of learning a latent flow model whose downsampling factor is $d$ (LFM-$d$) can be divided into two parts:

- •
a VAE-$d$ which compresses a protein density ${\mathbf{V}}\in\mathbb{R}^{D\times D\times D}$ to a latent code ${\mathbf{z}}\in\mathbb{R}^{\frac{D}{d}\times\frac{D}{d}\times\frac{D}{d}}$, $d$ is the compressing factor. The compressing factor governs the balance between perceptual compression and distortion (Blau & Michaeli, ; Rombach et al., ). We employed the publicly available code of Stable Diffusion (^7^77https://github.com/CompVis/stable-diffusion) by replacing 2D convolution layers with 3D layers. We conducted experiments with two settings: $d=4$ and $d=8$.
- •
a DiT (Diffusion Transformer) (Peebles & Xie, ) that learns the distribution of the latent codes. We experimented with a standard DiT architecture by replacing the 2D patchifying module with a 3D module (^8^88https://github.com/huggingface/diffusers). The DiT has $12$ layers, each of which has $1024$ dimensions, and the patchify factor is set to $2$.

Section F.4.2 shows that VAE-$8$ performs significantly worse in reconstrution tasks since it is over-compressed and causes loss of information. However, a larger downsampling factor produces smaller number of latent codes, which may ease the training of DiT. Therefore, we trained two DiTs on the latent codes of two VAEs. For each setting, we trained the DiT for $300$k steps.

#### F.4.2 Evaluation of VAE

Given a volume ${\mathbf{V}}\in\mathbb{R}^{D\times D\times D}$, we evaluate the reconstruction quality of VAE through auto-encoding Fourier shell correlation (AE-FSC) and auto-encoding L-1 loss between the input density and the reconstruction, defined by:

$$ AE-FSC $\displaystyle\triangleq\text{FSC}({\mathbf{V}},\text{VAE}({\mathbf{V}}))$ AE-L1 $\displaystyle\triangleq\|{\mathbf{V}}-\text{VAE}({\mathbf{V}})\|_{1}$ $$

Tab. 10 shows that VAE-$4$ outperforms VAE-$8$ in terms of reconstruction quality, since VAE-$8$ learns to compress a grid of $8\times 8\times 8=512$ voxels into a latent code, which is $8$ times harder than VAE-$4$ that operates on $4\times 4\times 4=64$ voxels.

**Table 10: Reconstruction quality of VAE**
| Model | AE-FSC ($\uparrow$) | AE-L1 ($\downarrow$) | Training Steps |
| --- | --- | --- | --- |
| VAE-$4$ | $0.9821$ | $0.0028$ | $588,480$ |
| VAE-$8$ | $0.8537$ | $0.0049$ | $1,097,872$ |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) Dataset for pretraining: You state that you remove "problematic cases" from the pretraining set. What are these problematic cases? Why are large maps with a length greater than 576 \\AA removed from the pretraining set? Could you please provide a complete list of the maps used for pretraining?
2) You only show 4 examples for the ab initio modeling task, one sample for each problem setting. Can you comment on the success rate? Are the shown approximate posterior samples, the best human selected ones or something we can expect to see on average when using the proposed approach on other data?
3) Lines 375/376: From Fig. 5 it is unclear why and to what extent DeepEMhancer "introduces artifacts and alters the overall shape." The FSC is very good and the visual appearance is similar to your result.
4) How fast is CryoFM compared to the other baselines? 

__Additional feedback__

1) The mathematical explanation of CryoFM should be improved. For example, Eq. (2) states $\\mathbf{x}_t | \\mathbf{x}_0 = (1-t)\\mathbf{x}_0 + t \\mathbf{x}_1$ implying $u(t, \\mathbf{x}_t | \\mathbf{x}_0) =  \frac{d}{dt} \\mathbf{x}_t|\\mathbf{x}_0 = \\mathbf{x}_1 - \\mathbf{x}_0$ and not $\\mathbf{x}_0 - \\mathbf{x}_1$. These sign errors occur at various places throughout the manuscript. 

2) Another problem can be found in line 172 which should read 
    $$\\mathcal{L} = \\mathbb{E}_{t, p_0(\\mathbf{x}_0), p_1(\\mathbf{x}_1)} \\left[ \\| v _{\\theta} (t, \\mathbf{x}_t) - (\\mathbf{x}_1 - \\mathbf{x}_0) \\|^2_2 \\right]$$
    In addition to the sign error, your loss is oversimplified and needs to indicate the expectation over $t \\sim \\mathcal{U}[0,1]$, $\\mathbf{x}_0 \\sim p_0(\\mathbf{x}_0)$, and  $\\mathbf{x}_1 \\sim p_1(\\mathbf{x}_1)$.

3) An error/typo can also be spotted in line 196: 
    $$\\nabla_{\\mathbf{x}_t} \\log p (\\mathbf{x}_t | \\mathbf{y}) \\approx \\nabla _{\\mathbf{x}_t} \\log p (\\mathbf{x}_t) - \\lambda_t 
\\nabla _{\\mathbf{x}_t}  \\| \\mathbf{y} - \\mathcal{A} \\hat{\\mathbf{x}}_0 (\\mathbf{x}_t) \\|_2^2$$
  Your version __adds__ the likelihood term rather than subtracting it, which only makes sense if $\\lambda_t \\leq 0$, but in Algorithms 1 and 3 we have $\lambda_t \\geq 0$.

4) ODE solvers require a time increment $\Delta t$ (explicitly or implicitly) which is missing in your version of Algorithm 1. So Algorithm 1 should read similar to

    \\begin{aligned}
    &\\mathbf{x}_1 \\leftarrow \\mathcal{N}(0, \\mathbf{I}) \\\\
    &\\textcolor{red}{\\Delta t \\leftarrow 1/N} \\\\
    &\\textbf{for }t \\in [1, 1- \\Delta t, 1 - 2\\Delta t, \\ldots, \\textcolor{red}{\\Delta t}] \\textbf{ do} \\\\
    &\\qquad \\mathbf{x}' _{t - \\textcolor{red}{\\Delta t}} \\leftarrow \\mathbf{x} _t  - v _{\\theta}(t, \\mathbf{x} _t) \\textcolor{red}{\\Delta t} \\\\
    &\\qquad l(\mathbf{x} _t) \\leftarrow \\| \\mathbf{y} - \\mathcal{A} \\hat{\\mathbf{x}} _0(\\mathbf{x} _t) \\|^2_2 \\\\
    &\\qquad \\mathbf{g} \\leftarrow \\frac{\\nabla _{\\mathbf{x} _t} l(\\mathbf{x} _t) }{ \\|\nabla _{\mathbf{x}_t } l( \mathbf{x}_t ) \\|_2} \\\\
    &\\qquad \\lambda _t \\leftarrow \\min \\{ \\lambda _{\\max}, t / (1-t) \\} \\\\
    &\\qquad \\mathbf{x} _{ t - \textcolor{red}{\\Delta t} } \\leftarrow  \\mathbf{x}' _{ t - \\textcolor{red}{\\Delta t} }  \\textcolor{red}{ - } \\lambda_t \\mathbf{g} \\\\
    & \\textbf{end for} \\\\ 
    & \\textbf{return } \\mathbf{x} _0
    \\end{aligned}

	Some additional comments on Algorithm 1: 
	* What does $\\mathbf{x} _{t-1}$ mean? Is $t$ an index or the real-valued time $t \\in [0, 1]$? In $\\mathbf{x} _{t-1}$ it seems to be an index, in $\\mathbf{x}_0, \\mathbf{x}_1$ it indicates the time.
	* In your version, we have $t = 0$ in the last integration step such that $\\mathbf{x}_{-\\Delta t}$ is computed. But it does not make sense to decrease $t$ further after reaching the data distribution at $t = 0$. So the loop should end at $t = \\Delta t$. 
	* Increasing the loss $l(\\mathbf{x}_t)$ does not make sense (most likely a typo or subsequent error of line 196). 
	* In Algorithm 3, the same errors occur as in Algorithm 1.

5) Again, please clarify the meaning of $t-1$ in $\mathbf{x}_{t-1}$ and avoid division by 0 in $1/(Nt)$ in the (unnecessary) last iteration.

6) Symbol $\theta$ is overloaded: In most cases, $\theta$ indicates the parameters of the foundation model. But in the discussion of the missing wedge dataset it indicates the tilt angle. 

7) It is confusing to use multiple versions of the flow. We see $v_t(\mathbf x_t)$, $v_\theta(t, \mathbf x_t)$ throughout the paper. Moreover, since the flow is a vector field, it should be bold face according to the ICLR guidelines. 

# Minor comments

1) The specific number of 38,626 cryo-EM structures mentioned in the abstract will most likely be outdated at the time when your paper is published.

2) In line 095, you write "... which is then able to generate a posterior distribution $p(x|y)$." This is incorrect: Sampling with $v_t$ can only __approximate__ the posterior distribution.

3) Typos/grammar: 
	* Page 7: "scaler"
	* Line 525: "We ... leave for future work."
soundness: 3
strengths: * CryoFM is a foundation model for density maps that can be combined with different types of data. 
* CryoFM achieves state-of-the art performance on a number of reconstruction tasks and is demonstrated to be superior to other baselines.
* CryoFM seems to be faster than other baselines.
summary: The paper introduces CryoFM, a foundation model for the distribution of high-quality cryo-EM density maps. CryoFM consists of two components: CryoFM-S models local details at high resolution, CryoFM-L models the overall shape at medium to low resolution. Both models are trained with data from EMDB by using flow matching. CryoFM can be combined with experimental data by using a flow posterior sampling method introduced by the authors. The impact of the foundation model is demonstrated for various reconstruction tasks (density modification, denoising, missing wedge problem, ab initio reconstruction) and in a number of ablation studies.
weaknesses: * The mathematical derivation of the algorithm seems to be a bit sloppy. 
* Tests with experimental data are missing. 
* No source code is provided (it would also help clarify some of the mathematical shortcomings of the paper). 

__Recommendation__

I recommend a weak accept. The results are quite convincing, although a systematic benchmark with more examples would be desirable. It is unclear how the experimental density maps were selected for pretraining. It would be desirable if the authors reported a list of all PDB entries that went into pretraining. That way it is unclear whether some of the complexes that were used to illustrate the power of CryoFM are already part of the training data. Also the mathematical treatment of the method is sketchy and at times inconsistent.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: - In Section 5.4, the authors show that CryoFM can be used for ab-initio modelling where a low-resolution density map is estimated from 2D projections. Can CryoFM also be used to estimate high resolution density maps from projections? Can CryoFM be for example incorporated in the "Refinement" step (done here with cryoSPARC)?


Note that I also raised some questions in the "Weaknesses" sections. They seemed more appropriate there.
soundness: 3
strengths: - The idea of learning a prior from data is promising. Work in other domains has shown that priors learned on large datasets can outperform classical, hand-crafted priors that, for example, model assumptions like smoothness or sparsity.
- The authors use the learned prior to solve four relevant inverse problems: two 3D denoising problems, the missing wedge problem (which deals with estimating missing data in the Fourier domain), and ab-initio modeling, which is the problem of estimating a 3D density from 2D projections of Cryo-EM.
- For all four inverse problems, the CryoFM-based approaches perform as well as or better than the baselines (see Table 1, Table 2, Figure 7), suggesting that CryoFM is a versatile and expressive prior.
- For ab-initio modeling, the authors tested CryoFM on two real-world datasets.
- All figures are of high quality.
summary: This work is in the field of cryogenic electron microscopy (cryo-EM), a microscopy technique that can be used to study biological structures and macromolecules (e.g. proteins). Cryo-EM produces 2D projections of a structure, which are then used to estimate a 3D map, parametrized as a voxel array, representing the structure.

This paper addresses the problem of how to incorporate knowledge from the large database of existing high-quality 3D cryo-EM density maps into the process of estimating and refining new density maps from cryo-EM data. In the language of inverse problems, this problem translates to the question of how to obtain a prior for cryo-EM density maps.

The authors propose CryoFM, a generative deep learning model that captures the prior distribution of high-resolution cryo-EM density maps. After pre-training the model on a large dataset of cryo-EM density maps, the authors show that it can be used as a prior to regularize inverse problems where a 3D density map must be estimated from corrupted measurements of the underlying object. As examples of such inverse problems, the authors consider 3D denoising tasks and the problem of estimating the 3D density of an object from a few 2D projections of the density. To learn the prior from the data, the authors use techniques from flow matching (similar to denoising diffusion models).

**Recommendation**: In my opinion, the paper is above the acceptance threshold: Judging from the experiments, CryoFM is a good and versatile prior for cryo-EM, and the authors have done good work in the promising direction of learning a cryo-EM prior from data (see "Strengths")! *However*, parts of the paper should definitely be revised before acceptance. I find parts of the paper really hard to understand, and context and details are missing in some places (see "Weaknesses").  
I will raise my score to "Accept" if my concerns related to the paper's clarity (see "Weaknesses") are sufficiently addressed.

**Post-Rebuttal:** I appreciate the authors' additional clarifications and increase my score from 6 to 8. Overall, the paper is an interesting case study of how to learn and use a prior for reconstruction and denoising in cryo-EM.
I also appreciate that the authors have made their training dataset publicly available (see "General Response"). I could not find the figshare link in the paper. I strongly encourage the authors to add the link to the paper, because I consider the dataset to be a valuable resource for future work in the direction of learning priors.
weaknesses: In many places, the paper is hard to understand and details and context are lacking. For example:

- I find the theory of the flow-matching approach for training the prior very difficult to understand. This is especially true for Section 3 and Section 4.3, which together make up about 1.5 pages. In particular, important concepts and equations are not explained, making it difficult to understand their meaning. For example:
	- In lines 159 to 162, the authors state that the learnable vector field generates a probability density path, but I could not find an explanation of what such a path is or how the vector field relates to a probability density function.
	- I could not find a definition or explanation of the variable (or function?) $\hat{\mathbf{x}}_0$ (lines 186-187), which appears many times in the paper and is an important part of the flow posterior sampling (algorithm 1).

 - I cannot judge the novelty of the variant of flow matching used in the paper. If the variant contains novel elements, the authors should explain them and highlight the differences to existing work more clearly. 

- Details and context are missing in some places in the experiments. For example:
	- The authors use EMReady, DeepEMhancer, and spIsoNet as baselines for the experiments in Section 5.1 and Section 5.2. They briefly introduce these baselines in lines 116 - 120, but I find this explanation insufficient to understand what the baselines actually do, why they were chosen, and how they relate to CryoFM.
	- In the denoising experiments in Sections 5.1 and 5.2, it is not clear what the test set is.
	- The motivation for studying anisotropic denoising (Section 5.2) is not entirely clear to me. Why does the noise affect "some orientations more than others"? Is this related to the preferred orientation problem, i.e. that some particles adopt certain orientations more often than others?

Other weaknesses:
- The authors propose a novel prior for solving inverse problems related to cryo-EM, but there is no discussion or experiment comparing CryoFM with existing priors, e.g. the Gaussian prior by Scheres et al. mentioned by the authors in line 38. One possible idea is to compare, for example, the Gaussian prior to CryoFM for regularizing ab-initio modeling.

- In Section 5.4, the authors use CryoFM for ab-initio modeling and compare it to cryoSPARC. It seems to me that cryoSPARC uses 2D projections of the particle directly, while CryoFM uses class averages. This difference in input data could be problematic when comparing the methods. This raises two questions:
	- Why did the authors use CryoSPARC as baseline? 
	- Is there another ab-initio reconstruction method that works with class averages?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. For the ab initio modeling approach, comparisons with established methods such as CryoDRGN2 and CryoFIRE would strengthen the evaluation of CRYOFM’s performance.
2. Regarding Figure 7, it would be interesting to see if using the same dataset with CryoSPARC actually improves performance.
3. Curious about effectiveness of the 2D to 3D reconstruction approach, particularly when only a limited set of 2D views is available
soundness: 3
strengths: 1. The development of a foundation model for cryo-EM is both significant and novel
2. The use of flow matching to capture the prior distribution, along with the flow posterior sampling method as a flexible prior for various downstream tasks, is technically sound
3. Overall, the paper is well organized
summary: This study introduces CRYOFM, a foundation model using flow matching to learn and generalize the distribution of high-quality density maps. CRYOFM serves as a flexible prior for multiple tasks in cryo-EM and cryo-electron tomography (cryo-ET), achieving state-of-the-art performance across applications without fine-tuning, underscoring its potential for broad use in these fields.
weaknesses: 1. While CRYOFM serves as a foundation model, it currently lacks demonstrated general applicability and versatility. As noted by the authors, further experiments across diverse tasks, such as 2D-to-3D reconstruction, would be beneficial to substantiate its broader utility.
2. CRYOFM relies on synthetic noise for denoising tasks. Testing on real datasets (realistic noise) would be necessary to validate its effectiveness in real-world applications.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Most important questions:
- For the ab initio reconstruction experiments: (1) why not use many noisy projections instead of a few averaged/denoised projections? This would be more realistic to modern practice and more similar to CryoSPARC. (2) why not use CryoFM for refinement as well as initialization? Since you have a high-resolution foundation model I expect that you might see improvement by using it at the refinement stage.
- Do the authors plan to release their code and pretrained foundation model to the research community? My preference for paper acceptance assumes that these will be released.

Minor questions:
- Does the FM in the title stand for “foundation model”, “flow matching”, or both?
- Line 209-210 says that proteins with helical structures are removed from the pretraining dataset. Then line 348 says that the proposed method enhances “the resolution of helices and loops”…why were helical structures removed from the pretraining dataset if they will be important in the test dataset?
- What is the motivation behind the specific data augmentations that were used for CryoFM-S and CryoFM-L?
- Figure 9 does not show results for ab initio reconstruction; why?
soundness: 3
strengths: There are many downstream protein-related modeling and reconstruction tasks that stand to benefit from an accurate model of the distribution of 3D protein density maps. The paper demonstrates applicability to several such tasks, and there are many more for which the proposed foundation model may be useful.
summary: The paper proposes a flow-based diffusion foundation model for 3D electron density maps of proteins as measured via cryogenic electron microscopy and tomography. This foundation model is then used as a prior for several 3D reconstruction and denoising tasks all framed as inverse problems: 3D density map denoising with spectral noise, 3D density map denoising with anisotropic noise, missing wedge restoration (for cryo-electron tomography), and ab initio reconstruction for cryo-electron microscopy from 2D projections.
weaknesses: - The motivation for using a flow-based model is not very clear
- Related work discussion could be improved. I would suggest mentioning Chroma, a foundation model for protein structures that was published in 2023: https://www.nature.com/articles/s41586-023-06728-8. My understanding is that the main difference for the proposed model is that it operates in the space of 3D electron density maps rather than discrete protein models, yet for many cryo-EM tasks either or both models might be useful (e.g. https://arxiv.org/abs/2406.04239 this recent paper uses Chroma as a diffusion prior for several protein inverse problems). The authors might also consider discussing DiffModeler (https://www.nature.com/articles/s41592-024-02479-0), which also proposes a protein diffusion model that can denoise 3D density maps (preprint was posted in January 2024 but the published version was just released a week ago).
- Also pertaining to related work, the paper would benefit from a richer description of the specific baselines that are used for comparison in the experiments. These baselines are introduced in lines 116-120 but there is not enough detail to really judge what is similar and different about the mechanism of each model in the comparison.
- The experiments on anisotropic noise are not clearly motivated. It would seem to be somewhere in between the realistic noise models for cryoEM and cryoET, but both of these are separately considered so I don’t see much added value to consider anisotropic noise unless there is a measurement setting where it arises (that is not already covered by other experiments).
- There are specialized versions of the algorithm used for several of the applications, yet these modified algorithm descriptions are in the appendix rather than the main text. I realize space is limited but it would be preferable to have at least a description of the algorithm modifications in the main text.
- I am very happy to see experiments on ab initio reconstruction, but unfortunately the results in this setting are not all that compelling since the final quality is similar or slightly worse than CryoSPARC. Please refer to some questions about this experiment in the questions section. 

Minor points
- The forward operator A is described as a “degradation operator”. In some cases this is accurate, such as when the measurements are noisy, but in other cases (such as ab initial reconstruction) my understanding is that A denotes projection rather than some form of degradation. Calling A a forward model/operator or measurement model/operator would be more consistent with the inverse problems literature.
- There is a typo on line 231; Transformer is misspelled.
- Figure 3: the colors for CryoFM and DeepEMhancer are too similar in the FSC curve. Another strategy would be to change the opacity or style of the lines so that both lines are visible even when they overlap.

Meta-review: The paper presents CryoFM, a generative model for cryoEM density maps. The model can be used to estimate a 3D density map from corrupted measurements. The paper considers 3D denoising tasks. 

All four reviewers recommend acceptance, and note that the paper provides an effective approach to learn and use a prior for cryo-EM imaging. 

The reviewers pointed to some weaknesses, in particular in the exposition, lack of context, lack of testing on real data, and a need to evaluate more diverse tasks. Most of those concerns were addressed in the rebuttal. For example, the authors improved clarity, and explain that they do not consider real-world data for a lack of ground-truth data. There are approaches to evaluate without ground-truth data, but the reconstruction problems considered in the paper are sufficiently interesting by itself. The reviewers see no major concerns after the rebuttal. 

Some of the reviewers and I find the expositions of the model as a foundation model a bit misleading. The model is a generative model and the paper considers reconstruction tasks, thus it would be accurate to present the model and method as a reconstruction approach based on generative modeling, as opposed to a foundation model. 

I recommend to accept the paper for its contribution to cyro-EM reconstruction.

Decision: Accept (Poster)