CryoSWD: Sliced Wasserstein Distance Minimization for 3D Reconstruction in Cryo-electron Microscopy

Abstract

Single particle reconstruction (SPR) in cryo-electron microscopy (cryo-EM) is a prominent imaging method that recovers the 3D shape of a biomolecule, given a large number of its noisy projections from random and unknown views. Recently, CryoGAN [1] cast SPR as an unsupervised distribution matching problem and solved it via a Wasserstein generative adversarial network (WGAN) framework. The approach bypasses the estimation of the projection parameters. The reconstruction criterion in CryoGAN is Wasserstein-1 distance. Despite the desirable properties of Wasserstein distances (WD) such as continuity and almost everywhere differentiability, they are difficult to compute and require careful tuning for a stable training. Sliced Wasserstein distance (SWD), on the other hand, has shown desirable training stability and ease to compute. Therefore, we propose to re-place Wasserstein-1 distance with SWD in the CryoGAN framework, hence the name CryoSWD. In low noise regimes, we show how CryoSWD eliminates the need to have a discriminator which is crucial in CryoGAN. However, coupling CryoSWD with a discriminator boosts its performance, especially in high noise settings. While performing as good as CryoGAN, CryoSWD does not require a gradient penalty term for stabilizing the training and imposing Lipschitz continuity of the discriminator.