Deep Neural-network Prior for Orbit Recovery from Method of Moments

TL;DR

The paper addresses orbit-recovery problems under noisy, group-distorted observations, focusing on multireference alignment and single-particle cryo-EM. It introduces an amortized method-of-moments (MoM) framework that couples moment-based recovery with neural-network priors to simultaneously estimate the unknown signal and the distribution of group actions from the moments. The approach demonstrates that moment-inversion maps can be learned (with supervised pretraining) to accelerate reconstruction in MRA and to enable low-dimensional ab initio models for cryo-EM, including reconstructions from simulated and biological volumes. It highlights improvements in noise resilience and convergence speed, and outlines future work on higher-order moments, more general cryo-EM settings, and scalable GPU implementations.

Abstract

Orbit recovery problems are a class of problems that often arise in practice and various forms. In these problems, we aim to estimate an unknown function after being distorted by a group action and observed via a known operator. Typically, the observations are contaminated with a non-trivial level of noise. Two particular orbit recovery problems of interest in this paper are multireference alignment and single-particle cryo-EM modelling. In order to suppress the noise, we suggest using the method of moments approach for both problems while introducing deep neural network priors. In particular, our neural networks should output the signals and the distribution of group elements, with moments being the input. In the multireference alignment case, we demonstrate the advantage of using the NN to accelerate the convergence for the reconstruction of signals from the moments. Finally, we use our method to reconstruct simulated and biological volumes in the cryo-EM setting.

Figures (11)

• Figure 1: Overview of our MRA pipeline: The encoder $\xi_\theta$ takes moments $(\hat{M}_F^1,\hat{M}_F^2)$ as input, and outputs $z_\rho\in \mathbb{R}^n$, approximating a discretized probability density $\rho(X_1)$, and $z_v\in \mathbb{R}^n$ that approximates a discretized Fourier signal $\widehat{v}(K_1)$. Next, we use $z_\rho$ and $z_v$ to create $\left(M_F^1[z_v, z_\rho](K_1), M_F^2[z_v, z_\rho](K_1, K_1)\right)$ via equation \ref{['eqn:recon moments']}, which we then compare with the inputs to the encoder, i.e., $(\hat{M}^1,\hat{M}^2)$ via the loss function $\mathcal{L}_\textnormal{recon}$\ref{['eqn:mra_loss']}.
• Figure 2: Overview of our cryo-EM pipeline: The encoder $\xi_\theta$ takes moments $(\hat{M}_F^1,\hat{M}_F^2)$ as input, and outputs $z_\rho\in \mathbb{R}^{|Q|}$, approximating a discretized probability density $\left(\rho(R) \right)_{R \in Q}$ for some fixed set of quadrature points $Q\subset{\text{SO}\left(3\right)}$. Next, we create copies of the grid $K_2$ rotated corresponding to the elements of $Q$ and input them to our neural representation $\widehat{v}_\phi$, which outputs corresponding slices of a running estimate of $\widehat{v}$. These slices $\left\{S\circ Q(j)\circ \widehat{v}_\phi (K_2)\right\}_j$ along with $z_\rho$ are used to create $\left(M_F^1[\widehat{v}_\rho, z_\rho](K_2), M_F^2[\widehat{v}_\rho, z_\rho](K_2, K_2)\right)$ via equation \ref{['eqn:quadrature_moments']}, which we then compare with the inputs to the encoder, i.e., $(\hat{M}_F^1,\hat{M}_F^2)$ via the loss function $\mathcal{L}_\textnormal{recon}$ in \ref{['eqn:cryo_loss']}. Optionally, $\xi_\theta$ can also be used to output an extra $z_v$, a latent variable of $\widehat{v}$ that can be inputted into $\widehat{v}_\phi$.
• Figure 3: Predictions for the distribution $\rho$ (Left) and volume $v$ (Right), outputted by trained encoders $\xi_\theta^\rho$ and $\xi_\theta^v$ respectively, for $\rho,v$ being mixture of $2$ Gaussians. The solid lines are the ground truth $\rho$ and $v$, while the dotted lines are the corresponding predictions by an NN.
• Figure 4: Plots of logarithms (with base $10$) of Sum of relative errors (defined in \ref{['eq:relative error moments']}) for $\hat{M}_F^1$ and $\hat{M}_F^2$ across 3000 iterations (Top), and Reconstruction error (defined in \ref{['eq:relative error v mra']}) across 3000 iterations (Bottom); averaged over $20$ reconstructions of $(\rho(X_1),\widehat{v}(K_1))$ pairs drawn from the family of a mixture of $2$ Gaussians. In both plots, the blue curve corresponds to the scenario where the encoder underwent supervised training, while the orange corresponds to the scenario where it did not.
• Figure 5: (Left) 1000 points sampled from a mixture of eight von Mises-Fisher random variables shown in different colors, and (Right) $100$-point $13$-design plotted on a 3D unit sphere.