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

Cryo-Electron Microscopy (cryo-EM) is a popular experimental technique to recover the 3D structure of macromolecular complexes, such as proteins, using extremely noisy images that contain particles posed in unknown orientations. We propose a new semi-amortized approach to the ab-initio reconstruction problem. In early stages, when uncertainty is high, poses are estimated using auto-encoding, followed by auto-decoding as uncertainty decreases. A multi-head encoder is adopted for amortization to infer multiple plausible poses for each image, encouraging exploration of pose space, while flexible auto-decoding iteratively update poses per-image using stochastic gradient descent. Empirical results on synthetic datasets demonstrate that our method is able to handle multi-modal pose distributions, and the use of auto-decoding yields faster and more accurate pose convergence compared to baselines. Finally, we show that our approach converges faster than state-of-the-art cryoAI and achieves higher resolution on experimental data.