CryoFM: A Flow-based Foundation Model for Cryo-EM Densities
Authors: Yi Zhou
Yilai Li
Jing Yuan
Quanquan Gu
Venue: ICLR 2025 Poster (2025)
Keywords: cryoEM; foundation model; flow matching
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

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 around 40k protein density maps at various resolutions. 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.

Published as a conference paper at ICLR 2025
CRYOFM: A FLOW-BASED FOUNDATION MODEL FOR CRYO-EM DENSITIES
Yi Zhou", Yilai Jing Yuan", Quanquan Gu²
ByteDance Research {zhouyi naive, yilai 1i, yuanjing.cc... eugene, quanquan. hatpoguuuannggbyttoo gu)ebytedance.com
ABSTRACT
Cryo-electron microscopy (cryo-EM) is a powerful technique in structural biol- ogy 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 around 40k protein density maps at various resolutions on EMDB However, cryo-EM data processing algorithms have yet to fully benefit from knowledge biomolecular density maps, with only a few recent models being data-driven but limited to specific tasks. In this study, present CRYOFM, a foundation model designed as a generative model, learning the distribution high-quality density maps and generalizing effectively downstream tasks. Built on flow matching, CRYOFM is trained to accurately capture the prior distribution of biomolecular density maps. Furthermore, we introduce flow posterior sam- pling 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 mate-of-theearr performance on most tasks and demonstrat- ing its potential as a foundational model for broader applications in these fields.
INTRODUCTION
Cryo-electron mieroscopy is an important technique in structural biology and drug discovery, al- lowing 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, 2015). A major component of cryo-EM data processing in- volves reconstructing 3D structures from noisy 2D projections of particles. can be formulated as an inverse problem (Bendory et al.) 2020) Singer & Sigworth 2020), which aims at recovering the Under this formula- tion, following Bayesian statistics, the objective is to sample a density from the posterior p(xly), ) which can be factorized to p(xly) x p(y| (xx)). Thus, the prior distribution becomes essential for guiding the reconstruction process and improving the accuracy of the resulting structures.
signal x E IR" (i.e., a clean protein density) from the observation y E m [
Scheres (012bb first introduced a Gaussian distribution as the prior p(x) in 3D reconstruction for cryo-EM, effectively functioning as a in many cases. Building on this, recent works have explored more sophisticated regularizers rather
low-pass filter, which has proven useful
than explicitly defining prior distribution (Punjani et al.] 2020 Tegunov et al.) 22221; Kimanius jtt al. 2021] Li al. 2022] (Schwab et al.) |024] Kimanius et al. 2024 Liu et al. 2024). 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., 22017] Ramírez-Aportela et al. 2020 Terwilliger et al. 2020a Kaur et al. 2021 Sanchez-Garcia al. 2021|He et al. (202)). Both share the common objective of improving the quality of cryo-EM maps, whether during refinement or in post- poos-processing Among all these approaches, a line of methods has emerged that utilizes data-dri driven models to directly learn p(xly) from data, showing strong potential (Sanchez-Garcis et al. 2021 He al., 2003 Kimanius et al. 2024). While powerful, these models are often designed for specific tasks, limiting their general applicability and versatility.
Equal contribution, Corresponding author.
FMMB (wwPDB Consortium 2023) is a public repository for cryo-EM volumes.
"The signal DD denotes clean protein density, where D is the side length. The ohservation have any arbitary shape, we use IR" for case of notation but without loss of generality.

---

Published as a conference paper at ICLR 2025
flows governed by the vector field Ve(+, inference: conditional flow x:
observation y
Po(Ka v)
prior
Vx. (Xg)
(3) training: unconditional flow X,
Reelihood tarm V., lag x.)
(()))) noise dist. (5- ()
data dist.
Figure 1: The overview of CRYOFM. the training stage, CRYOFM learns a vector field Ve (t, x,), whose corresponding probability flow generates the data distribution po (xx) of high-quality protein densities. In the inference stage, given an observation y. a likelihood term (y to convert the unconditional vector field (t, sample from the posterior distribution Pu (x) This enables signal restoration of the density map, resulting in improved resolution of the alpha helices in the shown case.
(y|()))) (xx) is incorporated conditional one (t, x₂ xyly. so that we can
Foundation models leverage prior knowledge and provide versatility across a wide range of appli- cations. In particular, they have revolutionized Al for protein structure prediction and design by utilizing vast datasets to tackle diverse tasks (Jumper et al., 2222 Back et al., 2021 Lin et al. 2023
Krishna et al. grabam et al.
2024 Hayes et al. 2024; Abramson et al. 2024| Wang et al.) 2024b). Chroma (In-
2023) and RFDiffusion (Watson et al. (203) are two foundation models for protein many structure-related tasks. However, exper- imental 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 mod- els. which integrate prior knowledge and offer flexibility across multiple applications. [Wang et al. (2024a) developed a condition diffusion model in the density space, which achieves superior results by refining the map for model building, but is limited to single downstream task. In this study, present a foundation model, CRYOFM, that directly models the prior distribution p(x) by learn- ing the distribution of high-quality cryo-EM density maps through flow matching. Notably, flow matching (Lipman et al.] 2022; Liu et al.) 2022a) is a well-established generative model, which achieves state-of-the-arr results many applications (Dao et al. 2023 Esser al. 2024 Yim al., 2023a). The key ingredient of flow matching is time-dependent vector field ve (f, xt). which can be employed to progressively denoise noise vector into data sample and closely related the score of the prior distribution Vx, log p(x:). In CRYOFM, when applied to downstream tasks, we draw sample from the posterior distribution with its score (Ho et al. 2020 [ong et al.) (002bb): Vx log p(x(y) = V log p(x) + V log p(y|x). In contrast to expert designed prior, data- 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 downstream task, we can combine it with the learned prior in a plug-and-pl fashion Zhang et al., 2021 Venkatakrishnan et al. 2013) sample from the posterior.
structure modeling, which are broadly applicable
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 medium to low resolution Fig. provides an overview of CRYOFM. In the training stage, we regress the vector field with high-quality density data downloaded from the EMDB database (wwPDB Consortium 2023). When applied to downstream tasks, we propose flow posterior sampling method to convert the vector field (l, x.) conditional one Ve x,ly), which can be used to sample from the posterior distribution p(x(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 summarized as follows: (1) we present the first flow-based generative model as foundation model that learns the distribution of high-quality cryo-EM density maps; (2) derive a flow posterior sampling algorithm, enabling CRYOFM to be effectively utilized prior for various downstream tasks; (3) in an unsupervised manner, we achieve better performance
in several downstream tasks compared DeepEMhancer (Sanchez- Garcia et al. 2021), 1), EMReady
Unless stated otherwise, CRYOPM refers CRYOFM-S throughout this paper.
2

---

Published as a conference paper at ICLR 2025
He et al. 2023) and spIsoNet (Liu et al. 2024 (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 match- ing for inverse problems. A more detailed discussion on these topics can be found in Section B
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 ob- served cryo-EM density maps (Terwiligger et al. 2020b). 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. 2021) and EMReady (He et al. 22023), learn the posterior p(xly) from data to recover high- frequency details. In contrast, self-supervised approaches, such as M (Tegunov et al. 2021) and splsoNet (Liu et al., 2024), 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 imaging technique used to reconstruct 3D volumes of biological specimens from 2D images captured at various tilt angles (Lučić et al. 2005). The resulting 3D volumes, known tomograms, provide detailed views of cellular structures, with smaller regions, called subtomograms, extracted to focus on specific structures like proteins (Wan & Briggs 2016). 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 Fourier space without information, leading anisotropic resolution and artifacts. Traditional approaches use signal processing and regularization techniques to address the missing wedge (Goris et al. 2012 Deng et al., 2016, Yan al.) 2019 Zhai al. (020), while recent deep learning methods have shown promise in tackling this issue by leveraging data-driven models (Liu
ct al. 2022b
Van Veen et al. 2024).
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. 1970). 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-noiss ratio (SNR) (Ludtke et al. 1999. Voss et al. 2010). However, since the introduction of cryoSPARC, modern methods now use raw particle images directly for model estimation, leveraging stochastic gradient descent (SGD) bypass the need for class averages (Punjani et al. 2017).
Diffusion/Flow Matching for Inverse Problem Inverse problems refer to class of problems where the goal is to recover the original sample from the forward model's observation, such as image super-resolution (Haris et al. 2018), inpainting (Ych et al. 2017]. deblurring (Kupyn et al. 2019). etc. Since a pretrained DDPM/Flow models the data distribution (prior distribution), it help discover the posterior distribution given a forward model (likelihood). Specifically, Chung et al. (022) developed diffusion posterior sampling for general inverse problems, removing the need for strong assumptions like the linearity of the forward operator. Another line of research examines the solution to the inverse problem when the forward operator is unknown (Chung al., [2023a] Kapon et al. 2024). For protein structure modeling, Levy et al. (2024) adopts Chroma (Ingraham et al.] 2023 as the diffusion prior for a dozen of downstream tasks, including atomic model refinement in cryo-EM.
3 PRELIMINARIES
In this section, we briefly review two key components of cryoFM. In the training stage, flow match- ing (Section 3.1) is employed as the framework to learn the distribution p(density). In the inference stage, given some observations (noisy densities, 2D projections, etc.), we use the posterior sampling (Section 3.2 technique to sample a density from obbervatio.)...
3.1 FLOW MATCHING
Flow matching (Lipman et al.) 2022 Liu et al. 2022a) defines the generative process of a data point Xo E IR. by continuously transforming a sample x I from the noise distribution, where the
3

---

Published as a conference paper at ICLR 2025
state an intermediate timestep & [0, 1] is denoted by x, E R". A flow is the path traversed by the state xg. which is characterized by an ordinary differential equation (ODE):
dx, = ve x, ddl,
(1)
where ve : [0.1] x IR + R a time-dependent vector field parameterized by a neural network parameterized by θ, which defines the transformation of the state X,. At each time step t, the state X; follows a distribution pr (x(),,, where P: [O, 1] x IR + I220 is a time-dependent probability density function, i.e., f pe px ) = 1. Since pr px(x) evolves from the noise distribution pi (xI) to the data distribution Po (xx), it is referred to as a probability density path generated by the vector field (t, An illustration of the concepts of flow matching can be found in Fig. 12
Given a set of data points sampled from P1 (X1) and po (Xo), flow matching aims to learn the vector field Ve from the data. Chen et al. (2018) directly solves Eq. ① via differentiable neural ODE solver. Lipman et al. (2022) and Liu et al. (2022a) suggest that it is more efficient to learn a manually designed conditional vector field ( Xn) 0, R I given data point Xg- For example,
the conditional flow can be a straight line between Xy and a noise sample X1
xx|x, (1 t)xo + tx1;
(2)
where the corresponding vector field is u, (x(x,|(x) xx) = x1 Xo. Finally, the conditional flow matching objective (Lipman et al., (2022) is defined as follows:
( (xo0/_)))xx=-((0))))))))))))))))
(t, x4) (x1 - -x))
which is equivalent to optimizing the original flow matching objective (Lipman et al. 2022
3.2 DIFFUSION POSTERIOR SAMPLING
Diffusion posterior sampling (DPS) has emerged as promising approach for solving inverse prob- lems (Song et al. 202Tb at Chung et al. 2023b). Given a measurement y E Rm derived from R RT the goal is sample x from the posterior p(xly). By Bayes' theorem, it is easy to show that the score of the posterior distribution is:
x IR" with a forward operator
Xx, log (xx,l) Xx, (x,,)( + V x, log (y(xx,))
(3)
However, computing log (eg log Sx0, p(xo pp (y) | )xxo intractable since it requires the integration over all possible Xo ~ po [xo xx)). Chung et al. (2022) proposed to use Laplace approximation of the likelihood term: (y((x))) 2 po (y (x,)) )). where xa (x(x) ) is the posterior mean estimated by the diffusion model. The conditional score can thus be approximated by:
Vx, (x Xx, log):(x.) + Ux, ( (xx()))
If we assume the observation distribution po(y/x) Gaussian conditioned on x, we can show that the gradient of the log- likelikel satisfies:
Xx, Po (x,))) & Vx. lly (x)
(4)
Putting all these things together, we can approximate the score of the posterior distribution by:
V logg(x//y)) R Vx. x, lx, Vx, Ily Ax,
where A is associated with the partition function of Gaussian, and we treat as hyperparameter to control the step size 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 posterior sampling algorithm for the downstream tasks.
4.1 PRETRAINING DATASET
Our training dataset consists of deposited sharpened density maps from the EMDB (WwPDB Con- sortium 2023), specifically those with: 1) a reported resolution better than 3.0 A, 2) structures resolved by single-particle cryo-EM, and 3) data entries that include half-maps, ensuring that reso- lution estimates based on the "gold standard" Fourier shell correlation (FSC) (Henderson al.
4

---

Published as a conference paper at ICLR 2025
We manually curated this subset by removing exceptionally large complexes, helical struc- tures, and problematic cases through visual inspection (see Fig. for examples). We also excluded density maps with side lengths greater than 576 Á. This curation resulted in a total of 3479 density maps, where 32 density maps were selected as test set and excluded from training. The density maps in the selected subset were lowpass filtered 1.5 Á/voxel and 3 Á/voxel for CRYOFM-S and CRYOFM-L, respectively. For CRYOFM-S, we applied random cropping to volumes of size 64 along with random rotations for augmentation, whereas data for CRYOFM-L were center cropped to volumes of size 128 and augmented solely with random rotations. Additionally, we rescale the density value avoid large numerical variance (Section C.2 describes the details).
2012).
4.2
ARCHITECTURE
X; D2
Pg xx) E ER
Data or feature
DownSampling
UpSampling
Transformer w/ NA
Transformer w/ GA
Figure 2: CRYOFM's architecture. The side length D of the input x; undergoes dimension reduc- tion through down-sampling layers and is then expanded back to its original size. To minimize computational cost 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.
major challenge in applying the vanilla Transformer architecture to 3D density x, E RRD2 the computational complexity, which is ( + CRYOFM processes the 3D input data through Transformer with Aierarchical architecture, based on HDiT (Crowson et al. (2024). As illustrated in Fig. 2 the hicrarchical structure downsamples the spatial dimension initial levels and upsam- ples it at final levels, implemented via PixelUnshuffe and PixelShuffle layers (Shi et al.) (016). 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 em- ploys Neighborhood Attention (NA) (Hassani et al. 2023) since a localized attention layer is able to capture local dependencies while considerably reducing the number of tokens attend.
4.3 FLOW POSTERIOR SAMPLING
Given a vector field ve (f, x²) that generates the prior data distribution (x(),, we aim to convert it vector field (t, Xx ly) that generates the posterior (x)yyly). First, according to Dao et al.
(0033; Song et al. (2021b), we can connect the vector field with the score function Vx, log pr
follows:
d¹ vx, log (x))),
(e x = f(())
2
where ft the drift term and is the diffusion coefficient in the forward-time SDE. addition, we define the conditional vector field as follows:
o vx, (xly) = (t. x,) 2
9 xx, (y)
(t, xx)y = fx(x))
2
where the second equality holds due to Eq. B3. Thus, substracting the score the likelihood term weighted by the diffusion coefficient generates vector field that models the posterior distribution.
Following the rectified flow (Dao et al. 2023), we set fx,,) = 1- and 8i + in the condi- tional flow in Eq. (2) and obtain the conditional vector field as follows:
I-t
ve x;ly) (t, I
Vx. (v)x))))x))
(5)
Similar to Section 3.2 by substituting Eq. into Eq. 5 we obtain:
t
ve(t, x;ly) = ve + 1 1- Vx. lly - A
(6)
*The computational complexity of Transformers. is OTT.2), where D is the number of tokens.
5

---

Published as a conference paper at ICLR 2025
Algorithm Flow Posterior Sampling
Require:
a pretrained vector field ve 10, 1] R" + R", forward operator A R - R" an
obseravation & IRm, number of steps N, maximum step size Xmax at each step the recovered signal Xo
Ensure:
AL <
X1 (0,
for t E [1, - At, 1 111--A4- 2At, ... At] do
=t t- X At (t,),
1(xt) + -lyy- 11A--xo (xx((x)) (x:
f (xx)]l
[x,]
Normalize the gradient for numerical stability
X < min Amnux I-6
Prevent the weighting term from being too large
X f X, At X/g
end for return Xo
where (x,) = x; vte x²) is the posterior estimated by the flow model. In practice, we incorporate some tricks to avoid numerical instability. The detailed algorithm of flow posterior sampling is described in Alg. [ The algorithm can be adapted for different tasks; refer the reader to Appendix G for the additional versions.
5 EXPERIMENTS
Fourier Shell Correlation (FSC) is a widely used metric that compares two density maps in Fourier space (Harauz & van Heell 1986). 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 metries: FSCAUC. FSCo.5, and Fail Rate (FR). Specifically, FSCAUC measures the overall correlation across all spatial frequencies, and FSCo.s represents the resolution of the reconstructed map the standard 0.5 cutoff (Rosenthal Henderson 2003). The Fail Rate (FR) identifies cases where the method fails to run or produces a result that significantly deviates from 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 E.
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-dependen with higher frequencies exhibiting greater noise variance. Given a density V E DDD in the Fourier domain, the forward model A : DxDDD - CDxDDD is:
A(V) = +e,
where e DxDDD noise volume, whose values on the spherical shell with the same radius are the same (v denotes the index of component in the frequency space):
c(v) ~ N 2 (V)).
In the experiment, we estimate 22 with the same underlying signal with uncorrelated noise. See Section F.I for details.
oose (v) from the half maps, which are two independent observations
We manually introduce spectral noise, resulting in resolutions ranging from 3.2 Á to 15.0 A. Deep- EMhancer and EMReady, both pretrained models that directly learn the posterior p(xly), ), are used as baselines. CRYOFM demonstrates superior robustness, with no failed cases across all experi- ments, while both baselines show higher fail rates, especially at medium resolutions. Additionally, as shown in Tab. Tan Tab. Section CRYOFM achieves higher FSCAUC and better FSCo.5 in
'The step size X in Alg. 1 is not equivalent to X in Eq. Co). This is because solving ODE requires multiplying the velocity by the length of time interval, i.c., X time-sopps

---

Published as a conference paper at ICLR 2025
Table 1: Results of the spectral noise denoising task, comparing CRYOFM with DeepEMhancer Sannhez-Garche et al. (2221) and EMReady (He et al.) 2023). In Sections 5.1 5.2 5.3 the forward
model acts degradation of the original data. The best attainable FSC 3.0 since the voxel size is 1.5 A.
Estimated resolution at different level of noise
3.2 A
4.3 A
6.1 À
FR.
FSCauc f
FSCa.s
FRI
FSCauc 1
FSCo.5 4
FR.
FSCAuc 1
FSCp.5 J
After degradation
0.8876
3.21
0.5527
4.96
0.4969
6.07
DeepEMMancer
0.06
0.92
3.0
0.34
0.66 +: 0.05
4.3
0.44
0.62 0.05
4.5 0.3
EMReady
0.03
0.73 + 0.02
3.1
0.16
0.59 + 0.08
=.9
0.19
0.35 + 0.09
5.3 = 0.9
CRYOFM
0.00
0.95 + 0.01
3.0 + 0.0
0.00
0.68 + 0.04
4.4 0.2
0.00
0.62 + 0.04
4.6 + 0.3
Degraded
CryoFM
FSC (vs GT)
Ground truth
0.75
EMD 18549
DeepEMhancer
EMReady
0.25
100
20
10
8
Resolution (A)
Figure 3: Result of the spectral noise denoising task. Two density maps from EMDB were added spectral noise that the estimated resolution is 4.3 Å. The degraded density maps (results after applying the forward model) were filtered by relion postooopssssoeesss for visual clarity.
cases, significantly improving degraded density maps and enhancing the resolution of helices and loops, as illustrated in Fig. 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
Anisotropic noise in cryo-EM occurs when noise distribution varies by direction, affecting some ori- entations more than others and leading to uneven reconstruction quality. As result, certain orien- tations have lower signal-to-noise ratios than oth- To approximate this degradation in simpli- fied manner, we amplify the spectral noise by a fac- tor when the particle orientation falls within spe- cific range of angles, simulating the increased un- certainty in less-sampled orientations. The forward model A DDD + DXDDD is given by:
(ii)
Anisotropic noise
Missing wedge
Figure 4: Forward operators for (i) the anisotropic noise and (ii) the missing wedge.
V +e, Va otherwise
if O S 0max
A(V)
where E DDDDDD is the same defined in Eq. 77 and 1 scalar that increases the noise for specific orientation angles, reflecting heightened uncertainty those directions. 0min and 8max the angles that controls the portion of the noise being amplified as illustrated in Fig. 4 In the experiment, we estimate o (v) and a from the half maps. See Section F2 for details.
We manually add varied anisotropic spectral noise to the test set density maps and use Deep- EMhancer, spIsoNet, and CRYOFM for signal restoration. As shown in Tab. 2 CRYOFM out- performs both baselines FSC metrics and has zero fail rate. Fig, 5 de that CRYOFM restores the corrupted signal without introducing significant bias, unlike DeepEMManncer, which in- troduces artifacts and alters the overall shape. While splsoNet preserves signals across frequencies, fails to improve correlation with the ground truth. This suggests that pretrained models that learn the posterior p(xly) like MeepEeaancer struggle to retain reliable signals, and self-supervised

---

Published as a conference paper at ICLR 2025
Table 2: Results of the anisotropic noise denoising task, comparing CRYOFM with DeepEMhancer
Manchez-Garche et al. 2021) and spIsoNet (Liu et al. 2024).
0 [- 15" 13"
0 1-30", +30")
0E [-15,, 1-11111,000 +15"
FR.
FSCAuC
FSCo. 4
FR.
FSCxue 1
FSCos 4
FR.
FSCAIC t
FSCo.5 4
Alter degradation
0.6623
4.15
0.6324
4.27
0.6111
4.38
DeepEMhancer
0.22
0.80 1: 0.03
3.2 4: 0.1
0.22
0.79 = 0.05
3.3 +: 0.2
0.22
0.77 #: 0.05
3.4
splsoNet
0.65
0.01
4.27 0.03
0.00
4.37 0.03
CRYOFM
0.00
0.88
0.03
3.1 + 0.1
0.00
0.84 0.03
3.2 + 0.1
0.00
0.81
+ 0.04
3.4
Degraded
CryoFM
FSC (vs GT)
1.0
Ground truth
0.75
EMD- 100011995 17958
0.5
DeepEMhancer
spisoNet
0.25
0
10
Resolution (A)
Figure 5: Result of the anisotropic noise denoising task. Two density maps from EMDB are added anisotropic spectral noise that the estimated resolution is 4.38 Å. The degraded density maps are filtered by relion postprocess for visual clarity.
methods like splsoNet lack a data-driven foundation, limiting their effectiveness. Additionally, splsoNet requires over hours 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 subtomo- grams. We simulate this effect by applying a wedge-shaped mask the Fourier domain, removing data from unmeasured orientations. The forward model A: CDDDxD DDDDDD is given by:
A(V) = [V. if min S SAS 0max 0, otherwise
where 8min and max typically set to -60 and +60 respectively, representing the common tilt angle limits experimental setups as illustrated in Fig. For posterior sampling, we slightly modify the flow posterior sampling algorithm to get Alg. The modification is based on the fact that it is not necessary to compute the gradient of the loss function, since we can directly combine the observed part y with the remaining part X2 to maximize the likelihood.
We apply the missing wedge effect to the test set, resulting in Tab. Bladd example in Fig. 6 CRYOFM effectively restores the original signal, reducing the missing wedge artifacts, as shown by both FSC and visual inspection. The FSC curve confirm strong correlation with the ground truth, particularly in the owwfrreuency range.
Table 3: Results of the missing wedge restoration task.
FSCAUC 1 (Missing Region)
FSCAUC t
After degradation 0.80 + 0.02
0.0000
CRYOFM
0.92 0.02
0.76 + 0.06
5.4
Ab initio MODELING
Ab initio modeling in cryo-EM involves reconstructing coarse density map from 2D particle pro- jections, which serves as a reference for refinement, Here, we simplify the problem by focusing on generating coarse density map from few clean 2D projections obtained via upstream 2D clas- sification. This approach mirrors earlier ab initio methods (Ludtke et al. 1999 Voss et al., 2010) before cryoSPARC's introduction of SGD (Punjani al. 2017). Given K projections, there exist.
8

---

Published as a conference paper at ICLR 2025
Degraded
CryoFM
FSC (vs GT)
1.0
Ground truth
EMD 18549
Resolution (A)
Figure 6: Result of the missing wedge restoration task. Two maps from EMDB are masked in Fourier space to simulate the missing wedge effect in cryo-ET. The central slices through the maps are also shown.
Synthotic data
Roal data
EMD-20795
EMPIAR-1002
EMPIAR-
105.247 parsiles
inpur class
137,700 rsput parecles
Ripurt cbis overages
vntoo (crySSPARC)
Ab walo (cryofFt)
Ab eido (eryoSARC)
Ab hrio (tryaRM)
Ground truch
input projections
Ab vnio (cryoFM)
(crySSPARC)
(cryOSPARC)
(crySPARC)
(crySPARR)
A
A
4,90 A
CryoFM
CryoSPARC
CryoFM
CryoSPARC
CryoFM
Figure 7: Result of the ab initio modeling task. One synthetic data and two real data are shown.
K forward operators, where the k-th forward model is A RDXDXD + DDD given by:
A (V) = P( V), k E 11,,, K],
where is a projection operator the real space, and ((()) E 8C(3) x TR2 is the pose of the 2D projections. In this task, (k)) unknown and can be pre-determined. We modify the flow posterior sampling to get Alg. 3 where the main modifications are: (i) searching the optimal pose iteratively the sampling process (Scheres) 2012a Punjani et al. 2017: Zhong et al. (021b). and (ii) computing the likelihood using correlation to avoid numerical issues. 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 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 those generated by cryoSPARC, demonstrating its potential for 3D reconstruction in cryo-EM.
5.5 ABLATION & DISCUSSION
In this section, 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 H
Moderate patchifying and downsampling parameters lead to efficient training with minimal performance degradation. As shown in Fig. g reducing both patchifying and downsampling pa- rameters lowers the test loss, since less downsampling causes less information loss. However, the

---

Published as a conference paper at ICLR 2025
Paech by
Downtmppe Dovrsimpe
Paeckity
27
Dowssungic
8.6
0.6
F5Cau
After Degridiian
Nermal
Norna laston
101t
05
Arsootropic Notia
Masin
Training Compute (Gitops)
Training Compute (Gllops)
No.s0
Wedgs
Figure 8: MSE loss on as function of total training compute. The estimation of training compute is consistent with Peebles Xie )2023)
Figure 9: Downstream
performance with respect to data normalization.
Amax
7.5
4.5
7.5
4.5
50
200
1000
LFM-4
LFM-8
CryoFM
timesteps
Figure 10: Averaged rank- ing (4) across downstream tasks for hyperparameters.
Unconditional sampling from models in different spaces. For latent flow models (LFM), the high-frequency information contains more noise.
Figure 11:
reduction becomes marginal while the training compute (Gflops) increases significantly. To achieve 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. Data normalization is the last step in data processing (as referenced in Section C.2) density with a value ranging from [0, is subtracted by 0.04 and then divided by 0.09. Fig. demonstrates that models trained with normalized data consistently improves the FSCAUC metrics compared to the one without normaliza- tion in three downstream tasks. The observation is in line with some DDPM-basch work (Yim al.) 2023b Rombach et al.) 2022), as a significant change the variance of the data can make the model difficult train (Karras et al. 2022).
More sampling steps boosts the performance in density restoration. Fig. 10 presents the aver- aged ranking across three downstream tasks for different time steps (#timesteps) and the maximum step sizes (Amax). The main factor influencing performance is the number of time steps. A smaller Amax may reduce the impact of the likelihood term and thus degrade performance. We choose the best combination: Amax and #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. 2022). Fig. illusstrats the unconditional sampling results from two latent flow models (LFM), where the nigh-frequency information more noisy. The observation concides with that of Crowson et al. 2024), where they found that latent diffusion models fail generate fine details. We suspect that the VAE is under-trained due to a small amount of data, and we will leave this 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, derive the forward 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 power- ful 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.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger. Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.086 08774, 2023.
Minkyung Back, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, R Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using three-track neural network. Science, 373(6557):871- 876. 2021.
Tamir Bendory, Alherto Bartesaghi, and Amit Singer. Single-particle cryo-electron microscopy: Mathematical theory, computational challenges, and opportunities. IEEE signal processing mag- azine, 37(2):58-7,, 2020.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models large datasets. arXiv preprint arXiv:2311. 2023.
Yochai Blau and Tomer Michaeli. The perception-distortin tradeoff. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6228-62277 2018.
Avishek Joey Bose, Tara Akhound-Sadegh, Kilian Fatras, Guillaume Huguet, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexan- der Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint 1xXiv:3310002391 2023.
Andrew Brock. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1096, 2018.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 146877 2022.
Hyungjin Chung, Jeongsol Kim, Sehui Kim, and Jong Chul Ye. Parallel diffusion models of operator and image for blind inverse problems. IEEE/CVF Conference Computer Vision and Pattern Recognition, 2023a.
Hyungjin Chung. Dohoon Ryu, Michael T McCann, Mare L Klasky, and Jong Chul Ye. Solving 3d inverse problems using pre-trained 2d diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 22542-22555 2023b.
Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Kaplan, and Enrico Shippole. Scalable high-resolution pixel-space image synthesis with hour- glass diffusion transformers. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learn ing Research, pp. 9550-9575. PMLR, 21-27 Jul 2024. URL https::/proceee: igg mll.. htm
press/ http://33533333333
Richard Anthony Crowther, DJ DeRosier, and Aaron Klug. The reconstruction of three-
dimensional structure from projections and its application to electron microscopy. Proceedings of the Royal Society of London. A. Mathematical and Physical Sciences, 317(1530)3319-444 1970.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran. Flow matching in latent space. arXiv preprint S0t::2307.00699 2023.
11

---

Published as a conference paper at ICLR 2025
Sacha De Carlo and J Robin Harris. Negative staining and cryo-negative staining of macromolecules and viruses for tem. Micron, 12(2):177-111, 2011.
Yuchen Deng, Yu Chen, Yan Zhang, Shengliu Wang. Fa Zhang, and Fei Sun. Icon: 3d reconstruction with mnisinnginformmttoon restoration in biological electron tomography. Journal of structural biology, 195( 100-112 2016.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1810v:881000480000 2018
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis, Advances in neural information processing systems, 0187880-8794, 2021.
John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. computational and applied mathematicx, 6(1):19-26, 1980.
Journal of
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Weijie Gan, Yuyang Hu, Jiaming Liu, Hongyu An, Ulugbek Kamilov, Block coordinate plug-and-play methods for blind inverse problems. Advances in Neural Information Processing Systems, 36, 2024.
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Bart Goris, Wouter Van den Brock, Kees Joost Batenburg, H Heidari Mezerji, and Sara Bals. Electron tomography based on total variation minimization reconstruction technique. Ultra- microscopy, 113: 120111000000. 2012.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367. 2018.
George Harauz and Marin van Heel. Exact filters for general geometry three dimensional recon- struction. Optik,
1986.
Muhammad Haris, Gregory Shakhnarovich, and Norimichi Ukita. Deep back-projection networks for super-resolution. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1664-1673, 2018.
Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi. Neighborhood attention trans- the former. In Proceedings the IEEF/CVF Conference on Computer Vision and Puttern Recognition
(CVPR), pp. 6185-6194, June 2023.
Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation, 2024. URL
wwtpabs////ooaoooooooo.......a /2312. 02139
Tomas Hayes, Roshan Rao, Halil Akin, Nicholas J Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q Tran, Jonathan Deaton, Marius Wiggert, al. Simulating 500 million years of evolution with a language model. bioRxiv, pp. 2024-07, 2024.
Jiahua He, Tao Li, and Sheng-Yor You Huang. Improvement of cryo-em maps by simultancous local and local deep learning. Nature Communications, 14(1):322 ::317, 2023.
Richard Henderson, Andrej Sali, Matthew I. Baker, Bridget Carragher, Batsal Devkota, Kenneth H Downing, Edward H Egelman, Zukang Feng, Joachim Frank, Nikolaus Grigorieff, et al. Outcome of the first electron microscopy validation task force meeting. Structure, 20(2):205-214, 2012.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851 2020.
12

---

Published as a conference paper at ICLR 2025
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pp. 13213 13213-13232. PMLR. 2023.
John B Ingraham, Max Baramov, Zak Costello, Karl W Barber, Wujie Wang, Ahmed Ismail, Vincent Frappier, Dana M Lord, Christopher Mig-Thow-Hing Erik R Van Vlack, et al. Illuminating protein space with a programmable generative model. Nature, 623(7989);1070-1078, 2023.
Arjen J Jakobi, Matthias Wilmanns, and Carsten Sachse. Model-based local density sharpening of cryo-em maps. Elife, 6:e27131, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterrization with gumbel-softmax. arXiv preprint 1rXiv:6611.00 144, 2016.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurmvv, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-588) 2021.
Roy Kapon, Guy Tevet, Daniel Cohen-Or, and Amit H Bermano. Mas: Multi-view ancestral sam- pling for 3d motion generation using 2d diffusion. Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1965-1974, 2024.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-26577, 2022.
Satinder Kaur, Josue Gomez-Blanco, Ahmad AZ Khalifa, Swathi Adinarayanan, Ruben Sanchez- Garcia, Daniel Wrapp, Jason McLellan, Khanh Huy Bui, and Javier Vargas. Local computa- tional methods to improve the interpretability and analysis of cryo-em maps. Nature communica- tions, 12(1):1240) 2021.
Bahjat Kawar, Gregory Vaksman, and Michael Elad. Snips: Solving noisy inverse problems stochas- tically. Advances in Neural Information Processing Systems, 34:21757-21766 2021.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593-23606, 2022.
Dari Kimanius, Gustav Zickert, Takanori Nakane, Jonas Adler, Sebastian Lunz, C-B Schönlieb, Ozan Öktem, and Sjors HW Scheres. Exploiting prior knowledge about biological macro- molecules in cryo-em structure determination. IUCrJ, 8(1):60-755 2021.
Dari Kimanius, Kiarash Jamali, Max E Wilkinson, Sofia Lövestam, Vaithish Velazhahan, Takanori Nakane, and Sjors HW Scheres. Data-driven regularization lowers the size barrier of cryo-em structure determination. Nature Methods, pp. 1-6, 2024.
Diederik P Kingma. Auto-encoding variational bayes. In International Conference Learning Representations, 2014.
Karsten Kreis, Tim Dockhor,, Zihao Li, and Ellen Zhong. Latent space diffusion models of cryo-em structures. arXiv preprint arXiv:2211.14169, 2022.
Rohith Krishna, Jue Wang, Woody Ahern, Pascal Sturmfels, Preetham Venkatesh, Indrek Kalvet, Gyu Rie Lee, Felix biomolecular modeling and design with rosettafold all-atom. Science, 384(6693) 2024.
Ivan Anishchenko, Ian R Humphreys, et al. Generalized
Orest Kupyn, Tetiana Martyniuk, Junru Wu, and Zhangyang Wang. Deblurgan-v2: Deblurring corders-of-magnin faster and better. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 8878-8887, 2019.
Charles Laroche, Andrés Almansa, and Eva Coupete. Fast diffusion em: diffusion model for blind inverse problems with application to deconvolution. In Proceedings of the IEEE/CVF Winter Conference Applications of Computer Vision, pp. 5271-5281, 2024.
13

---

Published as a conference paper at ICLR 2025
Andres E Leschziner and Eva Nogales. The orthogonal tilt reconstruction method: approach to generating single-class volumes with no missing cone for ab initio reconstruction of asymmetric particles. Journal of structural biology, 153(3):284 299, 2006.
Axel Levy, Eric R Chan, Sara Fridovich- Keil,, Frédéric Poitevin, Ellen D Zhong, and Gordon Wet- zstein. Solving inverse problems in protein space using diffusion-based priors. arXiv preprint arXiv:2406.04235 04239, 2024.
Yilai Li, Yi Zhou, Jing Yuan, Fei Ye, and Quanquan Gu. Cryostar: leveraging structural pri- ors and constraints for cryo-em heterogencous reconstruction. Nature Methods, Oct 2024. ISSN 7105,7105. doi: 10.1038//4119--02--22222 02486-11. URL https:://ccii http://00o..0r1111103133 541592- 020-922-24--228---
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 37966637) 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 2rX10:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 2022a.
Yun-Tao Liu, Heng Zhang, Hui Wang, Chang-La Lu Tao, Guo-Qiang Bi, and Z Hong Zhou. Isotropic reconstruction for electron tomography with deep learning. Nature communications, 13(1);6482, 2022b.
Yun-Tao Liu, Hongcheng Fan, Jason Hu, and Z Hong Zhou. Overcoming the preferred orientation problem in cryoem with self-supervised deep-learning. bioRxiv, 2024.
Vladan Lučié, Friedrich Förster, and Wolfgang Baumeister. Structural studies by electron tomogra- phy: from cells molecules. Annu. Rev. Biochem., 74(1):833 865, 2005.
Steven J Ludtke, Philip Baldwin, and Wah Chiu. Eman: semiautomated software for high- resolution single-particle reconstructions. Journal of structural biology, 128(1):82-97, 1999.
Nick Moran, Dan Schmidt, Yu Zhong, and Patrick Coady. Noisier2noise: Learning to denoise from unpaired noisy data. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12064-12072, 2020.
Piotr Nawrot, Szymon Tworkowski. Michal Tyrolski, Lukasz Kaiser, Yuhuai Wu, Christian Szegedy, and Henryk Michalewski. Hierarchical transformers are more efficient language models. In Ma- rine Carpuat, Catre-Catherine de Marneffe, and Ivan Vladimir Meza Ruiz (eds.), Findings of the Association for Compulational Linguistic:: NAACL 2022, pp. 1559-1571, Seattle, United States, July 2022. Association for Computational I Linguistics. doi: 10. f0.18653/v1/0222 . findi
Anddngs-naacl 17. naac 1aa.nggnnnna1111
URI https::/aacaa http:ppsee/cccaannoooooogoo/o/oe
antho wow.aooooo...o///aao0eo http://///oooaaooooooooooo//////
Eva Nogales and Sjors HW Scheres. Cryo-em: unique tool for the visualization of macromolecular complexity. Molecular cell, 58(4):677-689) 2015.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), pp. 4195-4205, October 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01992, 2023.
Ali Punjani, John L Rubinstein, David J Fleet, and Marcus A Brubaker. cryosparc: algorithms for rapid unsupervised cryo-em structure determination. Nature methods, 14(3):290-29,, 2017.
14

---

Published as a conference paper at ICLR 2025
Ali Punjani, Haowei Zhang, and David Fleet. Non-uniform refinement: adaptive regularization improves single-particle cryo-em reconstruction. Nature methods, 17(12):1214-15) 1221, 2020.
M Radermacher, Wagenknecht, A Verschoor, and J Frank. A new 3-d reconstruction scheme applied to the 50s ribosomal subunit of e. coli. Journal of microscopy, 141(1)RPP1RRP) 1986.
Erney Ramfre7-Aportela, Jose Luis Vilas, Alisa Glukhova, Roberto Melero, Pablo Conesa, Marta Martínez, David Maluenda, Javier Mota, Amaya Jiménez, Javier Vargas, et al. Automatic local resolution-based sharpening of cryo-em maps. Bioinformatics, 36(3):765-772, 2020.
Alexander Rives. Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma. et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 1 118(15)e2016229111 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of IEEF/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Peter B Rosenthal and Richard Henderson. Optimal determination of particle orientation, absolute hand, and contrast loss in single-paree particle electron cryomicroscopy. Journal of molecular biology, 333(4):721-745, 2003.
Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar. Theory and experiments vector quantized autoencoders. arXiv preprint arXiv: 1063, 2018.
Ruben Sanchez-Garcia Josue Gomez-Blanco, Ana Cuervo, Jose Maria Carazo, Carlos Oscar Sorzano, and Javier Vargas. Deepemhancer: deep learning solution for cryo-em volume post- processing. Communications biology, 4 (1):874, 2021.
Sjors H.W. Scheres. A Bayesian View on Cryo-EM Structure Determination. Journal of Molecular Biology, 415(():406-41) 418, 2012a. ISSN 00222836. doi: (0.1016/f.mbb22011111
Sjors H.W. Scheres. RELION: Implementation of Bayesian approach to cryo-EM structure de- termination. Journal of Structural Biology, 180(3):519-530, 2012b. ISSN 10478477. doi: 20.10166jjsh..221200..o
Johannes Schwab, Dari Kimanius, Alister Burt, Tom Dendooven, and Sjors HW Scheres. Dy- namight: estimating molecular motions with improved reconstruction from cryo-em images. Na- ture Methods, pp. 1-8, 2024.
Wenzhe Shi, Jose Caballero, Ferenc Huszár, Johannes Totz, Andrew Aitken, Rob Bishop, Daniel Rueckert, and Zehan Wang. Real-time single image and video super-resolution using effi- cient sub-pixel convolutional neural network. In 2016 JEEE Conference on Computer Vision and Pallern Recognition (CVPR), pp. 1874-1883, 2016. doi: 10. 10.110/CCPP..216.220
Amit Singer and Fred J Sigworth. Computational methods for single-partich electron cryomi- croscopy. Annual review of biomedical data science, 3(1):133-190 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International conference on machine learn- ing. pp. 2256-226622222 2265. PMLR, 2015.
Yang Song. Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:21/1.08005, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021b. URL
.et/
Eorum?id-Px111122
Dimitry Tegunov, Liang Xue, Christian Dienemann, Patrick Cramer, and Julia Mahamid. Multi- particle cryo-em refinement with m visualizes aaoosome-antiiionn complex at 3.5 in cells. Na- fure methods,
2021.
15

---

Published as a conference paper at ICLR 2025
Thomas C Terwilliger, Paul D Adams, Pavel V Afonine, and Oleg V Sobolev. Cryo-em map inter- pretation and protein model-building using iterative map segmentation. Protein science, 29(1): 87-99, 2020a.
Thomas C Terwilliger, Steven J Ludtke, Randy J Read, Paul D Adams, and Pavel V Afonine. Im- provement of cryo-em maps by density modification. Nature methods, 17(9):23-9277, 2020b.
Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks.
In International conference on machine learning, pp. 1747-1756. PMLR, 2016.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Dave Van Veen, Jesús G Galaz-Montoyy, Liyue Shen, Philip Baldwin, Akshay S Chaudhari, Dmitry Lyumkis, Michael F Schmid, Wah Chiu, and John Pauly. Missing wedge completion via unsuper- vised learning with coordinate networks. International Journal of Molecular Sciences, 25(10): 5473, 2024.
Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. Plug-and-play priors for model based reconstruction. In 2013 IEEE global conference on signal and information process- ing. pp.
IEEE, 2013.
Neil R Voss, Dmitry Lyumkis, Anchi Cheng, Pick- Wei Lau, Anke Mulder, Gabriel C Lander, Ed- ward Brignole, Denis Fellmann, Christopher Irving, Erica L Jacovetty, et al. A toolbox for ab initio 3-d reconstructions in single-particle electron microscopy. Journal of structural biology, 169(3):389-399) 2010.
W Wan and John AG Briggs. Cryo-electron tomography and subtomogram averaging. Methods in enzymology, 579:329-367, 2016.
Xiao Wang, Han Zhu, Genki Terashi, Manav Taluja, and Daisuke Kihara. Diffmodeler: large macro- molecular structure modeling for cryo-em maps using diffusion model. Nature Methods, pp- 1-11, 2024a.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Diffusion language models are versatile protein learners. arXiv preprint arXiv:2402. 18567, 2024b.
Joseph L Watson, David Juergens, Nathaniel Bennett, Brian L Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1089-11000 2023.
The wwPDB Consortium. EMDB the Electron Microscopy Data Bank. Nucleic Acids Research, 22DD1)D4566D4400 11 2023. ISSN 0305-1048. doi: h0.1093//ar/gkaa1019 URL https:
//doi.org/11 1093/nar/gkad1l
Rui Yan, Singanallur V Venkatakrishnan, Jun Liu, Charles A Bouman, and Wen Jiang. Mbir: A cryo-et 3d reconstruction method that effectively minimizes missing wedge artifacts and restores missing information. Journal of structural biology, 206(2):183-192, 2019.
Raymond A Yeh, Chen Chen, Teck Yian Lim, Alexander G Schwing, Mark Hasegawa-Johnson, and Minh N Do. Semantic image inpainting with deep generative models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5485-5493, 2017.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXiv:2310.05297 2023a.
Jason Yim, Brian I. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint X22.:2322.002777 2023b.
Xiaobo Zhai. Dongsheng Lei, Meng Zhang, Jianfang Liu, Hao Wu, Yadong Yu, Lei Zhang, and Gang Ren. Lottor: an algorithm for missing-wedge correction of the low-tilt tomographic 3d reconstruction of single-molecule structure. Scientific reports, 10419, 2020.
16

---

Published as a conference paper at ICLR 2025
Kai Zhang, Yawei Li, Wangmeng Zuo, Lei Zhang, Luc Van Gool, and Radu Timofte. Plug-and-ple image restoration with deep denoiser prior. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10):6300-631) 2021.
Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis. Cryodrgn: reconstruction of heterogeneous cryo-em structures using neural networks, Nature methods, 18(2):176-185, 2021a.
Ellen D. Zhong, Adam Lerer, Joseph H. Davis, and Bonnie Berger. CryoDRGN2: Ab initio neural reconstruction of 3D protein structures from real cryo-EM images. In 2021 IEEE/CVF Interna- lional Conference Computer Vision (ICCV), 4046-4055, Montreal, QC, Canada, 2021b. doi: 10. N0.100/9CCC44992222210000
IEEE. ISBN 978- 1666442-112-5.
Chunting Zhou, Lili Yu, Arun Babu, Kushal Tirumala, Michihiro Yasunaga, Leonid Shamis, Jacob Kahn, Xuezhe Ma, Luke Zettlemoyer, and Omer Levy. Transfusion: Predict the next token and diffuse images with one multi-modal model. arXiv preprint arXiv:2408. 1/09,, 2024.
A DISCUSSION ON FLOW MATCHING
A.1 CONCEPTS IN FLOW MATCHING
Fig. 12 illustrates three main concepts in flow matching. Here, the horizontal bar represents the time dimension, and each vertical slice represents iimedeepnneett probability distribution with only one dimension.
P(())
Po.2
PolKa)
R No.)
v(0.8.3..))
the flow of x,
80.
noise dist.
timestep
data dist.
field
The arrows denote the vector field v(J. x,), which represents how sample x, is transformed.
probobiliy density path
Each vertical slice denoces time-dependent probability distribution
P(K,). evolves from left right,
"path".
fow
trajectory of x, constructed continuous transformations defined by the vector field.
Figure 12: An illustration of three main concepts in flow matching: (i) the vector field (t, x:), (ii the probability density path Pe (xx), , and (iii) the flow of x,. The approximated posterior mean (x(x;) in Alg. 1 is also depicted on the right part.
A.2 CHOICE FLOW MATCHING FOR FOUNDATION MODELS
In this section, we discuss the choice of employing flow matching as the backbone of the cryo-EM foundation model. The main motivation is that (i) generative modeling trend in building modern foundation models, and (ii) flow matching excels at generative modeling.
Generative modeling is trend in building modern foundation models. ChatGPT (Achiam al., 2023) has disrupted the entire NLP (natural language processing) area with its extremely strong gen- erative power. It brings about paradigm shift from representation learning generative modeling. For representation learning, one of the most prominent models is BERT (Devlin) 2018), which per- forms well in language understanding but is not adept at performing generation. Similar trends
17

---

Published as a conference paper at ICLR 2025
can also be observed in computational biology. Two typical cases in this field are ESM-3 and Al- phaFold3. Built upon ESM (Rives et al. (021) (a BERT-style model), Hayes et al. (2024) developed ESM-3 that has the ability to generate protein sequences. A major change from A AlphaFold2 (Jumper et al. 2021) to AlphaFold3 (Abramson et al. 2024) is the diffusion module that can be utilized to generate ensemble of structures.
Flow matching excels at generative modeling. Diffusion-basel foundation models have had a significant impact generative modeling in Al for biology. AlphaFold3 (Abramson et al. 2024), RFDiffusion (Watson et al., 2023), and Chroma (Ingraham et al.) 2023) are all based on diffusion models (DDPM). ESM-3 (Hayes et al. 2024), although not explicitly stated in its manuscript, is closely related to discrete diffusion models. One of the crucial reasons for the success of DDPM re- sides in iterative refinement, which gradually refines the generated samples to more superior state. The iterative nature of DDPM enables it to outperform its counterparts with one-step generation methods, such as VAE and GAN. Flow matching can be regarded as a variant of DDPM and shares many similar properties with DDPM. We prefer flow matching over DDPM since it is easier to train and can converge faster. In the early stage of this study, we trained both DDPM and flow matching. We found that flow matching can yield similar results while incurring less computational cost.
B
ADDITIOIANT ONAL RELATED WORK
B. DENSITY MODIRCATION AND DENOISING IN CRYO-EM
Traditional density modification methods often apply Fourier space weighting, utilizing frequency- dependent scaling suppress noise and enhance signal (Jakobi et al. 2017 Ramirez-Aportela et al. 2020 Terwilliger et al., 2020a| Kaur et al., 2021I. These heuristic-
approaches, like Wiener- style deblurring (Ramírez-Aportela et al., 2020p, tend rely local resolution estimates and fil- tering strategies but may struggle with intricate structural details due to limited priors. On the deep learning side, Blush (Kimanius al.) 2024) further refine density maps by denoising half-maps dur- ing iterative refinement. However, they can introduce hallucinated details in lower-resolution maps. In contrast, methods such as M (Tegunov al. 2021) and spIsoNet (Liu al.) 2024) avoid exter- nal data, with M leveraging noise2noise (Moran et al. (2020) for denoising independent half-maps, and splsoNet 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. this paper, we selected DeepEMhancer Sanchez-Garcia et al. 2021), EMReady (He et al. 2023) and spisoNet as the baselines for the denoising tasks.
Mhancer DeepEMhancer is a deep learning- tool designed to enhance cryo-EM den- sity maps, improving their interpretability and aiding structural analysis. The method employs convolutional neural network (CNN) trained pairs of experimental cryo-EM maps and corre- sponding locally sharpened maps. This training enables DeepEMhancer to learn the mapping be- tween low-quality input maps and their high-quality counterparts. The pretrained model processes new cryo-EM maps to produce enhanced versions with improved clarity and detail, facilitating more accurate structural interpretations.
EMReady Similar to DeepEMhancer, EMReady is a computational method designed to enhance cryo-EM density maps by simultancously applying local and non-local denoising techniques. The method uses a Swin-Coni Conv-3DUNet and trained on pairs of experimental cryo-EM maps and corre- ponding synthetic maps from the atomic models. The pretrained model can enhance cryo-EM maps for better interpretations.
splsoNet spIsoNet is a self- supfrsupe deep learning method developed to address the preferred ori- entation problem in single-pane parttlle cyo.EM.. This issue arises when particles predominantly adopt specific orientations during imaging, leading anisotropic data and potential inaccuracies in 3D reconstructions. The method adopts a U-Net architecture and leverages the half maps in cryo-EM reconstruction to learn representations from the overrepresented orientations, recovering informa- tion for the underrepresented or missing views.
B.2 MISSING WEDGE PROBLEM IN CRYO-ET
Cryo-ET enables detailed visualization of macromolecular complexes and cellular structures their native environments. The missing wedge problem occurs due to the physical limitation of tilt angles
18

---

Published as a conference paper at ICLR 2025
during data acquisition, typically restricted to (-60", +60%), resulting in missing data in Fourier space and artifacts such as elongation along the missing axes. Traditional methods that use sig- nal processing techniques and regularization strategies (Goris et al. 2012 Deng et al. 2016 Yan et al. 2019 Zhai et al. 2020) are often based on heuristic assumptions and have limitations in fully recovering lost data. Recently, deep learning models have been applied this problem, offering improved recovery of complex patterns in tomograms (Liu et al. 2022b; Van Veen et al.) 2024) 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.
B.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.)) T 1986] Leschziner & Nogales, 2006), which provided indirect pose information, or negative stain (De Carlo & Harris 2011), which improved SNR but at the cost of nrihhffrequeng 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. 2010). While cryoSPARC's SGD ap- proach improved this by working directly with raw particles (Punjani et al., (017). certain structural features or heterogeneity observed in 2D class averages may still be lost in the final 3D reconstruc- tion, especially in challenging samples.
B.4 DIFFUIION/FLOW-BASED MODEL
Flow-based models (Lipman al. 2022 Liu et al. 2022a) and denoising diffusion probablistic
models (DDPM) -Sohl-Dickstein et al. 2013 Ho et al. (0200) are two modern deep generative mod- els. 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-aae results on many generative tasks, including image generation (Rombach al. [022]
Podell et al. 2023 Peebles & Xie 2023). video generation (Blattmann et al. (2023). and molecule
generation Yim et al. |023b; Abramson et al. 2024 Ingraham et al. 2023 Watson et al. 2023
Wang et al. 2024b), , etc. Recently, diffusion models have been adopted the cryo-EM field. Kreis
et al. (2022) traverses the latent space of cryoDRGN (Zhong et al. 2021a) with a diffusion model, while Wang et al. (2024a) 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 do- mains, including image generation (Esser et al. (2024) and moleculer generation (Yim et al. [002bb] Bose et al., (003).. The DDPM objective can also be unified into flow-based models by converting into a probability flow ODE (Song et al. 2021b).
B.5 VISION TRANSFORMERS FOR DIFFUSION MODELS
Diffusion transformers (Peebles & Xie 2023) have demonstrated significant scalability and gener- ative capabilities in image-relater lated tasks (Esser et al. 2024| Hoogebooo et al. 2023 Hatamizadehi ct al. 2024 Zhou et al. 2024). In particular, HDiT (Crowson et al.] 2024) leverages the inher- ent hierarchical nature of visual patterns in its model design. By integrating the characteristics of Diffusion Transformer (Peebles & Xie 2023) and Hourglass transformers (Nawrot et al.) 2022). and employing local attention mechanisms (Hassani et al. 2023), HDIT offers efficient model structure suitable for training the diffusion model in the data space.
EMD-14529
EMD-16694
EMD-10857
BMD-19483
EMD-19792
Figure 13: Examples of problematic cases removed from the training dataset after manual curation. Low contour levels are used to highlight the heterogeneous local resolution.
19

---

Published as a conference paper at ICLR 2025
64x64x66
64x64x664
Data or feature
DownSampling
+
x8
x4
UpSampling
16xl 6x
8x8xx8
16xl 6x1
Transformer
Input
Output
(a) CryoFM-S
Transformer GA
128x 128xx
128x 128x 28
x12
x2
32x32x32
6xx16x16
8x8x8
6xx16x16
32x32x32
Input
Output
(b) CryoRM-L
Figure 14: Overview of two 3D CRYOFM model details. (a) CRYOFM-S design for input shape 64 and (b) CRYOFM-L for shape 128
C ADDITIONAL DATA CURATION AND PROCESSING
C.I DATA CURATION
Examples of the problematic cases which excluded from the training after manual curation are shown in Fig. 13 These includes density maps with very heterogeneous local resolutions that significant portions of the maps are poorly resolved. The EMDB IDs of the training and testing data used in this paper have been uploaded to
C.2 DATA STANDARDIZATION
We preprocessed the deposited density maps normalize their values to a consistent range of [0, 1]. Since the absolute values of density maps are not intrinsically meaningful and the value ranges vary significantly, we applied a uniform processing procedure follows: Clipping: We clipped the values of each density map based on the minimum value, which is set to 0. and the 99.999th percentile value of the density map. Scaling: Subsequently, the clipped data were scaled to the [0, 1] range by adjusting according to the new minimum and maximum values obtained after clipping. Normalization: we used a mean of 0.04 and standard deviation of 0.09, which were calculated based the data that have been scaled to the [0, 1] range.
D ADDITIONAL DETAILED INFORMATION OF CRYOFM
D.1 IMPLEMENTATION DETAILS
Fig. 14 illustrate the model architectures for input dimensions of 64 and 1283. Fig. 15 shows the
details of different attention structure we used in the transformer block. Tab, details the specific model parameter configurations (aligned with the naming conventions used in Crowson et al. (0244) and the training hyperparameters.
employed the Fairseq (Ott et al. 2019) optimizer with default learning rate of le-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.
In all experiments,
D.2
LIKELIHOOD ESTIMATION
Given a data point Xg E R" we estimate its likelihood by solving a probability flow ODE (Song Lipman al. 2022 Chen al. 2018). We start with the continuity equation of
al. 202 lb
velocity helld Ve 00, x
d
(xe) + V ve((,x)) 0.
di
20

---

Published as a conference paper at ICLR 2025
Neighborhood Attention
Global Actention
Key
Value
Key
Value
Query
Query
Figure 15: Illustration of QKV structure of neighborhood attention (NA) and global attention (GA).
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
8xA100
8xA100
Patchifying
4
Levels (Local Global Attention)
1
2+
Depth
[4, 8]
[2, 2, 12] [320, 640, 1280] [5, 10, 20]
Widths
[768. 1536]
Attention Heads (Width / Head Dim)
[12, 24]
Attention Head Dim
64
64
Neighborhood Kernel Size
7
Integrating & E [0, 1] gives:
f
/' V ve x,)dI.
log log (x1)
Since pi a standard normal distribution, log (x1) can calculated exactly. The second term in the right-hand side of the equation can be computed by solving an ODE:
d
X fe
ve x+) (t, xx)
dt
=
v
with initial conditions Xo and fo 0. Since the divergence operator requires a lot of computational cost, we follow Grathwohl et al. (2018) to use the Hutchinson trace estimator to get an unbiased estimate:
d
x; Je
ve x1) z (ve(t,
dt
where z ~ No, I). In practice, we use the Runge- MunneeKKtta method (Dormand & Prince 1980 to solve the ODE to get fi, and the log probability is Po log Pi (x1) f1-
Normalization A commonly employed trick for training flow models is to linearly transform data point, so that the transformed data are centered around 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 datasct x e p64x64x64 if we normalize it by y (x-0xx 00.04)00.00 00.09 and compute the log- bog-probability of log p(y), then the probability of x can further calculated by:
1
p(x) log p(y) log
0.09
=logp(y) + 64 log
0.09
loog (ly) 631, 228
21

---

Published as a conference paper at ICLR 2025
Table 5: Additional results of the spectral noise denoising task, comparing CRYOFM with Deep-
EMhancer (Sanchez-Gaiee et al., 2021) and EMReady (He et al. 2023).
Estimated resolution different level noise
8.5 A
15.0 Á
FR
FSCAUC f
FSCo.5 +
FR.
FSCAUC t
FSCo.5 +
After degradation
0.3275
8.60
0.2014
15.98
DeepEMMancer
0.13
0.37 0.02
7.4 + 0.3
0.03
0.16 + 0.04
17.5 + 2.5
EMReady
0.00
0.37 0.04
7.7 + 0.8
0.03
0.22 0.02
13.3 + 1.6
CRYOFM
0.00
0.38 0.01
7.5
: 0.3
0.00
0.24 + 0.01
15.0 = 0.7
E
ADDITIONAL INFORMATION OF THE EXPERIMENTS
E. METRICS
ESCAUC The area under the curve (AUC) of the FSC curve (FSCAuc) provides an overall correlation across all spatial frequencies, offering comprehensive view of the alignment between the ground truth and the restored density map Harauu & van Heel 1986).
FSCo.5 The FSC resolution the 0.5 cutoff (FSCo.s) is a standard metric in cryo-EM, indicating the resolution of the reconstructed map relative to the ground truth (Rosenthal & Henderson 2003
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 FSCo.s between the result and the ground truth greater than 30 A. This metric helps identify cases where the reconstruction process fails to achieve meaningful results, ensuring more comprehensive evaluation the method's ' robustness.
Both FSCAuC and FSCo.s reported as the mean and standard deviation across the 32 density maps in the test set, excluding failed cases.
E.2 SPECTRAL NOISE DENOISING
The results of spectral noise denoising at additional resolutions presented in Tab.
F
ESTIMATION OF THE FORWARD OPERATORS
F.1 SPECTRAL NOISE POWER ESTIMATION
Given two half maps from a cryo-EM reconstruction, Y1 and Y2 (observations), where the signal and noise power are the same and the signal and noise uncorrelated, the obeservation models for and y2 are:
y1 V +G,
((v) NN(0, o (v)).
$2 V t 21 NO0 of (v)).
The so-called "tand-Standard" Fourier Shell Correlation (GSFSC) is the FSC between two half maps, which is defined as:
GSFSC(v) & VI(M) 52(x) [ [)))
where (v) represents the component of ý at frequency (i.e., all values on a spherical shell), and (>)) is the radial power spectrum.
From the GSFSC, the signal-to-noise ratio (SNR) can be computed as (Rosenthal & Henderson 2003):
E
signal (D)
GSFSC())
SNR(2)
[V)
GSFSC(V))
22

---

Published as a conference paper at ICLR 2025
Since y1 and Y2 share the same signal and noise each frequency, and the signal V and the noise e are uncorrelated, we have:
[[|l5(>))11 (()) <(v)]1 = Elo (1)) + Elo mosse (v)].
Thus,
Ellly (1)) (1
E [ signal (2) + Elo moiie (V)
Given the known SNR, we can derive:
Bio 2 (v) B 211)2(1)111
SNR(U)
1+ SNR(")
and
B1151(0)113) SNR(U) 21)))(())13] SNR(W)
B Bio signal (v))
1+ SNN((2)
1 (NN(()))
F.2 ANISOTROPIC NOISE POWER ESTIMATION
Anisotropic noise means that the noise is not uniform across all directions; certain directions, the noise power is stronger weaker than in others. In the context of spectral noise, anisotropic noise implies that the variance of the noise depends on the direction Fourier space. We denote by Ç (s) () the anisotropic angular components the sphere, and by the frequency component, corresponding to the radial component in spherical coordinates. We assume that the signal power o² igaa (v) is isotropic, meaning (E, 2) is constant for different values of
The noise power averaged over all directions for given frequency v is denoted as:
h(v)
1 No
Σ moise (5, (),,
where N2, represents the number of voxels in the frequency shell at v.
Our goal is to estimate the anisotropic noise power, ou (5, v). in spherical coordinates. We assume that for given angle Ç. the noise power muise (E, 1/) can expressed as:
Qusise (E, v) aEE posis (Eo.
where to represents an arbitrary reference direction in spherical coordinates, serving as the base- line for the noise power o P (E, v/), with a(E, 2) adjusting the noise power in other directions. In practice, a(, v) 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:
h((v) N,
1 (<(ENv)
Σ
n (Eu, v).
Simplifying, obtain:
2 (Co, v)
h(v))
Σ Ea(s, u)(s
N
where (v)) is the average noise power derived from the the half-map GSFSC in Section F.I Thus, we can solve for oo (En, v/) as:
h(v) No
misee (Co, D)
Finally, the anisotropic noise power for any angle is given by:
h(v) No Ega(E,
Nosso (E, = a(g,v)
23

---

Published as a conference paper at ICLR 2025
Algorithm 2 Flow Posterior Sampling for the Missing Wedge Problem
Require:
a pretrained vector field ve |0, 1] R" - R", forward operator A R" + R" an
obseravation & IRm, number of steps N the recovered signal X0
Ensure:
AL
x1 10,
for t E [1, - At, - 2At, At]] do
Ny f- Xt - t. (e,
X + + A
> Fill the observed part into Xo
- f !- Xx + Al X
end for return Xo
G ADDITIONAL ALGORITTMM
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. [ 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.
G.1 MISSING WEDGE RESTORATION
For certain tasks, such as missing wedge restoration, it is not necessary to compute the gradient with respect the loss function between y and Ax, since we can directly attain the optimum in analytical form. The missing wedge restoration is an in-painting task, where a noiseless part of the signal is already observed. AA is operator extracting part of the signal, hence the inverse of A just fills the observed part into the signal. In Alg. we formulate it as the pseudo-inverse operator Att m IR " (Kaaar al. 2021 2022).
G.2 Ab initio MODELING
For ab initio modeling, the k-thh observed volume (we will name it the projection later) is (M) (V) = PP V), where (k) E SO(3) x R² the pose. The posterior sampling of ab initio modeling more complex than the previous tasks in that the pose ((4) can not be easily pre-determined. The setting is closely related blind diffusion posterior sampling (Chung al. 2023a Kapon et al., 2024 Laroche et al. 2024 Gan et al. 2024), where the forward operator is unknown. We find 2(R) in the sampling process, by minimizing the discrepancy between the projection and the observation
(R) using a coarse-to- fo-rse--t---eeeeeee strategy (Zhong et al. 2021b). The discrepancy is defined by the negative correlation between projection p and an observation y:
p.y y
Ueor (p.y)
|p|| |
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 iol-posed compared others. The first one applies a stop gradient operator (denoted by "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. 2017 Roy et al. 2018 Jang et al. 2016):
p.y
Vag (p.y)
sg([pp]| Iyxl)
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 in the background area of the observation y. The background area is defined by a set of indices B(g.y) y) where the corresponding value smaller the q-th percentile of y.
bbg(g) (p.y)
Σ p
(EB(a.)
24

---

Published as a conference paper at ICLR 2025
Algorithm 3 Flow Posterior Sampling for Ab-initio Modeling
Require: a pretrained vector field ve [0, 1] x R" + IR", K obseravations y E R", number of steps N, maximum step size Amax at each step, the loss function with respect the likelihood
term
Ensure: the recovered signal Xo
At + N
X1 I)
for E [1,1 - At, 18.1-220 2At.... .... At] do
_1_ f Xx At ve(t,x,)
for K E 11, K] do
gik)
arg min Veor (P &o (xx)), y(k)) o x330
Search the optimal pose
(xx1
L (P(())) y(h))
Loss w.r.f the optimal pose
end for 1(x.)=
Ee 1(0) (x:)
[x,)(
g Vx,, (xx))la
X f min Amax: I-6
Xt-At t X, Ar X/g
end for return X0
The final loss function C can be combination of Pooor. v'ag and (bg(g)) 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.
Table 6: Hyperparameters for ab initio modeling
EMPIAR-ID stop gradient
q
10345 (synthetic)
0.7
10028
yes
0.0
10827
no
0.7
H ADDITIONAL ABLATION STUDY
H.1 DATA AUGMENTATION
Protein density data 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, expose the model to enough data with various orientations, appropriate rotational augmentation to the data. In this section, we discuss different data augmenta- tion strategies and how they affect the training and the performance implemente four different augmentation strategies, as shown in Fig. T6 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 24 possible orientations); random rotation: Randomly rotates data in the (0(3) space.
also need to apply
the downstream tasks. We
As seen in Fig. 16 (a), all strategies enabled the model's loss to converge stably on the training set. However, Fig. (b)
(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 Fig. (d), model exhibits similar behavior across different strategies. From the downstream task performance shown in Fig. 16 (e), the differences among the strategies were negligible. Given that our modeling targets a more holistie p(x), we choose random crop combined with random rotation our augmentation strategy.
H.2 PATCHIFYING & DOWNSAMPLING FACTOR
For three-dimensional data, the number of tokens is proportional to the cube of the edge length of the data. Applying patchification to the input and down-sampling within the model significantly im-
25

---

Published as a conference paper at ICLR 2025
(6) Training loss
(b) Test loss on ==200
(c) Test loss on t=800
0.65
0.60
MSE loss
0.55
MSE
0.50
0:32
0:45
0.30
0.40
20000
60000
80000
100000
20000
40000
50000
80000
100000
40000
80000
100000
stop
Step
stop
Center Grop
Crop Roc24
Random Grop
Rot24
Random Crop
(4) Likelihaod estimation log(p)
(e) Downstream performance
Canter
Rescoreni Random Ror24
Ros24
ESCAUE
0.8
Center Crap
Rando Crop
Type
Roc24
Basdom Crap
Train
Rando Rocation
0.5
-1.0
Preferedd
Like
le6
nelie
orieniation
wedge
Figure 16: Results different data augmentation strategies.
Table Different model configurations.
Default
Patchifying variants
Down-sampling variants
Short name
p4-d1
p2-d1
p8-d1
p4-d0
p4-d2
Patching
8
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
pacts 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 patchifying and down-sampling operations. As shown in Tab. 7 the default configuration
used is p4-dl, where p4 represents a patchifying factor of 4, meaning a 4 x 4 patch in the input is flattened and used token vector, and d1 represents that the data is down-sampled once in the model. We compare configurations with p = 2, 4, 8 and = 0, 1,2.
In Fig. 17, present the loss curves the test set during training for different model variants. Given the substantial differences in the computational costs among these model variants, the hori- zontal axis represents the training compute in Gflops. The estimation of training compute is: model Gflops batch size training steps 3 (Peebles & Xie 2023). 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. There- fore, we selected the p4-d1 configuration with comprehensive consideration of model performance and efficiency in downstream applications.
H.3 HYPERPARAMETERS FOR POSTERIOR SAMPLING
Posterior sampling process involves continuously moving forward and correcting the directions of prior p(x) and the likelihood p(xly). In this section, we present ablation experiments on the two parameters that significantly impact the effectiveness the posterior sampling algorithm: the number of total timesteps, and Amax which controls the extent to which is modified based on the observed data. We evaluated the performance using only the first four samples from the test set.
Tab. 8 and Tab. 9 present the ablation studies on these two parameters across three downstream tasks. The results show that for a fixed Xmax, increasing the number of sampling timesteps improves
26

---

Published as a conference paper at ICLR 2025
p8-dl
p4-d2
0.7
0.7-
p4-d1
p4-d1
Test loss
p2-d1
Test loss
p4-d0
0.6
0.6
0.5
0.5
0.4
0.4
10
10
10
10
10
10
Training Compute (Gflops)
Training Compute (Gflops)
Figure 17: MSE loss on test set as function of total training compute.
Table 8: Results with different posterior sampling parameters on spectral noise and anisotropic noise denoising task.
Spectral Noise Denoising
Anistropic Noise Demoising
3.2 Á
4.3 Å
6.1
8.5 A
15.0
+41
+-30
155
After degradation
0.8874
0.5526
0.4968
0.3278
0.2000
0.6626
0.6326
0.6113
Amax
# 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.1140
0.3577
0.31410
0.2156
0.3778
0.2802
0.2576
5.0
200
0.8994
0,6417
0.5808
0.4115
0.2678
0.7923
0.6694
0,5714
5.0
1000
0.9545
0.6843
0.6191
0.3794
0.2313
0.8879
0.8453
0.8053
10.0
50
0.4989
0.4312
0.3904
0.3381
0.2241
0.4012
0.3027
0.2568
10.0
200
0.9025
0.6397
0.5798
0.4103
0.2631
0.7926
0.6713
0,5503
10.0
1000
0.9536
0.6649
0.6183
0.3752
0.2314
0,8780
0.8329
0.7926
quality, with the best performance at 1000 timesteps. Moreover, a Amax of 5 yields optimal or near- optimal results. Therefore, defaulted to Amux of and 1000 timesteps for the posterior sampling parameters in our main results.
Table 9: Results with different sampling steps on missing wedge restoration task.
FSCAUC
After degradation
0.7988
Timesteps = 50
0.9295
Timesteps 200
0.9295
Timesteps = 1000
0.9310
H.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 vaxel space the latent space? Most of the early works (Goodfellow al., 2014 Van Den Oord et al., 2016; Brock 20188] Dhariwal & Nichol, 2021 for generative modeling learned the data distribution directly on the original image space. Rombach al. (2022) pointed out that learn- ing diffusion models on the latent space can generate high-quality images with much less computa- tional cost. The key is to learn an variational auto-encoder (VAE) (Kingmal 2014) which provides low-dimeniional representation that is equivalent to the original data space (Rombach et al.) [2022] Podell et al., 2023). Then, a Diffusion Transformer (DiT) (Peebles & Xie 2023) trained on the latent codes of the VAE.
We have trained several latent flow models (LFMs). An LFM-d model comprised of 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 may ease the training of the DiT with smaller size of latent codes.
27

---

Published as a conference paper at ICLR 2025
Fig. I1 shows the results of unconditional sampling from two latent flow models (LFM): LFM-4 and LFM-8, where the high-frequenge information contains more noise, and secondary structures such a-helicies density modeling, the flow models trained in the voxel space converge much faster and better than the latent flow models.
not be observed. From the preliminary result, concluded that: For protein
Our observation concides with that of Crowson et al. 2024), where they found that latent diffusion models may fail generate fine details. We attribute the failure of latent flow models the weak representational power of the learned variational auto-encoder (VAE) (Kingma (014). The protein density dataset for training the VAE is much smaller than the large-scale dataset in image synthesis tasks (Rombach et al.] 2022 Podell et al. 2023). The low-dimensional space of the VAE is poorly regularized and can be sensitive perturbations. small error the latent space may lead to large details.
peceptual loss in the voxel space, especially for cryo-EM tasks which is highly sensitive to fine
H.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 this work, 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 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:
VAE-d which compresses a protein density V DxDxD a latent code z IR 8 x x d is the compressing factor. The compressing factor governs the balance be- tween perceptual compression and distortion (Blau & Michaeli 2018 Rombach et al.) 2022 We employed the publicly available code Stable Diffusion by replacing 2D convolution layers with 3D layers. We conducted experiments with two settings: d and 8.
a DiT (Diffusion Transformer) (Peebles & Xie, 2023) that learns the distribution of the la- tent codes. We experimented with standard DiT architecture by replacing the 2D patchify- ing module with a 3D module and the patchify factor is set [ The DiT has 12 layers, each of which has 1024 dimensions,
shows that VAE-8 performs significantly worse in reconstrution tasks since it is over-compressed and causes loss of information. However, 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 300k steps.
Section H.4.2
H.4.2 EVALUATION OF VAE
Given a volume E DDDD,D we evaluate the reconstruction quality 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 A FSC(V, VAE(V))
AE-LI A V VAE(V)
Tab. shows that VAE-4 outperforms VAE-8 in terms of reconstruction quality, since VAE-8 learns 512 voxels into latent code, which is 8 times harder than VAE-4
compress grid of 8 that operates on x 4 x 4 64 voxels.
28

---

Published as a conference paper at ICLR 2025
Table 10: Reconstruction quality of VAE
Model
AE-FSC (1)
AE-L1 (4)
Training Steps
VAE-4
0.9821
0.0028
588, 480
VAE-8
0.8537
0.0049
1, 1,797, 872
29


Reviews:
[{"review_id": "FHO5fl60Xy", "reviewer": "Reviewer_a9QU", "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.", "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?\n2) 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?\n3) 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.\n4) How fast is CryoFM compared to the other baselines? \n\n__Additional feedback__\n\n1) 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. \n\n2) Another problem can be found in line 172 which should read \n    $$\\\\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]$$\n    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)$.\n\n3) An error/typo can also be spotted in line 196: \n    $$\\\\nabla_{\\\\mathbf{x}_t} \\\\log p (\\\\mathbf{x}_t | \\\\mathbf{y}) \\\\approx \\\\nabla _{\\\\mathbf{x}_t} \\\\log p (\\\\mathbf{x}_t) - \\\\lambda_t \n\\\\nabla _{\\\\mathbf{x}_t}  \\\\| \\\\mathbf{y} - \\\\mathcal{A} \\\\hat{\\\\mathbf{x}}_0 (\\\\mathbf{x}_t) \\\\|_2^2$$\n  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$.\n\n4) 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\n\n    \\\\begin{aligned}\n    &\\\\mathbf{x}_1 \\\\leftarrow \\\\mathcal{N}(0, \\\\mathbf{I}) \\\\\\\\\n    &\\\\textcolor{red}{\\\\Delta t \\\\leftarrow 1/N} \\\\\\\\\n    &\\\\textbf{for }t \\\\in [1, 1- \\\\Delta t, 1 - 2\\\\Delta t, \\\\ldots, \\\\textcolor{red}{\\\\Delta t}] \\\\textbf{ do} \\\\\\\\\n    &\\\\qquad \\\\mathbf{x}' _{t - \\\\textcolor{red}{\\\\Delta t}} \\\\leftarrow \\\\mathbf{x} _t  - v _{\\\\theta}(t, \\\\mathbf{x} _t) \\\\textcolor{red}{\\\\Delta t} \\\\\\\\\n    &\\\\qquad l(\\mathbf{x} _t) \\\\leftarrow \\\\| \\\\mathbf{y} - \\\\mathcal{A} \\\\hat{\\\\mathbf{x}} _0(\\\\mathbf{x} _t) \\\\|^2_2 \\\\\\\\\n    &\\\\qquad \\\\mathbf{g} \\\\leftarrow \\\\frac{\\\\nabla _{\\\\mathbf{x} _t} l(\\\\mathbf{x} _t) }{ \\\\|\\nabla _{\\mathbf{x}_t } l( \\mathbf{x}_t ) \\\\|_2} \\\\\\\\\n    &\\\\qquad \\\\lambda _t \\\\leftarrow \\\\min \\\\{ \\\\lambda _{\\\\max}, t / (1-t) \\\\} \\\\\\\\\n    &\\\\qquad \\\\mathbf{x} _{ t - \\textcolor{red}{\\\\Delta t} } \\\\leftarrow  \\\\mathbf{x}' _{ t - \\\\textcolor{red}{\\\\Delta t} }  \\\\textcolor{red}{ - } \\\\lambda_t \\\\mathbf{g} \\\\\\\\\n    & \\\\textbf{end for} \\\\\\\\ \n    & \\\\textbf{return } \\\\mathbf{x} _0\n    \\\\end{aligned}\n\n\tSome additional comments on Algorithm 1: \n\t* 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.\n\t* 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$. \n\t* Increasing the loss $l(\\\\mathbf{x}_t)$ does not make sense (most likely a typo or subsequent error of line 196). \n\t* In Algorithm 3, the same errors occur as in Algorithm 1.\n\n5) 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.\n\n6) 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. \n\n7) 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. \n\n# Minor comments\n\n1) 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.\n\n2) 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.\n\n3) Typos/grammar: \n\t* Page 7: \"scaler\"\n\t* Line 525: \"We ... leave for future work.\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* CryoFM is a foundation model for density maps that can be combined with different types of data. \n* CryoFM achieves state-of-the art performance on a number of reconstruction tasks and is demonstrated to be superior to other baselines.\n* CryoFM seems to be faster than other baselines.", "weaknesses": "* The mathematical derivation of the algorithm seems to be a bit sloppy. \n* Tests with experimental data are missing. \n* No source code is provided (it would also help clarify some of the mathematical shortcomings of the paper). \n\n__Recommendation__\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KFKiRJHn0B", "reviewer": "Reviewer_nPzD", "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.\n\nThis 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.\n\nThe 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).\n\n**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\").  \nI will raise my score to \"Accept\" if my concerns related to the paper's clarity (see \"Weaknesses\") are sufficiently addressed.\n\n**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.\nI 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.", "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)?\n\n\nNote that I also raised some questions in the \"Weaknesses\" sections. They seemed more appropriate there.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n- 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.\n- 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.\n- For ab-initio modeling, the authors tested CryoFM on two real-world datasets.\n- All figures are of high quality.", "weaknesses": "In many places, the paper is hard to understand and details and context are lacking. For example:\n\n- 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:\n\t- 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.\n\t- 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).\n\n - 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. \n\n- Details and context are missing in some places in the experiments. For example:\n\t- 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.\n\t- In the denoising experiments in Sections 5.1 and 5.2, it is not clear what the test set is.\n\t- 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?\n\nOther weaknesses:\n- 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.\n\n- 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:\n\t- Why did the authors use CryoSPARC as baseline? \n\t- Is there another ab-initio reconstruction method that works with class averages?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yEJGvu8UfR", "reviewer": "Reviewer_xQqb", "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.", "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.\n2. Regarding Figure 7, it would be interesting to see if using the same dataset with CryoSPARC actually improves performance.\n3. Curious about effectiveness of the 2D to 3D reconstruction approach, particularly when only a limited set of 2D views is available", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The development of a foundation model for cryo-EM is both significant and novel\n2. 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\n3. Overall, the paper is well organized", "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rIFzZbi9cg", "reviewer": "Reviewer_esMq", "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.", "questions": "Most important questions:\n- 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.\n- 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.\n\nMinor questions:\n- Does the FM in the title stand for “foundation model”, “flow matching”, or both?\n- 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?\n- What is the motivation behind the specific data augmentations that were used for CryoFM-S and CryoFM-L?\n- Figure 9 does not show results for ab initio reconstruction; why?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.", "weaknesses": "- The motivation for using a flow-based model is not very clear\n- 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).\n- 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.\n- 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).\n- 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.\n- 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. \n\nMinor points\n- 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.\n- There is a typo on line 231; Transformer is misspelled.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)