CryoSPIN: Improving Ab-Initio Cryo-EM Reconstruction with Semi-Amortized Pose Inference

Summary

This paper introduces an approach to ab-initio homogeneous reconstruction that handles multi-modal pose distributions with a tailored encoder and accelerates pose optimization with semi-amortization. The approach uses a shared CNN feature extractor with multiple pose predictor heads, predicting several plausible poses for each image to account for pose uncertainty. Unlike the computationally expensive implicit networks used by cryoDRGN and cryoAI, this method employs an explicit 3D decoder, speeding up reconstruction. The encoder-decoder architecture is trained with a "winner-takes-all" loss, where the 3D decoder generates multiple 3D-to-2D projections, and the one with the lowest reconstruction error determines the loss. This method achieves faster convergence and more accurate poses compared to relying solely on the encoder's predictions. Evaluations show that the approach outperforms cryoAI and is competitive with cryoSPARC on both synthetic and experimental datasets.

Strengths

- This paper proposes a two-stage ab-initio homogeneous reconstruction algorithm based on a novel multi-head transformer architecture to output multiple proposals of predicted poses. As a result, the performance on both synthetic and real datasets outperforms learning-based baselines. - A novel winner-takes-all loss has been introduced to specialize multi-heads of the decoder. Each head of the decoder can account for the pose estimation in a local region, mitigating the burden on the decoder. - This paper is well-structured and easy to follow, with appropriate references.

Weaknesses

Despite its strengths, I see some major weaknesses in this paper: 1. It is unfair to claim "our semi-amortized method is faster than the amortized method of cryoAI" (Line 79) as there are some minor modifications, like changing the decoder from an implicit representation to an explicit volume to reduce the computational cost. An explicit density volume or a feature volume can easily replace all coordinate-based neural networks. In this paper, there is no ablation study about the representation of the decoder or report on how much training and inference time has been reduced by this replacement. Explicit volume with multi-head predictions leads to higher memory costs during training. Advanced representations that can achieve a better trade-off between computational and memory costs, like TensoRF, triplanes, and hash encodings, are not well discussed or explored. 2. The improvement of results is incremental. In Figure 2, the performance improvements on synthetic datasets are very subtle. On the experimental results, it seems cryoSPARC, which uses a traditional ab-initio method, achieves the best performance. Additionally, why do the resolutions of EMPIAR-10028's FSC curve all meet the resolution limit? 3. Translation estimation is important for a complete ab-initio model. Why does this method not support translation estimation while cryoAI is capable?

Questions

Overall, I am not convinced that the proposed semi-amortized approach can fundamentally improve pose prediction results without any significant improvements on the image encoder side. All experiments conducted in this paper show only minor performance improvements or even no improvements compared to cryoSPARC, which has no deep learning but only traditional optimization, does it imply that the actual bottleneck is the weak image encoder? The image encoder (VGG in this paper) is not pre-trained on a diverse cryo-EM dataset but trained from scratch on a per-scene dataset in a self-supervised manner. How can it be ensured that the extracted features can be applied for accurate pose prediction? I believe the image encoder is also a performance bottleneck. A PCA visualization of the feature space of the encoder (or calculating the cosine similarity) would be very helpful to validate if images with similar poses can be clustered and images with different poses can be far from each other. It is important for this work to conduct a performance analysis for both the encoder and decoder to identify the actual problem. Unfortunately, I see no analysis in the current version of this paper. I am happy to improve my ratings if my main concerns (including weaknesses) can be addressed by the authors.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes, there is a discussion of the limitations in this paper, but I believe the limitation of the image encoder used in this work should be also discussed.

Summary

This work introduces a new method for 3D ab initio estimation in cryo-EM using amortized inference. It relies on an existing technique which predicts the 3D rotation of a cryo-EM image using a convolutional neural network. That estimated rotation is then used together with a neural representation of the 3D volume to regenerate the projected image. The network weights (both of the pose estimation network and the 3D neural representation) are then optimized to reduce this auto-encoding loss. The first important difference between the proposed framework and previous work is that this method produces several estimates using multiple heads and computes the loss with respect to the rotation estimates that yields the closest matching image. The second important difference is that after a few epochs of training, the pose estimation network is discarded, and the best pose estimates are refined using alternating minimization together with the 3D neural representation. The resulting algorithm is shown to perform well on both simulated and experimental datasets.

Strengths

The writing of the manuscript is clear, in particular with respect to the introduction of the problem and its context as well as the presentation and interpretation of numerical results. The method seems to be well thought through and successfully mitigates the important pitfalls of the 3D ab initio reconstruction problem.

Weaknesses

The exposition of the pose inference network and the 3D neural representation are quite superficial and deserve a more detailed explanation. For example, how are the 3D rotations obtained from the convolutional backbone output? Are they encoded in axis–angle parametrization? If so, how are the outputs properly constrained? Similarly, the 3D neural representation is described as leveraging the Hartley transform and somehow uses a decomposition of this transform into its mantissa and exponent. How is this relevant to the actual neural network architecture? The work would also benefit from more extensive testing on experimental datasets beside EMPIAR-10028. This is a relatively high-SNR dataset depicting a very large molecule (an 80S ribosome) and is not representative of experimental datasets in general.

Questions

– On like 54, what is meant by “we adopt an explicit paremeterization to further accelerate the reconstruction”? How does the parametrization affect convergence here? – The idea of “semi-amortized inference” is never actually explicitly described. Presumably this means switching from amortized estimation to pose refinement at some point. – How is L_{i,j} defined in eq. (6)? – Why can the proposed method not handle shifts? Should this not be just a matter of adding another output to the pose estimation heads? – For the synthetic and real datasets, the method switches from amortized inference to pose refinement after 7 and 15 epochs, respectively. How was this switching point determined and is there some way to automate this? – The labels (Left) and (Right) appear to be switched in the caption for Figure 4. – On line 227, “experiment” should be “experiments”.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have largely addressed the limitations of the work in the manuscript.

Summary

The submission addresses ab-initio cryo-EM reconstruction, where both image poses and the 3D structure are estimated. The authors adopt a multihead architecture to estimate multiple poses for each image, encouraging the exploration of pose space early in the reconstruction process. They then refine poses in an auto-decoding fashion using SGD. Experiments on synthetic and experimental datasets demonstrate acceleration and resolution improvement over baseline approaches.

Strengths

1. Compared to cryoAI, mapping the input image to multiple pose candidates in the auto-encoding stage can account for pose uncertainty and encourage exploration of the pose space during the initial stages of reconstruction. 2. The experimental results provided in the paper include both synthetic and real experimental data, as well as video results. 3. The proposed method offers a speed advantage over cryoAI.

Weaknesses

1. The paper's technical contribution is primarily limited to a multi-head architecture for estimating multiple plausible poses. 2. The results presented for cryoSPARC are significantly worse than my experience suggests (i.e., on Splicesome or EMPIAR-10028). CryoSPARC's results should not be inferior to those obtained by the proposed approach. 3. The paper cannot handle heterogeneous reconstruction, which limits its general applicability in structural biology research. 4. Import baselines are missing: RELION, cryoFIRE, cryoDRGN-BNB and cryoDRGN2.

Questions

1. The results presented for cryoSPARC are significantly worse than my experience. What is the setting for cryoSPARC? Are there any important assumptions? 2. Will the multi-head structure predict very similar poses, thereby undermining the assumption that this design encourages exploration of the pose space? 3. Can the authors provide a comparison of the training time and memory consumption between the proposed method and the baselines?

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

n/a