CryoGEN: Generative Energy-based Models for Cryogenic Electron Tomography Reconstruction

We sincerely thank the reviewer for their positive evaluation and recognition of our work's contributions. We apologize for any misunderstandings and would like to take the opportunity of rebuttal to clarify the strengths of our work.

1. Contribution and Originality:

   • To the best of our knowledge, we are the first to propose a comprehensive solution to the missing wedge problem in CryoET using Bayes' rule. Unlike most CryoEM algorithms, which often rely on known prior distributions, our approach addresses this as an unsupervised learning problem. This introduces additional complexity to the mathematical formulation.

   • Our identification of IsoNet's limitations is unique, emerging from a non-convex optimization perspective. The primary contribution of our work lies in the mathematical formulation and the algorithmic framework. While we employ standard techniques from the field of computer vision, this should be seen as a strength rather than a limitation, as these techniques verify the generalizability of our approach.

2. Theoretical Analysis: Formally analyzing CryoET from a theoretical perspective presents inherent challenges. While a comprehensive theoretical analysis is beyond the scope of this work, we offer key insights: although GANs are susceptible to mode collapse, they are highly effective in addressing mode averaging issues, as discussed in the motivation section. This is a primary reason for their inclusion in our approach. In the final version, we will provide a more detailed explanation and explore potential theoretical analyses.

3. In-Depth Analysis: Section 3 provides an illustrative analysis of why CryoGEN outperforms IsoNet, with Figures 5 and 6 highlighting the core motivations. Additional experiments have been conducted and included in the revised version. Figure 16 demonstrates CryoGEN's enhanced denoising capabilities. Furthermore, the comparison indicates that IsoNet consistently suffers from residual effects of the missing wedge problem across all experiments, while CryoGEN effectively mitigates this issue.

4. Comparison to SOTA methods: We have compared our method with the state-of-the-art (SOTA) methods, IsoNet and DeepDeWedge, the latter being published shortly before our submission. To the best of our knowledge, these represent the current SOTA. However, we are open to conducting additional experiments if the reviewer can suggest other relevant methods.

5. Large Datasets and Deep Models: Compared to the SOTA methods, CryoGEN exhibits greater robustness when applied to large datasets and deep models. This observation stems from IsoNet’s need for meticulous fine-tuning at each iteration. A plausible explanation is that IsoNet’s recursive training approach may accumulate errors over iterations, leading to instability. In contrast, our algorithm avoids recursive updates, ensuring faster training and enhanced stability. Furthermore, our experiments spanned a range of difficulty levels, from straightforward cases to the most challenging scenario involving the HIV ribosome.

6. Significance of CryoGEN for CryoET: We utilize the isotropy property of CryoET, assuming that distributions are uniform across all orientations. To the best of our knowledge, this is a novel approach that has not been formally proposed in the CryoET field before. Furthermore, we are the first to formally define and address the missing-wedge problem, a longstanding challenge in CryoET, with a consistent solution.

We thank the reviewer for acknowledging our contributions and offering positive feedback. We will address the reviews raised by the reviewers accordingly.

1. Discussion: We have partially addressed the limitations and potential future directions in our previous responses, and we will integrate these discussions into the revised version.

2. Typos: The identified typos have been corrected. We appreciate your attention to detail.

3. Number of Crops: We employ the same method proposed in Isonet. For the ribosome, the number of crops varies based on the input size, as the crop size is fixed at 96 for each dimension. Detailed training procedures are provided in Appendix A.6.

4. SNR of the data shown in Figure 7: The SNR can be considered effectively infinite since no noise is added. When the SNR is 0.01 or 0.001, the figures are too corrupted to contain any useful information, making learning impossible for all methods. However, we have conducted additional experiments with SNR = 0.2, and the updated results are provided in Appendix A.3.3.

We thank the reviewer for the detailed comments and helpful feedback. We also appreciate the acknowledgment of our contribution. We will address each of your remarks individually and hope to clarify any concerns.

1. Contribution and Originality:

   • To the best of our knowledge, we are the first to provide a comprehensive solution for the missing wedge problem in CryoET with Bayes rule. It brought to our attention that our approach was developed concurrently with the work CryoFM. However, Unlike CryoFM, which relies on CryoEM data with complete information to construct $p_x(x)$, we build $p_x(x)$ solely from the distribution of incomplete data $p_y$. This introduces additional complexity into the mathematical formulation.

   • Specifically, we leverage the isotropy property of CryoET: after constructing $p_y$, we define $\log p_x(x) \propto \frac{1}{|\mathcal{R}|} \sum_{R \in \mathcal{R}} \log p_y(\mathcal{T}_M \circ R \circ x)$. This rotational operation assumes that the distributions are uniform across all directions. To our knowledge, this represents a novel contribution that has not been previously proposed.

   • Furthermore, inpainting algorithms typically rely on stronger assumptions compared to ours. For instance, [1] trains its model on a dataset of complete images, which aligns with our initial observation where $p_x(x)$ is assumed to be known. However, this assumption is generally not valid in CryoET. We will highlight this distinction in our revised paper.

2. Motivation:

   The motivation discussed in Section 3 may have been addressed in general generative models but likely from different perspectives and aimed at solving distinct challenges:

   • For example, while GANs are known to suffer from mode collapse issues, they are also effective in addressing problems related to mode averaging. This strength is one of the reasons we incorporate GANs into our approach, although it is not the only solution.

   • In contrast, VAEs are prone to averaging issues due to their smoother latent space representation and the regularization inherent in their formulation. Specifically, the KL divergence term in the VAE objective encourages the posterior $q(z|x)$ to stay close to the prior.

   • Finally, our identification of IsoNet's limitations is unique, emerging from a non-convex optimization perspective. Our motivation complements these existing approaches while addressing challenges specific to our context. The primary contribution of our work lies in the mathematical formulation and the algorithmic framework. While we employ standard techniques from the field of computer vision, this should be seen as a strength rather than a limitation, as these techniques verify the generalizability of our approach.

3. Equation 8 is an implementation of Equation 3

   Yes, you are correct. Specifically, the order of the likelihood and prior is reversed between Equation 8 and Equation 3, and we will fix this in the revision. Thanks for pointing out.

4. Both the incorporation of an energy model and noise perturbation to the input are essential components

   • Energy Model: The energy model serves as a fundamental component in the procedure for estimating the prior $p_x(x)$ from $p_y(y)$. This process involves the following steps:

     • Defining the Energy Function $E(y)$: The goal is to construct an energy function where regions close to data points in $\mathcal{Y} = {y_i}$ have low energy, and regions far from the data points have high energy. To achieve this, we leverage GANs (refer to [2] for details) to assign low $E$ values to high-density regions and high $E$ values to low-density regions.

     • Constructing $p_y$: Once the energy model $E$ is trained, we construct $p_y$ using an energy-based representation, where $p_y(y) \propto \exp(-E(y))$, consistent with the Boltzmann distribution.

     • Constructing $p_x$: Using the isotropy property of CryoET, we define $\log p_x(x) \propto \frac{1}{|\mathcal{R}|} \sum_{R \in \mathcal{R}} \log p_y(\mathcal{T}_M \circ R \circ x)$, where $R$ denotes the rotation operation and $\mathcal{T}_M$ represents the application of the missing-wedge effect.

     By integrating these steps, we obtain $p_x(x)$, demonstrating the critical role of the energy model in our approach.

   • Noise Perturbation: As highlighted in the motivation section, if the cardinality of $X|y = \arg\max p(x|y)$ is greater than one, defining a parameterized function $g_\theta$ that maps a single observation to all possible values of $x$ becomes infeasible, resulting in a one-to-many mapping. To address this, we add $y + \epsilon$, where $\epsilon$ has full support over the real domain, effectively providing infinite cardinality and transforming the problem into a many-to-fewer mapping. This approach not only resolves the mapping issue but also enhances training stability, as discussed in [3].

5. Computational Efficiency

   We apologize for any confusion regarding computational efficiency. To clarify, our approach is faster during training while maintaining the same computational complexity during inference:

   • Training: IsoNet employs a recursive training process, whereas our method does not. This design choice makes our algorithm not only more robust during training but also over $\mathbf{\times 10}$ faster to converge. Moreover, IsoNet requires precise fine-tuning at each iteration, and its recursive nature can result in error accumulation and instability. In contrast, our approach avoids recursive updates, enabling faster training and greater stability.

   • Inference: For the prediction model, we use an architecture with the same computational complexity as IsoNet, ensuring that the time complexity during inference remains similar.

6. DeepDeWedge Results

   Apologies for the omission of results. DeepDeWedge was published shortly before our submission, but we have included XYZ slice images generated by DeepDeWedge in the appendix. We present the quantitative results below and they have been incorporated into the revised paper. Moreover, DeepDeWedge does not alter the formulation and uses IsoNet as a backbone, thereby inheriting all of its drawbacks.

   ---

   Table: Quantitative evaluation of image quality for tomography reconstructions using different methods, comparing PSNR and SSIM metrics (higher values indicate better performance for both metrics) on sphere, prism, and Vipp1 assembly datasets.

   Data State	Sphere PSNR	Sphere SSIM	Prism PSNR	Prism SSIM	Vipp1 Assembly PSNR	Vipp1 Assembly SSIM
   Corrupted	21.12	0.8113	14.82	0.6931	26.68	0.8000
   Iso-corrected	22.98	0.8770	19.11	0.8857	27.12	0.8191
   Dewedge-corrected	23.17	0.8824	21.10	0.9278	28.75	0.8758
   CryoGEN-corrected	29.19	0.9706	32.69	0.9949	30.65	0.9199

[1] Lugmayr, Andreas et al. “RePaint: Inpainting using Denoising Diffusion Probabilistic Models.” 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (2022): 11451-11461.

[2] Murphy, Kevin Patrick. Probabilistic Machine Learning: Advanced Topics. MIT Press, 2023. Chapters 24 and 26.

[3] Arjovsky, Martin, and Léon Bottou. "Towards Principled Methods for Training Generative Adversarial Networks." International Conference on Learning Representations, 2017, https://openreview.net/forum?id=Hk4_qw5xe.

Thank you for the clarification. The authors' responses have addressed most of my concerns. So I changed my score from 5 to 6. While the motivation presented in Section 3 still appears very general to me, I am largely convinced of the technical novelty of the proposed method. One remaining confusion is that, why is "leveraging the isotropy property of CryoET" considered a specific technical contribution? Also I suggest that the authors include a discussion of image inpainting approaches in the related work section, as these share similar problem settings and techniques with cryo-ET reconstruction.

3. Constructing The Distribution $p_x$

   One of the key challenges—and a main contribution of our work—is constructing the distribution $p_x$ from a distribution $p_y$. We outline this process step-by-step:

   • 1. Defining the Energy Function $E(y)$: Our objective is to define an energy function such that points near any data point in $\mathcal{Y} = \{y_i\}$ have low energy, while points distant from all data points have high energy. For this, contrastive learning presents itself as an effective choice. This method assigns lower $E$ values to regions with high data density and higher $E$ values elsewhere, which naturally leads us to use GANs as a tool (refer to [2] for more details).
   • 2. Constructing $p_y$: The next question is how to construct $p_y$ given a trained energy model $E$. We address this by representing the distribution using energy-based models as $p_y(y) \propto \exp(-E(y))$, in line with a Boltzmann distribution.
   • 3. Constructing $p_x$: Leveraging the isotropy property of CryoET, we define $\log p_x(x) \propto \frac{1}{|\mathcal{R}|} \sum_{R \in \mathcal{R}} \log p_y(\mathcal{T}_M \circ R \circ x)$, where $R$ denotes the rotation operation.

   By integrating all these steps, we arrive at the final algorithm.

4. Streak-like Artifacts

   For the experimental results presented in Figure 9, there is indeed no ground truth for comparison. However, we argue that the original results shown in IsoNet exhibit strong streaking artifacts. Streaking in imaging and signal processing refers to unwanted lines or artifacts in images, often caused by incomplete data. In CryoET, these artifacts arise due to the missing wedge problem, where limited angles are used for capturing images, resulting in elongated, streak-like artifacts in the reconstruction. In contrast, our method does not exhibit these streaking artifacts at all compared to IsoNet.

5. IsoNet’s Instability

   We apologize for not explicitly addressing IsoNet’s instability. This is a general observation we have made, as IsoNet requires careful fine-tuning at each iteration. An intuitive explanation is that IsoNet’s recursive training approach can lead to error accumulation, which in turn causes instability. Specifically, in IsoNet, training $g_{\theta_t}$ depends on the predictions generated by $g_{\theta_{t-1}}$ (see our previous response - Intuition Behind Consistency Loss). As a result, any intermediate failure in model $g_\theta$ can propagate through subsequent cycles, ultimately compromising the entire training process. This is similar to an open-loop control system, where errors accumulate due to the lack of feedback. In contrast, our algorithm does not rely on recursive updates, resulting in faster training and greater stability.

6. Addressing Mode Collapse in GANs

   The original GANs are known to encounter both mode collapse and training instability issues, but we address these as follows:

   • Mode Collapse: Mode collapse is less problematic for reconstruction in CryoET, as our goal is to find a single solution or a set of solutions $\tilde{x}$ that maximize $p(x|y)$ given the observation. While mode collapse may reduce diversity, it does not compromise quality.

   • Training Instability: By introducing noise $\epsilon$, our method aligns with the theoretical framework in [3] (the foundational paper on Wasserstein GANs), effectively addressing our specific one-to-many mapping problem illustrated in motivation section while significantly improving training stability.

7. Incorporation of Rotations

   The incorporation of rotations is used to derive $p_x$ from $p_y$ by leveraging the isotropy property of CryoET.

8. Experimental Results

   We evaluated our model and IsoNet on noisy synthetic data of Vipp1 assembly, showing that our method achieves better reconstruction of the missing wedge with reduced noise, as illustrated in Figure 16. Additionally, Figure 20 presents the visualization of a single ribosome. We also demonstrate the enhanced denoising capabilities of CryoGEN, while the comparison reveals that IsoNet consistently exhibits residual effects of the missing wedge problem across all experiments.

References:

[1] Chen, Wenlin, et al. "Diffusive Gibbs Sampling." International Conference on Machine Learning, 2024.

[2] Murphy, Kevin P. "Probabilistic machine learning: Advanced topics - Chapters 24 and 26." MIT press, 2023.

[3] Arjovsky, Martin, and Léon Bottou. "Towards principled methods for training generative adversarial networks." International Conference on Learning Representations, 2017.

We thank the reviewer for the detailed comments and instructive feedback. We also appreciate the recognition of our contribution. We will address your remarks one by one and hope to clarify your concerns. We apologize for any lack of clarity in our current version.

1. Motivation Clarification:

   • Gaussian Mixture Prior

     • The key observation here is that for each observation $y$ (incomplete data with missing wedge), there may be multiple solutions that align the observed outcome. By "align", we mean that for a given $y$, $X|y = \arg\max p(x|y)$ can represent a set of solutions rather than a single one. In such cases, the log posterior distribution $\log p(x|y)$ can exhibit multiple maxima—a fairly common scenario.

     • Our goal is to determine a parameterized function $g_\theta$ (usually a neural network) that maximizes $\log p(g_\theta(y)|y)$, which is the primary objective of this paper. In contrast, IsoNet’s objective is to directly minimize $\sum_{x_i \in X|y} ||g_\theta(y) - x_i||^2_2$, where each $x_i$ is a maximum of $\log p(x|y)$.

     • The distribution $p(x|y)$ is not necessarily a Gaussian mixture; rather, it is generally assumed to be non-convex. Since the exact form of $\log p(x|y)$ is unknown, we adopt the Gaussian mixture prior as a general assumption—an approach that has proven effective in current generative tasks within deep learning (For example, diffusion models approximate the data distribution as a Gaussian mixture by convolving the discrete data with Gaussian distributions at different levels of variance, which has demonstrated strong performance).

   • Lack of Ground Truth for $\log p_x(x)$

     In CryoET reconstruction tasks, there is no ground truth for $\log p_x(x)$, so we approach this as an unsupervised learning task, unlike the concurrent work, where $p(x)$ is assumed to be known. However, we propose a way to express $p_x(x)$ by first approximating the distribution $p_y$ based on observations of the incomplete data $\{y_i\}$ and then defining $\log p_x(x) \propto \frac{1}{|\mathcal{R}|} \sum_{R \in \mathcal{R}} \log p_y(\mathcal{T}_M \circ R \circ x)$. This formulation leverages the isotropy and rotational symmetry properties of CryoET, representing a key contribution of our work.

   • Addressing One-to-Many Mapping

     As mentioned earlier, if the cardinality of $X|y = \arg\min p(x|y)$ is greater than one, it is impossible to find a parameterized function $g_\theta$ that can map a single observation to all possible values of $x$, as this creates a one-to-many mapping. To address this, we use $y + \epsilon$, where $\epsilon$ has full support over the real domain, giving it infinite cardinality and effectively converting the problem into a many-to-fewer mapping.

2. Intuition Behind Consistency Loss

   This is the original objective function used in IsoNet. Although we lack reference or ground truth for the missing wedge, we do have it for other parts, enabling us to proceed with a pseudo-supervised approach. The steps of IsoNet can be outlined as follows:

   • Step 1: First, we fill in the missing wedge for the incomplete data using the prediction model $g_\theta$ to generate $\tilde{x}$.

   • Step 2: Next, we rotate the sample and apply a different missing wedge via $\mathcal{T}_M$ to produce $\tilde{y}$.

   • Step 3: Since we have the label for the artificially created missing wedge, we can train $g_\theta$ in a supervised manner by minimizing $||\tilde{x} - g_\theta(\tilde{y})||_2^2$.

   At the start of training, we incorporate this objective into our algorithm as well. The reasoning is that, when sampling from a distribution $p(x)$ using MCMC, a reasonable starting point for the Markov chain is $\int_x x p(x) \, dx$, even though this is not a minimum (see [1] for example). In our case, rather than sampling, we directly learn the mapping from $p_y$ to $p_x$ via $g_\theta$, but the intuition remains similar. Nonetheless, including this objective function still aids in improving the convergence of training $g_\theta$.

I thank the authors for their detailed answers to my questions. I would like to sincerely apologise for my late reply. Here are some comments on their answers:

• Point 4:

For the experimental results presented in Figure 9, there is indeed no ground truth for comparison. However, we argue that the original results shown in IsoNet exhibit strong streaking artifacts. Streaking in imaging and signal processing refers to unwanted lines or artifacts in images, often caused by incomplete data. In CryoET, these artifacts arise due to the missing wedge problem, where limited angles are used for capturing images, resulting in elongated, streak-like artifacts in the reconstruction. In contrast, our method does not exhibit these streaking artifacts at all compared to IsoNet.

I did not notice streaking artefacts in the IsoNet or DeepDeWedge reconstructions. I encourage the authors to provide pointers to streaking and/or ringing (as claimed in line 444) artefacts in the IsoNet and DeepDeWedge reconstructions shown in Figure 9 and Figure 19 (a).

In my opinion, the reconstruction obtained with the CryoGEN method does not look better than the IsoNet or DeepDeWedge tomograms. It just looks different. This is of course subjective because there is no ground truth.

• Point 5:

We apologize for not explicitly addressing IsoNet’s instability. This is a general observation we have made, as IsoNet requires careful fine-tuning at each iteration. An intuitive explanation is that IsoNet’s recursive training approach can lead to error accumulation, which in turn causes instability. Specifically, in IsoNet, training $g_{\theta_t}$ depends on the predictions generated by $g_{\theta_{t-1}}$ (see our previous response - Intuition Behind Consistency Loss). As a result, any intermediate failure in model $g_\theta$ can propagate through subsequent cycles, ultimately compromising the entire training process. This is similar to an open-loop control system, where errors accumulate due to the lack of feedback. In contrast, our algorithm does not rely on recursive updates, resulting in faster training and greater stability.

I appreciate the authors' explanation. However, as the paper contains neither this discussion, nor an experiment that demonstrates IsoNet's claimed instability, I feel that the statement that "IsoNet relies on recursively updating its predictions, which can result in train- ing instability and potential model collapse." should be removed from the abstract.

• Point 7:

The incorporation of rotations is used to derive p_x from p_y  by leveraging the isotropy property of CryoET.

I could not find any discussion of this "isotropy property" in the paper. As the authors rely on the "isotropy property" to motivate the random rotations during training, it should be discussed explicitly in the paper (e.g. in Section 4.1.1 and/or Section 4.2).