Mixture of neural fields for heterogeneous reconstruction in cryo-EM

Authors: axel  levy
rishwanth  raghu
david  shustin
adele  rui
yang  peng
huan  li
oliver  biggs
gordon  wetzstein
ellen d zhong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:13:12

Abstract: Cryo-electron microscopy (cryo-EM) is an experimental technique for protein structure determination that images an ensemble of macromolecules in near-physiological contexts. While recent advances enable the reconstruction of dynamic conformations of a single biomolecular complex, current methods do not adequately model samples with mixed conformational and compositional heterogeneity. In particular, datasets containing mixtures of multiple proteins require the joint inference of structure, pose, compositional class, and conformational states for 3D reconstruction. Here, we present Hydra, an approach that models both conformational and compositional heterogeneity fully ab initio by parameterizing structures as arising from one of K neural fields. We employ a hybrid optimization strategy and demonstrate the effectiveness of our approach on synthetic datasets composed of mixtures of proteins with large degrees of conformational variability. We additionally demonstrate Hydra on an experimental dataset imaged of a cellular lysate containing a mixture of different protein complexes. Hydra expands the expressivity of heterogeneous reconstruction methods and thus broadens the scope of cryo-EM to increasingly complex samples.

Main paper: 
 introduction: 
 Structural information is key to understanding the function of macromolecular complexes, making protein structure determination a crucial tool in basic structural biology and rational drug design. Among experimental structure determination methods, cryogenic electron microscopy (cryo-EM) is unique in its capability to reveal dynamic information about large macromolecular complexes in near-native states.In single particle cryo-EM, a set of biomolecular complexes (i.e. particles) is flash-frozen and imaged with a transmission electron microscope. Each collected image consists of a noisy, randomly oriented projection of an individual particle of unknown identity or composition, frozen in an unknown state. Reconstruction algorithms processing these data without information from upstream algorithms are called ab initio heterogeneous reconstruction methods. Classical reconstruction techniques use a Bayesian approach to optimize a finite number of voxel-based representations [50, 45] . These algorithms enabled biologists to process datasets made of a mixture of different protein compositions and, thus illuminate the molecular details of fundamental biological processes. These methods, however, tend to aggregate all the proteins of the same composition in a single class, despite these proteins being trapped in different conformational states, due to thermal fluctuations prior to the freezing step. Another line of work has extensively studied the possibility of reconstructing continuous motion from cryo-EM datasets, using linear combinations of voxel arrays [43] , neuralbased representations [68, 69, 25, 26] , Gaussian mixture models [5, 20] , or combining a voxel array with a flow field [44] . These methods sometimes leverage a structural prior about proteins (e.g., proteins are made of a fixed number individual atoms) and enable the reconstruction of molecular motion at high resolution. However, none of these methods can handle datasets containing different types of proteins, thereby strongly limiting their application, especially in the context of in situ cryo-EM.As of today, there exist no approaches to simultaneously reveal compositional (discrete) and conformational (continuous) heterogeneity in cryo-EM datasets, potentially limiting the structures that can be revealed from the data. Unraveling this information poses a nontrivial problem that sequential strategies cannot solve. Due to their low signal-to-noise-ratio, cryo-EM images cannot be clustered depending on the type of protein they contain prior to the reconstruction. Because orientations are unknown, the problem cannot be solved by handling compositional heterogeneity before conformational heterogeneity as consensus poses (i.e., orientations) are inaccurate for large motions.Here, we introduce Hydra, a neural-based method for ab initio heterogeneous reconstruction in cryo-EM. Inspired by the success of implicit neural representations in cryo-EM, we extend the neural field representation of DRGN-AI [26] with a mixture model of K neural fields. Using a new likelihood-based loss function, we simultaneously optimize orientations, conformations and class assignments and circumvent the pitfalls of sequential approaches. We demonstrate that our method allows neural-based methods to handle strong compositional heterogeneity and enables the simultaneous reconstruction of compositional and conformational heterogeneity with state-of-the-art accuracy. In an experimental cryo-EM dataset of a cell lysate mixture, we reveal three compositional states in a single pass, fully ab initio. We therefore make the following contributions:• We develop a mixture of neural fields model for ab initio heterogeneous reconstruction in cryo-EM; • We demonstrate that our method improves the expressiveness of neural-based methods for handling strong compositional heterogeneity; • We enable simultaneous reconstruction of conformational and compositional heterogeneity beyond the limits of existing methods; • We demonstrate the reconstruction of multiple protein complexes from an experimental dataset of an unpurified sample.2 Related Work 2.1 Heterogeneous Reconstruction in Cryo-EM Discrete Variability. Cryogenic electron microscopy offers the potential to reveal the structure of macromolecules in heterogeneous samples, where multiple types or multiple conformations are mixed together [24] . The first reconstruction algorithms handling 3D variability modeled the set of particles as a finite set of static structures. RELION [50] popularized the Bayesian approach to tackle heterogeneous reconstruction and the later introduced multi-body refinement tool [34] offered the possibility to segment static density maps and model continuous motion as a combination of rigid transformations. The software suite was recently improved with the Blush regularization tool [22, 4] , leveraging a data-driven prior to enable the reconstruction of small protein-nucleic acid complexes (≤ 40 kDa). CryoSPARC [45] accelerated the inference with stochastic gradient descent and, in the 3DVA [43] extension, modeled continuous motion with a linear combination of density maps. These methods represent the state of the art for cryo-EM reconstruction but do not have the capability to represent complex non-linear motion and tackle the discrete heterogeneity (e.g., with "multi-class ab initio") before the continuous variability (e.g., with 3DVA or [21] ), leading to inaccurate pose estimation for states exhibiting large motion.Non-Linear Methods for Heterogeneous Reconstruction. In the past years, significant progress has been made in revealing non-linear dynamics from cryo-EM data. Manifold embedding [11] and Laplacian methods [33] are among the first attempts to model non-linear continuous motion but have only been applied to a small number of macromolecular complexes [8, 7] . HEMNMA [20] used a decomposition over the low-energy normal modes of an atomic model, thereby leveraging a prior over the structure and the dynamics of macromolecules. Herreros et al. [17] proposed the use of 3D Zernike polynomials and eliminated the need for pseudo-atomic models. CryoDRGN [69] used neural networks to continuously represent deformable density maps as well a variational auto-encoding framework to estimate conformational states. E2GMM [5] , cryoFold [71] , and DynaMight [51] followed this encoder-decoder framework and used Gaussian mixture models to represent density maps, thereby reducing the memory footprint of previous non-linear methods. 3DFlex [44] introduced the use of a parametric flow field to smoothly deform canonical 3D density maps, leveraging the knowledge that energetically favorable deformations tend to preserve the local geometry of proteins. Although these methods demonstrated the ability to reconstruct molecular motions and heterogeneity, they all need a coarse initialization of the density map, or the poses to be provided by an upstream reconstruction algorithm. In practice, this ab initio step is error prone in the presence of large conformational motions.Ab Initio Heterogeneous Reconstruction. Recent works investigated the problem of reconstructing an ensemble of density maps where poses are unknown. The preliminary cryoDRGN method [68] tackled the ab initio reconstruction problem by combining traditional pose search algorithms with the encoder-decoder neural-based architecture, which was refined in cryoDRGN2 [70] . Multi-CryoGAN [15] sidestepped pose estimation by casting the reconstruction problem as a distribution matching problem and successfully revealed 3D variability in synthetic cryo-EM datasets. Rosenbaum et al. [49] showed that the auto-encoding framework could be applied to jointly estimate poses and conformations of atomic models, and Levy et al. [25] demonstrated a fully autoencoding framework for ab initio heterogeneous reconstruction on real benchmark datasets. DRGN-AI [26] recently introduced a hybrid pose search strategy combined with an autodecoding architecture to handle low-signal datasets and demonstrated ab initio heterogeneous reconstruction on challenging datasets containing a significant number of junk images. These methods extend neural-based reconstruction to scenarios where poses cannot be reliably estimated by any upstream algorithms, but have a limited capability to represent mixtures of biomolecular complexes, due to the limited capacity of a single neural representation.Here, we propose to represent the landscape of accessible structures via an ensemble of K neural networks, each specialized in representing the variability within a single compositional state. By jointly estimating structures and orientations, our method handles datasets that exhibit strong compositional heterogeneity and large motions. 
neural fields for large and dynamic scenes: 
 Dynamic Scenes. In graphics, neural networks have also been used as light-weight, differentiable representations of continuously defined signals (e.g., occupancy fields [31] , signed distance functions [37, 6] , surface light fields [66] , latent representation of appearance [52, 53] , radiance fields [32, 1] , and light fields [54] ). In order to handle time-dependency and represent dynamic scenes, two approaches have been explored, as described in [39] . Similar to the cryo-EM method 3DFlex [44] , deformation-based approaches apply a spatially-varying deformation to some canonical radiance field [38, 42, 59] . Methods following this strategy recover detailed dynamic scenes but are unable to model any topological variations. Similar to cryoDRGN [69] , modulation-based approaches directly condition the radiance field of the scene on some latent vectors encoding temporal or dynamic information [13, 27, 65, 12] . These techniques are capable of modeling arbitrary deformations, topological changes, and other complex phenomena, but require additional regularization strategies to avoid trivial, non-plausible solutions. HyperNeRF [39] was introduced as a combination of both approaches and enabled photo-realistic reconstruction of dynamic scenes with topological changes.Large-Scale Scenes. Due to their limited capacity, neural networks have essentially been used to represent single objects or small-scale scenes. Several methods [46, 47] have looked into the possibility of using an ensemble of neural fields to represent large scenes such as a neighborhood in the city of San Francisco [57] . Another approach is to provide extra capacity with a coarse 3D grid of latent codes [27, 55, 40] or a block-coordinate multi-scale decomposition [28] . These works focus on static reconstruction, and if dynamic objects are present in the scene, these are simply masked out [57] . Ost et al. [36] enabled the representation of complex, dynamic multi-object scenes by decomposing them into their static and dynamic parts and learning one neural radiance field per dynamic object.Here, we use the modulation-based approach to handle the continuous motion of proteins and an ensemble of neural fields to increase the representational capacity of our model, allowing it to reconstruct compositional mixtures of proteins. In contrast with the works cited above, we need to estimate which compositional state each image belongs to -which is done using a variational approach -while simultaneously optimizing the ensemble of neural networks and individual image poses.A mixture of experts (MoE) model uses an ensemble of neural networks to process a given input [19] . Doing so, the input space can be partitioned into subspaces on which only one of the neural networks needs to become "expert". The mechanism with which the input space gets partitioned is usually referred to as the gating mechanism. Masoudnia and Ebrahimpour [29] partitions MoE approaches into two groups, depending on the gating mechanism and both how and when this mechanism is involved: (1) mixtures of explicitly localized experts (MELE) cluster the input data before the experts' training phase starts while (2) mixtures of implicitly localized experts (MILE) jointly optimize the expert networks and the gating mechanism. In cryo-EM, clustering methods operating directly on images are usually ineffective because of the high level of noise and because of the presence of other nuisance variables (orientation and conformation), making the MELE approach irrelevant.Jacobs et al. [19] examined the use of different error functions (or loss functions) to optimize the expert networks and the gating mechanism, and the best performance was obtained using the negative log-probability of a Gaussian mixture model [35] . Several architectures have been explored, and one of the most applied is the mixture of MLP-experts (MME) [63, 9] , where multi-layer perceptrons (MLPs) are used for both the experts and the gating networks. Our approach can be viewed as an application of the MILE method where, as in Jacobs et al. [19] , the error function corresponds to the negative log-probability of a Gaussian mixture model. Due to the high level of noise, the gating mechanism ignores the content of input images and relies on an autodecoding framework.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The limitations are mentioned in the discussion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 In this section, we first describe the image formation model in single particle cryo-EM (3.1). We then define a 3D variability model that models both compositional and conformational heterogeneity (3.2). Based on these two first sections, we describe our system with a latent variable model (3.3) and explain our optimization method (3.4). The method is schematically described in Figure 1 .In cryo-EM, a purified solution of macromolecules is flash-frozen inside a thin layer of vitreous ice. Each molecule gets trapped in a random orientation with respect to the microscope and in a random conformational state, approximately following the Boltzmann distribution at ambient temperature [2] . The sample is exposed to an electron beam and individual 2D projections are extracted from micrographs via a step known as "particle picking". The reconstruction task therefore starts with a given set of N images (or particles). Each image I i can be modeled as [60, 50] where C i models the the Contrast Transfer Function (CTF), V i is a scalar 3D field (R 3 → R) known as the "electron scattering potential" or the "density map",EQUATIONϕ i = (R i , t i ) is a "pose" (R i ∈ SO(3), t i ∈ R 2 )and P projects V i on a 2D grid:EQUATION)η i models isotropic Gaussian noise (η i ∼ N (0, σ 2 )).In a typical experiment, N can vary between 10 5 and 10 7 ; The signal-to-noise ratio can vary between 10 -1 and 10 -2 .Structural heterogeneity among the macromolecules can originate (1) from continuous motion along a small number of degrees of freedom, or (2) from discrete compositional changes. We refer to the first kind of heterogeneity as "conformational heterogeneity" and to the second one as "compositional heterogeneity".To model this mathematically, we make the assumption that all the density maps V i belong to a finite union of low-dimensional manifolds:EQUATIONwhere K ∈ N, d ∈ N and, for all k, θ k ∈ Θ is an unknown parameter. In other words, we assume that there exist K compositional states and that the conformational motion of state k can essentially be described with d degrees of freedom.We statistically describe the set of observed images with a latent variable model. Here, the latent variables are the poses ϕ i , the conformations z i and the class identities k i , while {θ k } K k=1 is a set of shared parameters.Given the image formation model described by Eq. 1 and the heterogeneity model described by Eq. 3, each observed image can be seen as a sample from a multivariate mixture model:EQUATIONWe parameterize a probability distribution over the class identity, such that:EQUATIONwhere s i ∈ R K is a free parameter called the "score". For each k, we parameterize point estimates of the continuous latent variables:EQUATIONUnder the model in Equation ( 4) and the variational parameterization described above, the negative log-likelihood of an image is given byEQUATIONIn our implementation, each low-dimensional manifold V(.; θ k ) is implemented with a residual MLP.We provide further details on the architecture of the neural networks in the supplementary materials.We aim at minimizing the negative log-likelihood of the set of observed images,EQUATIONover({ϕ i,k }, {z i,k }, {s i }, {θ k }).All the parameters are optimized using a combination of gradientbased optimization and exhaustive search. Here, we use an autodecoding framework and do not amortize the inference of latent variables with an encoder, following the observation in Levy et al. [26] that encoders tend to memorize images in highly-noisy setups.We handle the minimization over {z i,k }, {s i } and {θ k } with stochastic gradient descent. The minimization over the poses {ϕ i,k } is more challenging, due to the existence of local minima. We take inspiration from the two-step pose estimation strategy introduced in DRGN-AI [26] and adapt it to cope with the simultaneous representation of multiple maps. First, optimization over poses is done with hierarchical pose search (HPS) in alternation with stochastic gradient descent on the other variables. On a given random minibatch of indices I ⊂ {1, . . . , N },EQUATIONwhere Ω is a predefined grid in the space of poses (SO(3) × R 2 ). See [26] and the supplementary for more details on how the minimization can be done efficiently with a hierarchical strategy. This robust but computationally-expensive strategy is used for a predefined number of steps, usually between 5 × 10 5 and 2 × 10 6 . After that, most poses are located in the basin of attraction of their global optimum and switching to stochastic gradient descent (SGD) makes computation more efficient while improving pose accuracy (poses are not constrained to belong to a predefined grid of fixed resolution during SGD).At the end of optimization, the estimated class of image I i is simply given by the index k i of the largest entry in s i ∈ R K .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: We provide this information in experiments section and in the supplementary material.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
experiments: 
 We run three experiments to evaluate Hydra. In Section 4.1, we show that our method improves the expressiveness of neural-based methods and can reveal strong compositional heterogeneity fully ab initio. In Section 4.2, we use Hydra to reveal compositional heterogeneity in an experimental dataset containing protein complexes of diverse sizes, in a single run. In Section 4.3, we demonstrate the simultaneous reconstruction of compositional and conformational heterogeneity. The classification accuracy is evaluated for each method using the adjusted Rand index (ARI) [18] . To evaluate each method's reconstruction quality, we use the mean area under the Fourier shell correlation (FSC) curve for 20 images per class (we report ±1 standard deviation). We bold the best result, and underline the second best result.We evaluate Hydra on tomotwin3, a synthetic dataset of 3,000 images emulating a protein sample containing multiple species with static structures. We selected the 6th, 7th, and 8th largest proteins by atomic weight from the TomoTwin training dataset [48] , which correspond to entries 6up6, 6id1, and 4cr2 in the RCSB PDB (Protein Data Bank) [3] . We rendered 1,000 synthetic images of each specimen following the cryo-EM image formation model. For each protein structure, we simulated density maps with the ChimeraX molmap command [30, 58] . We padded the density maps to a box size of D = 384 pixels and centered them. We sampled 1,000 orientations per ground truth class uniformly from SO(3). We projected the density maps using Eq. 1 and applied a translation vector uniformly chosen from a 30-pixel-wide box centered at the origin. We applied a CTF sampled from EMPIAR-11247 [10] and added Gaussian noise (SNR = 0.01). We downsampled each image to a box size of D = 128, yielding 3,000 images in total (samples shown in Figure S1 ).We compare the performance of Hydra to several state-of-the-art ab initio methods for resolving compositional heterogeneity in cryo-EM datasets, including DRGN-AI [26] , cryoDRGN2 [70] , and cryoSPARC [45] . Hydra only uses 2 dimensions for the conformational space, while DRGN-AI and CryoDRGN2 are evaluated with a more expressive 8-dimensional space. We test Hydra using the correct number of classes (K = 3, the true number of specimens present in the sample) and using an over-parameterized setup (K = 5 in Figure 2 and K = 7 in Figure S2 ). For DRGN-AI and CryoDRGN2, we classify each image using k-means clustering on the conformational space with 3 clusters. We evaluate per-image FSC for each class using the dataset's ground truth class labels.In Table 1 , we report results for each configuration of Hydra and other methods, (best of three replicas with distinct seeds for Hydra and DRGN-AI, full results in Table 3 ). We observe that other methods with implicit neural volume representations, including DRGN-AI and CryoDRGN2, fail to capture the compositional heterogeneity of the dataset. Hydra matches the classification quality of cryoSPARC, which models independent, discrete mixtures. Hydra also outperforms all methods on per-image FSC, a distributional volume reconstruction quality metric based on the Fourier Shell Correlation (FSC) [68] . We provide additional results with a larger version of this dataset in Figure S3 and additional metrics (including roto-translation accuracy) in Table 4 .We evaluate our method on an experimental dataset of a protein mixture from red blood cell lysate subjected to density-gradient centrifugation followed by chemical cross-linking. Through an exhaustive, expert-driven data processing pipeline in cryoSPARC [45] , we determined that the dataset consisted of substantial compositional heterogeneity from three main component protein complexes: the membrane proteins ryanodine receptor (RyR) and mitochondrial respiratory chain complex III (CIII), and a dimeric complex of the soluble valosin-containing protein (p97). A sweep of K-values from cryoSPARC ab-initio reveals that K ≥ 5 is necessary in order for cryoSPARC to reconstitute distinct RyR, p97, CIII, and junk classes (Figure S4 ).In Figure 3 , we compare the performance of DRGN-AI [26] with Hydra. Figures 3d and S5 show that DRGN-AI successfully partitions the dataset into two clusters corresponding to the RyR and non-RyR particles; however, DRGN-AI is unable to learn distinct shapes for all three discrete structures and instead appears to learn structural artifacts from the non-RyR particles that match the overall RyR shape. By contrast, Hydra with K = 4 successfully separates the particles into the three component protein complexes and a fourth junk class (Figures 3 and S6 ). As a baseline for comparison, we also carried out the typical workflow of first generating poses from a consensus reconstruction in cryoSPARC, followed by fixed-pose DRGN-AI to separate the heterogeneity. As can be seen in Figure S7 , DRGN-AI recovers a high resolution RyR density but fails to learn non-RyR protein structures when training from poses from homogeneous reconstruction. A multi-step reconstruction with DRGN-AI does not either reveal the CIII class (Figure S8 ). We prepare a synthetic dataset exhibiting both compositional and conformational heterogeneity (ribosplike). For each of the pre-catalytic spliceosome, 80S ribosome, and SARS-CoV-2 spike protein, 5,000 images are generated from different conformational states of the macromolecule along a 1-dimensional trajectory of motion [41, 64, 61] yielding a dataset of 15,000 images (See SI Section F for additional details and Figure S1 for sample images). As seen in Figure 4 , our method with K = 3 nearly perfectly separates the particles into their three classes. Within each neural field, the conformational change of the dominant particle type is well captured both by the conformational space. Due to the trimeric structure of the spike protein, we observe three continuous manifolds in the conformational space, though the majority of spike particles are aligned to one of the symmetries. Tables 2 and 5 compare the performance of Hydra with other ab initio baselines, demonstrating superior performance in classification accuracy as measured by the Adjusted Rand Index (ARI) and volume reconstruction quality (per-image FSC). It additionally achieves the lowest pose error amongst all methods. Representative reconstructions for baseline methods are shown in the supplementary material (Figure S9 ). 
discussion: 
 This work introduces Hydra, a method that expands the expressiveness of neural-based models for heterogeneous reconstruction in cryo-EM. Hydra models structures as arising from 1 of K neural fields and is designed to capture heterogeneity in datasets containing a mixture of multiple species. Notably, our method runs fully ab initio, i.e. does not rely on pre-estimated poses, which are not straightforward to obtain in cases of strong compositional heterogeneity. We evaluated Hydra on both synthetic and experimental datasets, showed improved performance over previous neural-based methods on compositionally heterogeneous datasets, and demonstrated that it could simultaneously handle compositional and conformational heterogeneity.Our method can be seen as an instance of a mixture of experts model, where the gating mechanism is handled by an autodecoding framework, due to the low signal-to-noise ratio in the input data. However, predicting classes with a neural network can be viewed a classification task and may be easier than pose or conformation estimation. We view the possibility of using a (potentially pretrained) neural network for classification, thereby making use of amortized inference, as an exciting avenue for future work.One important limitation of our approach is the need to specify the number of classes K prior to reconstruction. Other methods that cope with compositional heterogeneity, like RELION [50] and cryoSPARC [45] , possess the same requirement, and this is solved in practice by running the algorithm several times with a sweep on the hyperparameter K. As this process is time and energyconsuming (a single reconstruction can run for up to 4 GPU-days on high-end GPUs), more efficient methods for hyperparameter selection, for example, with an adaptive and/or hierarchical strategy, is an impactful direction for future work. We would also like to note that, in its current implementation, Hydra must give the same latent dimension (d) to all the classes, but this constraint could be relaxed with, for-example, a dictionary-based representation for the conformations.With Hydra, we broaden the scope of datasets that cryo-EM reconstruction algorithms can process. Our method would be especially suitable for data collected in situ (i.e., inside the cell), where mixtures of dynamic biomolecular complexes coexist. This data is usually acquired by cryogenic electron tomography (cryo-ET), where the sample is progressively tilted throughout the data collection process.Extending our method to enable subtomogram averaging of cryo-ET data is another potential future direction that we hope can enable future discoveries in structural biology and facilitate the design of new therapeutic compounds.We evaluated Hydra on tomotwin3 with a 2-dimensional conformational space. We perform HPS on 100,000 images (33 epochs), followed by 100 epochs of SGD pose optimization. We set the batch size to 64 during SGD pose optimization. The score table learning rate is set to 0.01, and σ is set to 0.1. We evaluate Hydra using both the correct number of classes (K = 3) and an overparametrized configuration (K = 5). All other parameters are set to default values.We train CryoDRGN2 v3.3.0 for 30 epochs using an 8-dimensional latent space, an encoder width of 1024, 3 encoder layers, and a decoder width of 1024. We evaluate DRGN-AI using an 8-dimensional conformational space, 100k images (33 epochs) of HPS followed by 100 epochs of SGD pose estimation. cryoSPARC ab initio was run with K = 3 classes. All other parameters are set to default values.When calculating volume metrics, we compare the reconstructed density map to a downsampled ground truth density map (D = 128 pixels). All experiments on tomotwin3 are run on one NVIDIA A100 GPU. Our experiments required 4h00min for K = 3 and 6h20min for K = 5. DRGN-AI ran in 1h20min and cryoDRGN2 in 1h50min.Simulation of conformational and compositional heterogeneity. For the pre-catalytic spliceosome, we obtain a trained cryoDRGN model from Zenodo2 and generate 500 density maps of box size D=256 at equally spaced points along the first principal component of the latent space [41] [69] .For the 80S ribosome, we likewise obtain a trained cryoDRGN model from Zenodo and generate 500 density maps of box size D=256 at equally spaced points along a linear trajectory connecting the embeddings of particles 34570 and 60629 in the latent space [64, 69] . For the SARS-CoV-2 spike protein, a cryoDRGN model was trained on a filtered subset of the dataset of Walls et al. [62] , for 25 epochs with a circular image mask of dimension 64, latent dimension of 8, and all other hyperparameters at default. Then, we generate 500 density maps of box size D=256 at equally spaced points along the first principal component of the latent space [61] .We next match the pixel sizes of the density maps by standardizing them to that of the spliceosome.In particular, we downsample the ribosome density maps to D=228 and pad back to D=256, and downsample the spike protein density maps to D=198 and pad back to D=256. Then, we generate and apply a soft mask for each density map, at thresholds of 0.03, 0.05, and 0.1 for the spliceosome, ribosome, and spike protein, respectively, and with 25Å of dilation and a 15Å cosine edge. Next, we normalize density map intensity values such that the signal ratios between different particle types match approximate "true" signal ratios. We calculate the true signal ratios by running the ChimeraX molmap command on the PDBs of each particle (spliceosome 5NRL, ribosome 3J79, spike 6VXX and 6VYB), summing the density within each density map, and calculating signal ratios (where the signal of the spike protein is taken to be the average over the two PDBs). Then, the intensities of the 500 ribosome density maps and 500 spike protein density maps are scaled such that the ratios of total intensities across all density maps for ribosome to spliceosome, and spike to spliceosome, match the true ratios.Finally, we generate images from all the density maps. For each density map, 100 projection images are generated with uniformly sampled rotations from SO(3) and uniformly sampled inplane translations within [-20, 20] pixels. CTF is applied to each image, with values drawn (with replacement) from the experimental CTF values of the spliceosome dataset [41] . Noise is added to all images at an SNR of 0.1, with the variance of the signal computed over all 15k particles. Lastly, the images are downsampled from D=256 to D=128.We train Hydra for 100,000 images of HPS before 100 epochs of SGD. We train cryoDRGN2 for 90 epochs. All other hyperparameters for all methods are set to their defaults. For single-class DRGN-AI and cryoDRGN2, the predicted classes for Adjusted Rand Index calculation are obtained by k-means clustering the latent space with k=3. Image-FSC is computed over 50 images equally spaced along the true conformational trajectory, per each of the 3 particle types (150 images total). Experiments for Hydra were run using 2 NVIDIA V100 GPUs, and experiments for baselines were run using 1 NVIDIA V100 GPU.Dataset details. 85,656 particles were manually picked from one round of 2D classification on cryoSPARC of a larger dataset of 148,596 particles. The particles were downsampled from a box size of D = 300 pixels and 1.66 Å/pixel to D = 150 pixels at 3.32 Å/pixel.DRGN-AI and Hydra. All training for DRGN-AI and Hydra were carried out on 4 A100 NVIDIA GPUs with 80GB memory. For ab initio DRGN-AI, we ran DRGN-AI with default parameters (latent dimension d = 8), with 500k images for HPS followed by 100 epochs of SGD. We verified that single-class DRGN-AI could separate the RyR and non-RyR particles by selecting the particles belonging to the non-RyR cluster (as visualized in the latent space) for downstream single-class DRGN-AI analysis, and a second run of single-class DRGN-AI with default parameters, 500k images of HPS and 100 epochs SGD on the non-RyR particles recovered a p97 density map and a non-p97 cluster in the latent space (followed by a third round of DRGN-AI with the same parameters on the non-p97 particles). For Hydra, we trained the model with 2 million images for HPS followed by 100 epochs of SGD and latent dimension d = 2. After an initial sweep of the hyperparameters for learning rate and high-entropy prior σ using K = 3 or K = 4, the optimal parameters for this dataset were determined to be a learning rate of 0.1 and σ = 10.0. In order to reduce the resource demand of Hydra, we lowered the hypervolume dimension to 128 and reduced the number of hidden layers in the hypervolume to 2, as well as decreased the SGD batch size to 64. For fixed-pose DRGN-AI with consensus cryoSPARC poses, poses from a cryoSPARC homogeneous refinement job were used as input for optimization by single-class DRGN-AI with default parameters.CryoSPARC processing. We performed a sweep of K-values for 3D classification using cryoSPARC ab initio from K = 1 through K = 8 (all other parameters set to default values), and determined that K < 5 is insufficient to separate the junk particles from protein. K = 6 ab initio yielded the best separation of particles into the four different classes, with three different RyR density maps produced.We performed heterogeneous refinement on the best ab initio K = 6 replicate for comparison to Hydra. In accordance with best practices for data processing, we performed ab initio with the downsampled D = 150 particles but used undownsampled particles boxed at D = 300 to generate the most accurate poses during refinement against the low-resolution density maps generated by ab initio.We provide the following supplementary files:• Three movies illustrating the continuous motion of the ribosome, spliceosome and spike, as shown in Figure 4 . The movies were obtained by sampling 9 points on a linear trajectory of the latent space, for each compositional state. • The 10 volumes shown in Figure 4 for the ribosome and the spliceosome, in mrc format. Reconstructed maps on the left and confusion matrix on the right. Hydra is able to accurately reconstruct the three states in the dataset but four classes end up empty. Guidelines:Predicted Class• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: To the best of the authors' knowledge, the paper respects the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We do not foresee any potential malicious applications or uses of the work described in this paper.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: The paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The structures from the PDB are properly credited. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.https://zenodo.org/records/4355284 
 
 Appendix:  


Decision: This paper received mixed reviews.  The authors submitted a rebuttal with some additional experimental results and there was an active discussion period with many of the issues noted in the initial reviews being addressed.

The main issues raised were around datasets and novelty.  The issues around datasets were largely resolved and clarified during the discussion phase.  The question of novelty is somewhat subjective with two reviewers (RoiP and SenM) arguing that the work makes a sufficient contribution.  After reviewing the comments and the paper, the AC agrees with RoiP and SenM that the work presents a novel and significant improvement over existing work and therefore recommends acceptance.

Reviews: review: paper_summary: paper_summary: This submission presents Hydra, a method for handling heterogeneous cryo-EM reconstruction. Hydra can model both conformational and compositional heterogeneity and can perform ab initio reconstruction. To achieve this, it parameterizes structures as arising from one of K neural fields. In the optimization pipeline, the conformations, poses, class probabilities, and neural fields are optimized to maximize the likelihood of the observed images. The authors demonstrate the reconstruction of multiple protein complexes from an experimental dataset.
strength_weakness: strength_weakness: 1. Compared to DRGN-AI, the proposed approach primarily changes the single neural representation to multiple ones.
2. The determination of K seems tricky.
3. The paper lacks comparisons with conventional approaches such as cryoSPARC and RELION, especially qualitative comparisons. More results, preferably video results, are needed.
4. It is unclear whether the pose predictions are accurate. More visualization analysis would be helpful.
questions: questions: 1. How should the value of K be determined? How would the proposed model behave if K is too large?
2. Will the multiple neural representation design lead to excessive partitioning, that is, splitting complete independent structures into unreasonable parts?
3. How does the proposed approach compare to conventional methods like cryoSPARC and RELION in terms of reconstruction quality and efficiency?
limitations: limitations: n/a
paper_summary: paper_summary: The paper presents a neural network-based methodology for modeling both compositional and conformational protein states in cryo-electron microscopy (Cryo-EM) 3D reconstruction. 

In particular, the authors propose a fully *ab initio* approach, named Hydra, which enables the joint inference of poses, conformations, and class identities.

The novelty of this approach lies in its ability to capture both discrete (compositional) and continuous (conformational) heterogeneity within Cryo-EM datasets, without relying on pre-computed pose estimations. Previous methods have either struggled with accurate pose estimation, relied on coarse initializations and upstream algorithms, or had limited capacity to represent complex biomolecular mixtures.

Moreover, the authors validate their proposed approach by comparing it against other popular methodologies, using both synthetic and real datasets, showing the potential of Hydra to capture both compositional and conformational heterogeneity.
strength_weakness: strength_weakness: Although the manuscript is generally well-written, it may be hard to follow for someone who is not an expert in the field. In particular, there are some areas where accessibility to a broader readership might be possible:
* For me, reading some previous work was necessary to understand the context of this work sufficiently to appreciate and understand the contributions of this manuscript. A clearer introduction tailored to a broader readership could make this submission more self-contained, which I would find desirable.
* The introduction lacks a straightforward definition of the taxonomy used throughout the paper. While the authors introduce the context of their research, in particular regarding cryo-EM, I think that the reader should be introduced to the concepts of “poses”, “conformational states”, and “compositional states”, and then to why they are relevant to the presented research and future users. Casting this into a concise but clear way will be much appreciated by future readers, I believe. After that, the authors address how their technical contributions addressed the main challenges presented to them. In my opinion, such changes will allow non EM experts to appreciate the presented work much better and potentially allow other fields to benefit from the same/similar ideas and methods.
* The proposed method seems to rely heavily on the DRGN-AI approach from Levy et. al. Although the authors explicitly state in the introduction that this paper represents an extension of that method, throughout the paper it’s not always clear whether the methodological choices described are novel contributions or are unaltered from the previous method. For example, in Section 3.4, the sentence “We use the pose estimation strategy introduced in [25]” may lead the reader to think that the authors are experimenting with a new pose estimation approach from the literature, while it was already used in the DRGN-AI (or at least I believe so). 
* While results are clear and easy to follow, they only present standard deviations in Table 1 and not in the other tables. Moreover, it is unclear why the CryoSPARC result in Table 1 does not report any standard deviation.
* Will the authors provide a public code repository enabling others to replicate the presented results and use the method in the context of their data and experimentation?
questions: questions: - Do you think that the choice of using only 2 dimensions to represent the conformational space in Hydra may have affected the performance of your method? I think that the choice is excellent for increasing the interpretability of results since it eliminates the need to use techniques like PCA to visually understand the learned latent space. My intuition is that thanks to the K neural fields, the conformal representation may be better partitioned in lower dimensional spaces, but I wonder if you also have considered a higher number of dimensions.

- In Table 2, both Img-FSC and Pose errors are reported. It is unclear to me if the pose error and the Image FSC may be somehow correlated, as, intuitively, a wrong estimated pose could increase the reconstruction error. If that was the case, then I would be surprised to see such a small gap in Img-FSC with respect to CryoDRGN2 compared to the improvement in pose error.

- In Table 2, Hydra shows the lowest Pose error, while DRGN-AI shows the highest despite using a similar pose estimation approach. Can this performance difference be ascribed to the mixture model, or there are other factors to take into account, e.g. the dimensionality of the conformational space?
limitations: limitations: - The main limitation of this work is the need to know the number of classes K in advance. This can be addressed by an exhaustive search of the optimal value of K, however, this further worsens the second main limitation that is related to the computational cost of the method. The authors address both limitations and propose a possible future direction to reduce computational cost.

- Another point of view of the previous limitation is related to scalability. Given a fixed computational budget, since a new neural field is required for each additional protein class, the possible choice of the number of desired classes K is limited, potentially reducing the applicability of the method in datasets that contain a higher number of different macromolecules.
paper_summary: paper_summary: This work describes a new method for ab initio heterogeneous reconstruction in cryo-EM using mixtures of neural fields. This generalizes previous approaches, such as CryoDRGN and DRGN-AI, which attempts to reconstruct 3D molecular densities using a single neural field representation. The resulting method is able to handle both compositional (discrete) and conformational (continuous) heterogeneity, with each mixture component handling the continuous variability of each distinct compositional state. The performance of the method is evaluated on two synthetic and one experimental dataset.
strength_weakness: strength_weakness: The most important issue is the lack of experimental validation for the combined estimation of compositional and conformational heterogeneity. While this is tested in the third experiment (Section 4.3), this is only on a synthetic dataset. As the authors are no doubt aware, however, the behavior of a reconstruction algorithm can be quite different when applied to real data. It is therefore encouraging that the authors present results on an experimental dataset (Section 4.2), but this only covers compositional heterogeneity (and not conformational). That being said, validating the full method on an experimental dataset would make a stronger case for the proposed work.
questions: questions: – On line 49, please explain the principal difference between DRGN-AI and its predecessors. This is particularly important since the proposed method is quite closely related.
– What is meant by “tackle the discrete heterogeneity and the continuous variability *in sequential order*”? Is this referring to clustering the data, estimating the pose, and then estimating the continuous variability for each cluster? This could perhaps be clarified.
– Another approach to manifold embedding for continuous heterogeneity (line 82) is discussed in Moscovich et al., 2020.
– Line 130 should have “used” instead of “use”.
– The last line in the caption for Figure 1 is missing “Section” in front of “3.3”.
– If the dimension d is the same for each component in the mixture (eq. 3), how does this work when the number of degrees of freedom varies with each discrete structure?
– In eq. 4, σ_2 should be σ^2.
– Please discuss the choice of using point estimates for the pose and latent variables in eq. 6. Why is this approximation reasonable? This can be especially sensitive in the ab initio case where the can be great uncertainty in the pose estimation stage. Why do we not need to account for this here?
– Why use the area under FSC to evaluate the accuracy of the reconstruction (Table 1)? What are the advantages compared to gold-standard resolution estimates?
– Does the “strategy” on line 216 refer to the pose (eq. 9) estimation or to the overall estimation of pose and other variables? In other words, do we run ~1e6 iterations of pose estimation of ~1e6 iterations of pose estimation alternated with SGD?
– In line 227, what is meant by “a single pass”? The described algorithm seems to loop over the data several times over. Single-pass approaches usually refers to moment-based methods which pass over the data once to calculate certain statistics and then use those statistics to reconstruct the molecule.
– What is the SNR of the images described in Section 4.3? This information is important enough that it should be supplied in the main text.
– Please provide running times for the various experiments. The metric of 4 GPU-days is cited in the Discussion, but it is not clear how this relates to the number of images (is this for the 60 000-image dataset or the 3 000-image dataset?).
limitations: limitations: As stated above, the main limitation of the work is its lack of validation on experimental data (for both compositional and conformational variability). It is also not clear how computationally intensive the implementation is and how this can be mitigated.
paper_summary: paper_summary: In this paper, author(s) porpose Hydra, an *ab initio* approach to model conformational and compositional heterogeneity. They ahieve this by parameterizing structures of proteins as a mixture of *K* neural fields.
strength_weakness: strength_weakness: Novelty:

- The hierarchical pose search (HPS) method for pose estimation is proposed in DRGN-AI [1].
- Using an ensemble of representations to model heterogeneity in protein cryo-EM reconstruction has been adopted in many previous works [2][3].

[1] Levy, Axel, et al. "Revealing biomolecular structure and motion with neural ab initio cryo-EM reconstruction." *bioRxiv* (2024): 2024-05.  
[2] Punjani, Ali, and David J. Fleet. "3D variability analysis: Resolving continuous flexibility and discrete heterogeneity from single particle cryo-EM." *Journal of structural biology* 213.2 (2021): 107702.  
[3] Kimanius, Dari, Kiarash Jamali, and Sjors Scheres. "Sparse Fourier backpropagation in cryo-EM reconstruction." *Advances in Neural Information Processing Systems* 35 (2022): 12395-12408.
questions: questions: How would Hydra perform on EMPIAR-10076, one widely used dataset with complex compositional heterogeneity?
limitations: limitations: Authors discussed the limitation in their paper.
paper_summary: paper_summary: This paper introduces a novel method, Hydra, for ab initio heterogeneous cryo-EM reconstruction. Different from existing approaches, Hydra separately models conformational and compositional heterogeneity by integrating K-parameterized neural fields to represent cryo-EM density maps. Furthermore, Hydra employs a hybrid optimization strategy to optimize particle poses, heterogeneity, and density map representations concurrently. The authors assess the efficacy of Hydra on three datasets comprising various protein complexes (two synthetic, one experimental). Extensive experimental results indicate that Hydra outperforms three baseline methods regarding reconstruction quality, particle classification, and pose estimation accuracy.
strength_weakness: strength_weakness: **Synthetic Dataset**

The number of images in the tomotwin3 synthetic dataset is too small to match the real cryo-EM setting; there are only 3000 particle images with a very low SNR of 0.01 (tomotwin3). In a real scenario, there would be 50,000 to more than 100,000 particles with conformational variability. Without reporting the 3D resolutions of the results, it is very hard to quantitatively evaluate the reconstruction results. Additionally, this synthetic dataset lacks conformational variability, which hinders the evaluation of conformational variability recovery. The combination ratio of the three types of particles is not explored either; it remains unclear if Hydra would be sensitive to a class with only a small ratio of particle numbers. In Section 4.3, the dataset settings that include pre-catalytic spliceosome, 80S ribosome, and SARS-CoV-2 spike protein are unrealistic. To sum up, I recognize a significant gap between this synthetic dataset and real datasets, and I feel the experiment is not sufficient to evaluate Hydra adequately.

**Reconstruction Quality Evaluation**

For qualitative evaluation, the differences between the various states in Figure 4 are very subtle, making it difficult to judge whether the surface changes are due to different conformations or the result of applying different thresholds to the density map. Also, I argue that cryoDRGN-AI and cryoDRGN2 also account for conformational heterogeneity without explicitly classifying particles; their qualitative results should also be compared in Figure 4.

For quantitative evaluation in Table 1 and Table 2, the authors only use Img-FSC to compare reconstruction quality. To the best of my knowledge, prior work such as cryoSPARC and cryoDRGN reports widely used 3D resolution calculated by thresholding the FSC curve between two half maps of the reconstruction for experimental datasets or between the reconstructed 3D density map and the synthetic ground truth density map (available for synthetic datasets).

**Choice of Metrics for Pose Evaluation**

The authors use the median Frobenius norm as the metric to compare pose error, which ignores the translation part. Referring to DRGN-AI, the in-plane and out-of-plane angle error, along with the translation error, should be reported. Additionally, showing the angle distribution of poses for comparison could be beneficial.

**Ablation Study**

This paper misses the ablation study of the number of K. I would like to know how to determine K in real cases. In this paper, the authors run CryoSPARC multiple times to determine the best K for H.

**Miscellaneous**

1.	The chirality of the reconstruction results for the pre-catalytic spliceosome in Figure 4 appears to be incorrect.
2.	The construction of the ribosplike dataset involves a mixture of three different proteins. In real single-particle cryo-EM experiments, this scenario seems rare as the purified samples are carefully prepared and should not contain or only contain a very small ratio of undesired particles that can be easily filtered out in 2D classification. Maybe in cryo-ET, Hydra can perform one-for-all sub-tomo averaging?
3.	What’s the protocol to run CryoSPARC on a synthetic dataset? If it performs so well in Table 1, what is the advantage of Hydra?
4.	This paper seems rushed. In Lines 57 – 62, the first three contributions have basically the same meaning; please consider rephrasing them.
questions: questions: My major concerns have been listed in the weakness section in this review, along with some of my confusion and questions. I list the important ones:
- Why is the synthetic dataset limited to 3000  particle images in a very low SNR of 0.01 for the tomotwin3 dataset? In my experience, this limited number usually causes poor reconstruction results.
- Can you provide the 3D resolutions of the reconstruction results to facilitate a more quantitative assessment?
- Have you tested how Hydra handles classes with a small ratio of particle numbers, given that the combination ratio of the three types of particles is not explored?
- The differences between the various states in Figure 4 are very subtle. How do you ensure that the observed surface changes are due to different conformations and not due to varying density map thresholds?
- Could you include qualitative comparisons with cryoDRGN-AI and cryoDRGN2, as these methods also address conformational heterogeneity without explicit particle classification?
- How do you sample the different conformation states in the conformational latent space? (like the white circular points in the latent space in Figure 4.)
-The median Frobenius norm metric ignores translation errors. Could you also report the in-plane and out-of-plane angle errors, along with the translation error, similar to DRGN-AI?
- What is the protocol for running CryoSPARC on a synthetic dataset, and what advantages does Hydra offer if CryoSPARC performs well in Table 1?
- How is the runtime efficiency of the Hydra algorithm compared to other algorithms?
limitations: limitations: The authors have addressed the limitations of Hydra in the main paper.

Comments: title: Clarification to response -- demonstration on real data
comment: Thanks for the response. I will consider it for my final decision.
Thank you for your thorough rebuttal and for your clarifications. I will keep my score at 7 as I believe this is a good paper that deserves to be published as part of the proceedings.
Thank you for your suggestion regarding the order of Section 4.2 and 4.3. We will swap these sections for clarity. On the experimental dataset, we have provided a comparison to cryoSPARC and DRGN-AI in Figure S5 and S6, and we agree with the reviewer that this comparison could be emphasized. We will move parts of Fig. S5 and S6 to the main paper and can further add a discussion on the level of expert processing required in the related works and the discussion of the results. Thank you for the feedback!
Apologies we mis-stated that 3DVA does not handle compositional heterogeneity in our last response – we meant to clarify that 3DVA is not designed to handle a mixture of different species, a strong form of compositional heterogeneity. While 3DVA can technically handle compositional heterogeneity (as any density-based method does not explicitly constrain the reconstructed density maps), 3DVA uses a linear subspace of the volumes for its model of heterogeneity, and thus has an inductive bias towards modeling continuous motions of a single complex (similar to cryoDRGN and DRGN-AI). We would like to emphasize that, in Fig. 9 of 3DVA [1], the reconstructed complex is an assembling 50S ribosome where the discrete “classes” share most of their density. Finally, but perhaps most relevant, 3DVA is applied in the fixed pose setting, typically at the end of the processing pipeline after several steps of 2D or 3D classification, where poses are typically aligned to a single static reference structure, which would not apply in the case of a mixture of proteins.

In contrast, Hydra is the first method to adopt an ensemble of independent neural representations to model heterogeneity. This approach is similar to the ubiquitous 3D classification (a discrete mixture model of voxel arrays), but is far more expressive due to its neural representation augmented with latent scores. Thus, we demonstrate Hydra on datasets containing a mixture of different species that have conformational motions, where _ab initio_ pose estimation is necessary.  


[1] Punjani, A., & Fleet, D. J. (2021). 3D variability analysis: Resolving continuous flexibility and discrete heterogeneity from single particle cryo-EM. Journal of structural biology, 213(2), 107702.
Thank you for bringing up this discussion point. Our main motivation for Hydra is to propose a method designed to tackle a more challenging class of datasets for cryo-EM – those containing a mixture of distinct, dynamic protein complexes. We note that this is a new experimental setting for cryo-EM, thus we hope that Hydra will inspire new experimental protocols as well as other future reconstruction methods designed for this form of extreme heterogeneity.  

Hydra performs _ab initio_ reconstruction, i.e. joint inference of poses and structure, to address this form of extreme heterogeneity. This contrasts with traditional tools that require input poses (e.g. cryoSPARC 3DVA, cryoSPARC 3DFlex, and most other heterogeneity analysis tools). Since poses are typically obtained from an upstream consensus reconstruction, this assumes that all images can be aligned to a static reference structure, which does not hold when analyzing a mixture of distinct complexes.

Hydra is the first method to adopt an ensemble of independent neural representations to model heterogeneity. This approach is inspired by the ubiquitous 3D classification (a discrete mixture model of voxel arrays), but is far more expressive due to its neural representation augmented with latent scores. In Fig. 4, for example, we show that Hydra reveals the conformational heterogeneity of the ribosome, the spliceosome and the spike protein, while 3D classification in cryoSPARC can only reconstruct three static states (Fig. S7.c).
Thanks for the author's rebuttal, I have read it carefully as well as other reviewers' opinions.

I noticed that the proposed method has been compared with cryoSPARC. What I want to hear is some insights about the advantages of the proposed method compared with these traditional methods, and why structural biologists should use the proposed method instead of sticking to traditional methods?

I hope to hear the author's answer during the discussion stage to decide my final rating.
I sincerely appreciate authors' clarification. But based on my understanding, 3DVA can handle compositional heterogeneity. In the fig. 9 of the bioRxiv version of 3DVA paper, they showed results on EMPIAR-10076, a dataset with complex compositional heterogeneity.
Thank you for the clarification! It might be easier to follow if the order of the results in Section 4.2 and Section 4.3 is reversed. 

As the authors mentioned, in an experimental dataset, methods like DRGN-AI and cryoSPARC struggle to reveal the subtleties of the total distribution of compositional heterogeneity without expert data processing. This is the most convincing case, in my opinion, and should be further emphasized in the paper, with additional evaluations based on it. I recommend that Hydra be carefully compared with the most competitive baselines, cryoSPARC and DRGN-AI, in terms of performance, training and inference speed, and the level of expert effort required. 

While my concerns about the real dataset have been somewhat addressed, the paper still has some previously mentioned weaknesses that I cannot ignore.I will consider improving my rating before the discussion deadline, taking into account all the comments from the other reviewers.
Thank you for reading our response. We realize there was a misunderstanding in your comment, and we apologize for the lack of clarity.

In our work, we consider two synthetic datasets for validation of our method: tomotwin and ribosplike. The first contains a mixture of 3 complexes that we use for evaluation of hyperparameters (which we updated during rebuttals to contain a realistic number of images, and with resolution metrics); the latter contains a mixture of 3 complexes with conformational motions. However, **we do additionally showcase the method on a real dataset**, collected from a lysate sample where we are able recover the structures of the membrane protein ryanodine receptor (RyR), the mitochondrial respiratory chain complex III (CIII), and the dimeric complex of the soluble valosin-containing protein (p97). We don't emphasize it since there is no ground truth to validate the conformations, however we feel that this dataset demonstrates that the method will be transferrable to real application settings.

To answer your comment on the technical contributions, we would like to emphasize that Hydra is the first method to use a mixture of neural fields for cryo-EM reconstruction. We show that it can handle a novel and complex problem setting in cryo-EM reconstruction: the simultaneous estimation of conformational and strong compositional heterogeneity, in an _ab initio_ setting.

We also want to clarify that the new resolution estimates for the volumes are close to the Nyquist resolution for this dataset (9.2 A vs 9 A). We emphasize that these experiments were performed with 128x128 images for computational considerations for hyperparameter evaluation.

Finally, thank you for the suggestions to extend this work to cryo-ET sub-tomogram averaging (STA). We completely agree this would be an amazing showcase and hope that this work can inspire research in this direction, but we feel that adapting to a new imaging modality and data type is beyond the scope of this work.
I sincerely appreciate the authors' huge efforts in responding to the review. I have also carefully considered the perspectives provided by other reviewers. While some of my concerns, such as the choice of evaluation metrics and the size of the synthetic dataset, have been addressed, I still find the overall quality of the paper to fall below the standard expected for NeurIPS. The main technical contributions remain unclear, and the experimental setup is lacking, particularly due to the absence of a real dataset and the low-quality reconstruction results, which are insufficient for 3D model building. Consequently, I am unable to raise my rating at this time. I recommend that the authors extend this work to cryo-ET sub-tomogram averaging, where the motivation is stronger and more substantial real-world experiments can be performed.