Wasserstein K-Means for Clustering
Tomographic Projections

Given a 3D molecule with an electrostatic potential function $\rho:\mathbb{R}^{3}\rightarrow\mathbb{R}$ (henceforth “particle”), the images captured by the electron microscope are modeled as noisy tomographic projections along different orientations. A tomographic projection of the particle onto the x-y plane is defined as $\left(\mathcal{T}\rho\right)(x,y)=\int_{\mathbb{R}}\rho(x,y,z)dz$. More general tomographic projections are given by a rotation $R$ of the potential function $\rho$ followed by projection onto the x-y plane,

where $R\in\text{SO}(3)$, and we define $(R\rho)(x,y,z)=\rho(R^{T}(x,y,z))$. Particle images in cryo-EM are typically modeled as follows,

where $h$ is a point spread function that is convolved with the projection and ${\boldsymbol{\epsilon}}$ is Gaussian noise. See Figure 1 for an example of a tomographic projection. We define the viewing angle of our particle as a unit vector ${\bf v}$ that represents orientation of our particle $R\in\mathrm{SO}(3)$ modulo in-plane rotations (which can be represented as members of $\mathrm{SO}(2)$).

The imaging process in cryo-EM involves high-levels of noise. To improve the signal-to-noise ratio (SNR), it is common to cluster the noisy projection images with ones that are similar up to an in-plane rotations. This clustering task is called “2D classification and averaging” in the cryo-EM literature and the clusters are called classes. The results of this stage have multiple uses, including:

• •

  Template selection from 2D class averages for particle picking (Frank and Wagenknecht, 1983; Chen and Grigorieff, 2007; Scheres, 2015).

• •

  Particle sorting (Zhou et al., 2020), discarding images that belong to bad clusters.

• •

  Visual assessment and symmetry detection.

• •

  Ab-initio modeling, where an initial 3D model (or set of models) is constructed to be refined in later stages (Greenberg and Shkolnisky, 2017; Punjani et al., 2017).

Existing methods for 2D clustering include the reference-free alignment algorithm that tries to find a global alignment of all images (Penczek et al., 1992), clustering based on invariant functions and multivariate statistical analysis (Schatz and Van Heel, 1990), rotationally invariant k-means (Penczek et al., 1996) and the Expectation-Maximization (EM) based approach of Scheres et al. (2005). All of these methods use the $L_{2}$ distance metric on rotated images.

Let us consider a simple centroid+noise model, where each image in a cluster is generated by adding Gaussian noise to a particular clean view ${\boldsymbol{\mu}}$ that is the (oracle) centroid,

In that case the $L_{2}$ metric is (up to constants) nothing but the log-likelihood of the image patch, conditioned on $\boldsymbol{\mu}$,

Under the centroid+noise model, the commonly-used $L_{2}$ distance metric seems perfectly suitable. However, real particle images have many other sources of variability, including angular differences, in-plane shifts and various forms of molecular heterogeneity. For these sources of variability, the $L_{2}$ distance metric is ill-suited. See Section 4 for more on this.

In this work we propose to use the Wasserstein-1 metric $W_{1}$, also known as the Earthmover’s distance, as an alternative to the $L_{2}$ distance for comparing particle images. In Section 2 we describe a variant of the rotationally invariant k-means that is based on a fast linear-time approximation of $W_{1}$ and in Section 3.2 we demonstrate superior performance on a synthetic dataset of Ribosome projections. In Section 4 we analyze the behavior of the $W_{1}$ and $L_{2}$ metrics theoretically with respect to angular differences of tomographic projections. In particular we show that the rotationally-invariant $W_{1}$ metric has the nice property that it is bounded from above by the angular difference of the projections. On the other hand, the $L_{2}$ metric shows no such relation.

Algorithm 1 describes a generic rotationally-invariant k-means algorithm based on an arbitrary image patch distance metric $d$. The choice $d=L_{2}$ gives the rotationally-invariant k-means of Penczek et al. (1996). By supplying $d=W_{1}$ we get a new rotationally-invariant k-means algorithm based on the $W_{1}$ distance. We choose $p=1$ for the $W_{p}$ distance as it admits a fast linear-time wavelet approximation as we will see in the next section. For both choices of the metric, we initialize the centroids using a rotationally-invariant k-means++ initialization which we describe in Algorithm 2. When we denote $r\in\mathrm{rot}$ we mean that $r$ performs an in-plane rotation of an image. We approximate the space of all in-plane rotation angles by a discrete set of angles.

Computing the distance $W_{1}(I_{1},I_{2})$ (also known as the Earthmover’s distance) between two $N\times N$ pixels can be formulated as a linear program in $O(N^{2})$ variables and $O(N)$ constraints. Given $n$ images, $k$ centers, and $m$ rotations for each image and $t$ iterations of k-means, we have to compute $O(n\times m\times k\times t)$ distances. Computing the $W_{1}$ distance between two $100\times 100$ size images using a standard LP solver takes on average of 5 seconds to compute using the Python Optimal Transport Library of Flamary and Courty (2017) on a single core of a 7th generation 2.8 GHz Intel Core i7 processor. Computing the exact $W_{1}$ distance over all the rotations of all the images over all the iterations is prohibitively slow.

Fortunately, the $W_{1}$ distance admits a wavelet-based approximation with a runtime that is linear in the number of pixels (Shirdhonkar and Jacobs, 2008). Let $\mathcal{W}I_{i}$ be the 2D wavelet transform of $I_{i}$. We can approximate the $W_{1}$ distance between two images $I_{1},I_{2}$ by a weighted $\ell_{1}$ distance between their wavelet coefficients,

where $\lambda$ goes over all the wavelet coefficients and $\lambda_{s}$ is the scale of the coefficient. This metric is strongly equivalent to $W_{1}$. i.e. there exist constants $0<c<C$ such that for all $I_{1},I_{2}\in\mathbb{R}^{N^{2}}$

Different choices of the wavelet basis give different ratios $C/c$. We chose the symmetric Daubechies wavelets of order 5 due to the quality of their approximation (Shirdhonkar and Jacobs, 2008). The sum in Eq. (8) was computed with scale up to $6$ using the PyWavelets package (Lee et al., 2019).

We built a synthetic dataset of 10,000 tomographic projections of the Plasmodium falciparum 80S ribosome bound to the anti-protozoan drug emetine, EMD-2660 (Wong et al., 2014), shown in Figure 1. To generate each image, we randomly rotated the ribosome around its center of mass using a uniform draw of $\mathrm{SO}(3)$, projected it to 2D by summing over the $z$ axis and resized the resulting image to $128\times 128$. Finally we added i.i.d. $\mathcal{N}(0,\sigma^{2})$ noise to the image at different signal-to-noise (SNR) levels. Given a dataset of images $S\in\mathbb{R}^{N\times N\times n}$, the SNR is defined as

We produce three datasets to run experiments on by adding noise at SNR values $\{1/8,1/12,1/16\}$.

We performed rotationally-invariant 2D clustering on our three datasets using rotationally-invariant k-means++ (Algorithms 1, 2) with the $W_{1}$ and $L_{2}$ distance using $k=150$ clusters.

All in-plane rotations at increments of $\pi/100$ radians were tested in the Algorithm 1. To quantify the difference in performance, we computed the distribution of within-cluster viewing angle differences for all pairs of images assigned to the same cluster. For a molecule like the Ribosome that has no symmetries, we expect these angular differences to be concentrated around zero, since large angular differences typically give large differences in the projection images. For all SNR levels, we see that $W_{1}$ gives better angular coherence than $L_{2}$. Note that some of these distributions have a small peak towards 180 degrees which physically represents our algorithms confusing antipodal viewing angles.

In Figure 3 we show the cluster means of the 8 largest clusters at SNR $1/16$. Visually, we can see that the algorithm based on $W_{1}$ produces sharper mean images.

Finally, we examine the occupancy of the clusters. The $W_{1}$ algorithm provides more evenly sized clusters, whereas for the $L_{2}$ algorithm we see a few very populated centers and a large dropoff in occupancy in the other clusters.

When clustering images with Gaussian noise, the averages of larger clusters will tend to be less noisy, since the noise variance is inversely proportional to the cluster size. Due to the lower noise levels, more of the images will be assigned to the larger clusters, making them even larger. This “rich get richer” phenomenon has been observed in cryo-EM (Sorzano et al., 2010). It can explain the large differences in occupancy visible in the top panels of Figure 4, despite the fact that the angles were drawn uniformly. The Wasserstein distance is more resilient to i.i.d. noise and this may explain the uniformity in the resulting cluster sizes seen in the bottom panels of Figure 4.

Finally, clustering with the $W_{1}$ distance does not increase the runtime of the clustering algorithm by much. We include timing results per iteration of k-means with each metric, along with the number of iterations it took for the algorithm to converge in Table 1.

To examine the effect of noise on the $W_{1}^{R}$ and $L_{2}^{R}$ distances, we plot the $W_{1}^{R}$ and $L_{2}^{R}$ distances against the viewing angle difference between projections of the ribosome. In Figure 5 we can see that $W_{1}^{R}$ continues to give meaningful distances under noise compared to the $L_{2}^{R}$ distance.