Gum-Net: Unsupervised Geometric Matching for Fast and Accurate 3D Subtomogram Image Alignment and Averaging

Abstract

We propose a Geometric unsupervised matching Network (Gum-Net) for finding the geometric correspondence between two images with application to 3D subtomogram alignment and averaging. Subtomogram alignment is the most important task in cryo-electron tomography (cryo-ET), a revolutionary 3D imaging technique for visualizing the molecular organization of unperturbed cellular landscapes in single cells. However, subtomogram alignment and averaging are very challenging due to severe imaging limits such as noise and missing wedge effects. We introduce an end-to-end trainable architecture with three novel modules specifically designed for preserving feature spatial information and propagating feature matching information. The training is performed in a fully unsupervised fashion to optimize a matching metric. No ground truth transformation information nor category-level or instance-level matching supervision information is needed. After systematic assessments on six real and nine simulated datasets, we demonstrate that Gum-Net reduced the alignment error by 40 to 50% and improved the averaging resolution by 10%. Gum-Net also achieved 70 to 110 times speedup in practice with GPU acceleration compared to state-of-the-art subtomogram alignment methods. Our work is the first 3D unsupervised geometric matching method for images of strong transformation variation and high noise level. The training code, trained model, and datasets are available in our open-source software AITom¹.

1. Introduction

Given a transformation model, geometric matching aims to estimate the geometric correspondence between related images. In two and three dimensions, geometric matching is widely applied to fields such as pattern recognition [1, 2], 3D image reconstruction [3, 4], medical image alignment and registration [5, 6], and computational chemistry [7]. Finding the global optimal parameters consistent with a geometric transformation model such as affine or rigid transformation has a fundamental bottleneck. The parametric space needs to be exhaustively searched but the computational cost is infeasible [8]. Many popular methods have been proposed that alleviate the computational cost by detecting and matching hand-crafted local features [9, 10, 11] to estimate the global geometric transformation robustly [12, 13, 14, 15].

Recently, end-to-end trainable image alignment attracts attention. There are two major advantages over traditional non-trainable methods: (1) a properly trained convolutional neural network (CNN) model can process a large amount of data in a significantly shorter time and (2) with increasing amount of data collected, the deep learning model performance can be improved progressively by better feature learning [16].

In this paper, we focus on an important geometric matching application field, cryo-electron tomography (cryo-ET). In recent years, cryo-ET emerges as a revolutionary in situ 3D structural biology imaging technique for studying macromolecular complexes in single cells, the nano-machines that govern cellular biological processes [17]. Cryo-ET captures the 3D native structure and spatial distribution of all macromolecular complexes together with other subcellular components without disrupting the cell [18]. Nevertheless, cryo-ET data is heavily affected by a low signal-to-noise ratio (SNR) (example input data and mathematical definition in Supplementary Section S3) due to the complex cytoplasm environment and missing wedge effects². Therefore, the macromolecular structures in the 3D tomogram need to be detected and recovered for further biomedical interpretation.

A subtomogram from a tomogram is a small cubic subvolume generally containing one macromolecular complex. Subtomogram alignment is the most critical cryo-ET data processing technique for two reasons: first, high-resolution macromolecular structures can be recovered through subtomogram averaging based on alignment. Second, the spatial distribution of a certain structure can be detected through alignment. To recover the structure, subtomograms containing the same macromolecular structure but in different poses must be iteratively aligned and averaged. Subtomogram averaging improves resolution by reducing noise and missing wedge artifacts [19]. Subtomogram alignment is a considerably more challenging geometric matching task than related tasks such as 3D deformable medical image registration from two aspects: first, there is strong transformation variation because the structure inside a subtomogram is of completely random orientation and displacement. Second, medical images are relatively clean tissue images whereas subtomograms are cellular images with a low SNR (around 0.01 to 0.1) due to the complex cytoplasm environment and the low electron dose used for imaging [20] (example input data in Supplementary Section S3).

Given the 3D rigid transformation model, subtomogram alignment computes the six parameters (three rotational and three translational). We and others have proposed methods [21, 22] to approximate the constrained correlation objective function [23] as heuristics to limit the computational time to a feasible range. However, it is possible nowadays to collect a set of tomograms in several days containing millions of subtomograms [24]. Existing state-of-the-art subtomogram alignment methods [21, 22] generally align a pair of subtomograms on the scale of several seconds, which is too slow for processing such a large amount of data. Moreover, their accuracy is limited because they are approximation methods.

We propose Gum-Net (Geometric unsupervised matching Network), a deep architecture for 3D subtomogram alignment and averaging through unsupervised rigid geometric matching. Integrating three novel modules, Gum-Net inputs two subtomograms to estimate the transformation parameters by extracting and matching convolutional features. Gum-Net achieved significant improvement in efficiency (70 to 110 times speedup) and accuracy (40 to 50 % reduction in alignment error) over two state-of-the-art subtomogram alignment methods [21, 22]. The improvements from proposed modules were demonstrated in three ablation studies.

Main contributions.

Our work is the first 3D unsupervised geometric matching method for images of strong transformation variation and high noise level. We integrated three novel modules (Figure 1): (1) we observe that as the max pooling and averaging pooling operations in the standard deep feature extraction process seek to achieve local transformation invariance, it is not suitable for accurate geometric matching, because the feature spatial locations need to be preserved to a large extent during feature extraction. Therefore, we introduce a feature extraction module with spectral operations including pooling and filtering to preserve the spatial location of extracted features. (2) We propose a novel Siamese matching module that improves spatial correlation information propagation by processing two feature correlation maps in parallel. (3) We incorporate a modified spatial transformer network [25] with a differentiable missing wedge imputation strategy into the alignment module. We achieved fully unsupervised training by feeding into random pairs of subtomograms regardless of their structural class information. Therefore, in contrast to other weakly-supervised geometric matching methods [26, 27, 28, 29, 30], no supervision such as instance-level or category-level matching information is needed.

Figure 1:

Gum-Net model pipeline. The model is unsupervised and feed-forward. The model inputs two subtomograms s_a and s_b (underlying structures are shown in isosurface representation) and outputs the transformed subtomogram ${\widehat{s}}_b$ to geometrically match s_a, in addition to the transformation model parameters ϕ^tr and ϕ^rot. The dash-line denotes that the parameters are shared between the two feature extractors.

2. Related Work

2.1. 2D image alignment based on CNN

2D image alignment usually consists of two steps: (1) obtaining image feature descriptors and (2) matching feature descriptors according to a geometric model. Recently, some methods have employed pre-trained [31] or trainable [32, 33] CNN-based feature extractors. Specifically, [34] proposed a hierarchical metric learning strategy to learn better feature descriptors for geometric matching. However, all the networks are combined with traditional matching methods.

In 2017, Rocco et al. proposed the first end-to-end convolutional neural network for geometric matching of 2D images [35]. This fully supervised model utilizes a pre-trained network [36] to extract features from the two images to be matched. Then a correlation layer matches the features followed by a network to regress to the known transformation parameters for supervised training. Later, they extended this model to be weakly-supervised for finding category-level [27] and instance-level correspondence [26]. Other weakly supervised methods have been proposed for similar tasks including semantic attribute matching [28], simultaneous alignment and segmentation [29], and alignment under large intra-class variations [30]. However, they still require additional training supervision such as matching image pairs on the instance level or category level.

2.2. Unsupervised optical flow estimation

Optical flow estimation describes the small displacements of pixels in a sequence of 2D images using a dense or sparse vector field. Early unsupervised methods have used the gated restricted Boltzmann machine to learn image transformations [37, 38]. Recent CNN-based methods applied techniques such as frame interpolation [39], occlusion reasoning [40], and unsupervised losses in terms of brightness constancy [41] or bidirectional census [42]. Although these methods are all unsupervised, they require their input images to be highly similar with only small pixel shifts.

2.3. Unsupervised deformable medical image registration

3D image registration is the 3D analog to the 2D optical flow estimation. Deformable image registration has been extensively applied to 3D medical images such as brain MRI [43, 44], CT [45, 46], and cardiac images [47, 48]. Recent works present unsupervised CNN models based on spatial transformation function [49, 50, 51] or generative adversarial networks [52, 53]. Similar to optical flow estimation, these methods require the input pair of fixed and moving volumes to be similar. The information from the two volumes is integrated by stacking them as one input to the CNN models. However, simply stacking the input image pairs works poorly when there is strong transformation variation because the image similarity comparison is spatially constrained to a local neighborhood [54].

2.4. Non-learning-based subtomogram alignment

Early works have used exhaustive grid search of rotations and translations with fixed intervals such as 1 voxel and 5° to align subtomograms [55, 56, 57]. To reduce the computational cost of searching the 6D parametric space exhaustively, high-throughput alignment proposed in [21] applied the fast rotational matching algorithm [58]. Fast and accurate alignment proposed in [22] also used the fast rotational matching algorithm and takes into account more information including amplitude and phase into their procedure. Another approach is to collaboratively align multiple subtomograms together based on nuclear-norm [59].

In this paper, we focus on pairwise subtomogram alignment and compared our method against the two most popular subtomogram alignment methods as baselines [21, 22].

3. Method

Our model is shown in Figure 1 (detailed architecture in Supplementary Section S2). Two subtomograms (3D grayscale cubic images) s_a and s_b are processed using feature extractors with shared weights to produce two feature maps v_a and v_b. Then a Siamese matching module computes two correlation maps c_ab and c_ba. At a specific position (i,j,k), c_ab contains the similarity between v_a at that position (i,j,k) and all the features of v_b, whereas c_ba is similarly defined. c_ab and c_ba are processed with the same network architecture and are later concatenated to estimate the transformation parameters. The 6D transformation parameters, which consist of ϕ^tr = {q_x, q_y, q_z} for 3D translation and ϕ^rot = {q_α, q_β, q_γ} for 3D rotation in ZYZ convention, are feed into a differentiable spatial transformer network to compute the output, a transformed subtomogram ${\widehat{s}}_b={\mathcal{T}}_{\varphi }(s_b)={\mathcal{T}}_{{\varphi }^{\text{tr}}}{\mathcal{T}}_{{\varphi }^{\text{rot}}}(s_b)$ with the missing wedge region imputed (Section 3.3). A spectral data imputation technique is integrated into the spatial transformer network to compensate for the missing wedge effects. In the training process, we do not have the ground truth transformation parameters to regress to as in [35]. Therefore, to assess the geometric matching performance, our objective is to find 3D rigid transformation parameters to maximize the cross-correlation between s_a and ${\widehat{s}}_b$ in an unsupervised fashion. The cross-correlation-based loss is back-propagated to update the model weights.

3.1. Feature extraction module

Feature extraction is a dimensionality reduction process to efficiently learn a compact feature vector representation of interesting parts of raw images. There are various popular feature extraction techniques such as DenseNet [60], InceptionNet [61], and ResNet [62]. Subsampling methods such as max pooling and average pooling are used in these convolutional neural networks to reduce feature map dimensionality and facilitate computation. Compared with max pooling and average pooling, spectral representation for convolutional neural networks preserves considerably more spatial information per parameter and enables flexibility in the pooling output dimensionality [63]. 2D spectral pooling layers that perform dimension reduction in the frequency domain have been proposed based on discrete Fourier transform (DFT) [63], discrete cosine transform (DCT) [64], and Hartley transform [65]. However, these methods are designed for 2D images and do not take into account image noise.

We propose a 3D DCT-based spectral layer with pooling and filtering operations. Since our inputs are 3D noisy images, the novel filtering operation is for feature map high-frequency noise reduction, and pooling operation for feature map dimension reduction. We choose the DCT because it stores only real-valued coefficients and compacts more energy in a smaller portion of the spectra compared to the DFT [66].

For an input feature map $v\in {\mathbb{R}}^{L\times W\times H}$, its 3D type-II DCT is defined as [67]:

$$\begin{gathered} \mathbb{C}(v{)}_{lhw}=\frac{8}{LWH}{ϵ}_l{ϵ}_h{ϵ}_w\sum _{i=0}^{L-1}\sum _{j=0}^{H-1}\sum _{k=0}^{W-1}{v}_{ijk}\cos \left(\frac{l\pi (2i+1)}{2L}\right) \\ \cos \left(\frac{h\pi (2j+1)}{2H}\right)\cos \left(\frac{w\pi (2k+1)}{2W}\right), \end{gathered}$$ (1)

where ${ϵ}_l=\{\begin{array}{cc}\frac{1}{\sqrt{2}}\hfill & \text{for}l=0\hfill \\ 1\hfill & \text{otherwise}\hfill \end{array}\phantom{\}}$, ∀l ∈ {0, …, L – 1}. ϵ_h and ϵ_w are similarly defined ∀h ∈ {0, …, H – 1}, ∀w ∈ {0, …, W – 1}. The inverse transform $\mathbb{C}$ of 3D type-II DCT is well defined as 3D type-III DCT [67]. Therefore, the pooled and filtered representation in the frequency domain can be transformed back through type-III DCT to the spatial domain as the output of the layer.

We use the DCT to perform subsampling in which the input is transformed to the frequency domain and cropped there. The output with reduced dimensionality is computed by transforming the cropped spectrum back into the spatial domain. The spectral pooling operation has been shown to achieve better spatial information preservation per parameter in terms of the l₂ norm as compared to the max pooling operation [63]. Figure 2 shows the image reconstruction from max pooling, average pooling, and DCT spectral pooling at different subsampling factors. Compared to other pooling operations, a major advantage of using spectral pooling & filtering layers for geometric matching tasks is that the spatial location of features in two images are significantly better preserved for accurate matching. For example, during max pooling, the maximum from the receptive field is selected to achieve local rotation and translation invariance with the intuition that the exact location of a feature does not matter to the final classification. By contrast, during the feature extraction step for geometric matching, the exact feature spatial location is critical and the information loss will lead to inaccurate downstream matching.

Figure 2:

Image reconstruction from the max pooling, average pooling, and DCT spectral pooling scheme at different subsampling factors. DCT spectral pooling retains substantially greater spatial information of features from the original image and offers arbitrary output map dimensionality.

We implement the 3D DCT spectral pooling & filtering as differentiable layers in the feature extractor. The low-pass filtering is also performed by masking out high-frequency regions dominated by noise. The forward and back-propagation procedure of the 3D DCT spectral pooling & filtering layer is outlined in Algorithm 1 and 2.

The arbitrary output size of spectral pooling & filtering layers offers another major advantage for geometric matching tasks. If the output two feature maps are of size L × W × H with C channels, the Siamese correlation layer (Section 3.2) will create two correlation maps, each of size L ×W × H with (LWH) channels. The output feature map size from the feature extraction module to the Siamese matching module needs to be carefully manipulated, especially for 3D images. If the output feature map is too small, such as 3 × 3 × 3, there is too much information loss for matching. If the output feature map is too large, such as 20 × 20 × 20, the resulting correlation maps will be of size 20 × 20 × 20 × 8000, which is too large to be processed. Unlike max pooling or average pooling layers which aggressively reduce each dimension to half of the size and remove 87.5 % of the information, spectral pooling & filtering layers can gradually reduce the feature map size to the desired feature extraction module output size. Therefore, no additional spatial cropping or padding layer is needed to control the feature map size.

Algorithm 1 DCT spectral pooling & filtering

$$\begin{array}{cc}\hfill & \text{OutOutput FnProcedure}\hfill \\ \hfill & \text{Feature map}v\in {\mathbb{R}}^{L\times W\times H}\hfill \\ \hfill & \text{Output size}{L}_1\times W_1\times H_1\hfill \\ \hfill & \text{Cropping size}{L}_2\times W_2\times H_2\hfill \\ \hfill & \text{Feature map}\widehat{v}\in {\mathbb{R}}^{L_1\times W_1\times H_1}u\leftarrow \mathbb{C}(v)\hfill \\ \hfill & u\leftarrow \text{Crop}u\text{to size}{L}_2\times W_2\times H_2\hfill \\ \hfill & \widehat{u}\leftarrow \text{ZeroPad}u\text{to size}{L}_1\times W_1\times H_1\hfill \\ \hfill & \widehat{v}\leftarrow \mathbb{C}(\widehat{u})\hfill \end{array}$$

Algorithm 2 DCT spectral pooling & filtering back-propagation

$$\begin{array}{cc}\hfill & \text{OutOutput FnProcedure}\hfill \\ \hfill & \text{Gradient w.r.t layer output}\frac{\partial L}{\partial \widehat{v}}\text{Gradient w.r.t layer input}\frac{\partial L}{\partial v}y\leftarrow \mathbb{C}(\frac{\partial L}{\partial \widehat{v}})\hfill \\ \hfill & y\leftarrow \text{Crop}y\text{to size}{L}_2\times W_2\times H_2\hfill \\ \hfill & \widehat{y}\leftarrow \text{ZeroPad}y\text{to size}L\times W\times H\hfill \\ \hfill & \frac{\partial L}{\partial v}\leftarrow \mathbb{C}(\widehat{y})\hfill \end{array}$$

3.2. Siamese matching module

The matching of extracted features from images is usually performed as an independent post-processing step [68, 69, 70, 71, 72]. The 2D correlation layer proposed in [35] achieved the state-of-the-art for integrating the matching information from two images. It is essentially a normalized cross-correlation function $G:{\mathbb{R}}^{H\times W\times C}\times {\mathbb{R}}^{H\times W\times C}\to {\mathbb{R}}^{H\times W\times (HW)}$. One of the input feature maps v_a is first flattened into shape $v_a\in {\mathbb{R}}^{N\times C}$, where N = HW, in order to keep the output correlation map 2D. Then for each feature (pixel) in v_a and v_b, the dot product is computed over all the channels (as feature descriptors) to obtain the correlation, which is later normalized. Nevertheless, to control the dimension of the output correlation map, all axes of one input feature map are broken and later cast into the channels of the output whereas the other input feature map is preserved.

We propose a novel Siamese matching module for pairwise 3D feature matching. To better utilize and process the feature correlation information, we design a Siamese correlation layer. Different from the correlation layer in [35] which computes only c_ab, the Siamese correlation layer is intuitive and symmetrically designed, which computes two correlation maps c_ab and c_ba. Each of them preserves the spatial coordinates of one input feature map. The use of two correlation maps propagates more feature spatial correlation information for the transformation parameter estimation. Element at a specific position lwhc is defined as:

$$\begin{gathered} (c_{ab}{)}_{lwhc}=\frac{\langle v_{a_{n:}},v_{b_{lwh:}}\rangle }{\sqrt{{\sum }_{i,j,k}\langle v_{a_{n:}},v_{b_{ijk:}}\rangle }}, \\ (c_{ba}{)}_{lwhc}=\frac{\langle v_{b_{n:}},v_{a_{lwh:}}\rangle }{\sqrt{{\sum }_{i,j,k}\langle v_{b_{n:}},v_{a_{ijk:}}\rangle }}. \end{gathered}$$ (2)

The two correlation maps are feed into a pseudo-Siamese network consisting of convolution layers and convolved separately but later concatenated for one fully connected layer. After another fully connected layer, the Siamese matching module outputs the estimated rigid transformation parameters ϕ^tr and ϕ^rot. Detailed model architecture can be found in Supplementary Section S2.

3.3. Unsupervised geometric alignment module

Existing subtomogram alignment methods optimize a matching metric [21, 22, 57, 73]. In practice, preparing the subtomogram alignment ground truth for training is extremely time-consuming (need to exhaustively search the 6D parametric space). Therefore, the deep model should be unsupervised for this task. To achieve this goal, we propose an unsupervised geometric alignment module utilizing the spatial transformer network [25] with spectral data imputation designed specifically for subtomogram data.

In a tomogram with fixed voxel spacing (around 1nm), a certain type of macromolecular structure does not scale or reflect. Therefore, we restrict ourselves to 3D rigid transformation. Denoting the transformation matrix generated by the 3D rigid transformation parameters as M_θ [74] and the 3D warping as ${\mathcal{T}}_{\varphi }:{\mathbb{R}}^3\to {\mathbb{R}}^3$, we have:

$$\begin{gathered} \left(\begin{array}{c}\hfill x_i^s\hfill \\ \hfill y_i^s\hfill \\ \hfill z_i^s\hfill \\ \hfill 1\hfill \end{array}\right)={\mathcal{T}}_{\varphi }(x_i^t,y_i^t,z_i^t)=M_{\theta }\left(\begin{array}{c}\hfill x_i^t\hfill \\ \hfill y_i^t\hfill \\ \hfill z_i^t\hfill \\ \hfill 1\hfill \end{array}\right) \\ =\left[\begin{array}{cccc}\hfill {\theta }_{11}\hfill & \hfill {\theta }_{12}\hfill & \hfill {\theta }_{13}\hfill & \hfill {\theta }_{14}\hfill \\ \hfill {\theta }_{21}\hfill & \hfill {\theta }_{22}\hfill & \hfill {\theta }_{23}\hfill & \hfill {\theta }_{24}\hfill \\ \hfill {\theta }_{31}\hfill & \hfill {\theta }_{32}\hfill & \hfill {\theta }_{33}\hfill & \hfill {\theta }_{34}\hfill \\ \hfill 0\hfill & \hfill 0\hfill & \hfill 0\hfill & \hfill 1\hfill \end{array}\right]\left(\begin{array}{c}\hfill x_i^t\hfill \\ \hfill y_i^t\hfill \\ \hfill z_i^t\hfill \\ \hfill 1\hfill \end{array}\right), \end{gathered}$$ (3)

where ($x_i^t$, $y_i^t$, $z_i^t$) is the target coordinates on the transformed output 3D image and ($x_i^s$, $y_i^s$, $z_i^s$) is the source coordinates on the input 3D image. θ is an element of the transformation matrix. The 3 × 3 orthogonal rotation matrix is from θ₁₁ to θ₃₃. The displacement along each axis is specified by θ₁₄, θ₂₄, and θ₃₄. The 3D warping is differentiable and therefore able to be trained end-to-end.

In order to compensate for the missing wedge effects and thus to decrease the bias introduced, we integrate a spectral data imputation strategy from our previous work [75] into the spatial transformer network. For a subtomogram, we use its current estimated transformation to compute the rotated missing wedge mask m, as an indicator function to represent whether the Fourier coefficients are valid or missing in certain regions, and impute the missing ones with those from its transformation target subtomogram s_a. We can form a transformed and imputed subtomogram ${\widehat{s}}_b$ such that:

$$\begin{gathered} (\mathcal{F}{\widehat{s}}_b)(\xi )=\{\begin{array}{cc}[\mathcal{F}{s}_a](\xi )\hfill & \text{if}m(\xi )=0\hfill \\ [\mathcal{F}{\mathcal{T}}_{\varphi }(s_b)](\xi )\hfill & \text{if}m(\xi )=1\hfill \end{array}\phantom{\}}, \\ m(\xi )=\{\begin{array}{cc}0\hfill & \text{if the Fourier coefficient at}\xi \text{is missing}\hfill \\ 1\hfill & \text{if the Fourier coefficient at}\xi \text{is valid}\hfill \end{array}\phantom{\}}, \end{gathered}$$ (4)

where $\mathcal{F}$ is the Fourier transform operator, $\xi \in {\mathbb{R}}^3$ is a Fourier space location, and m(ξ) is the rotated missing wedge mask according to ϕ^rot. Since the magnitude of Fourier transform is translation-invariant, we only need to rotate m(ξ) without using ϕ^tr [23]. The imputation operation facilitates the unsupervised geometric matching task because only when the optimal alignment is obtained, the imputed data results in the highest consistency with the transformed subtomogram.

We note that since the rotation of the missing wedge mask m is implemented along with the transformation of the input subtomogram in the differentiable spatial transformer network and the inverse discrete Fourier transformation is well defined, this spectral data imputation step is differentiable in a similar manner as Algorithm 2.

Loss function.

Pearson’ correlation and its variants are widely used for assessing the alignment between two subtomograms [21, 22, 23, 57, 73] because of its simplicity and effectiveness. We implement it as a loss function to Gum-Net:

$$\mathcal{L}=1-\frac{{\sum }_{i=1}^N(s_{a_i}-{\overline{s}}_a)({\widehat{s}}_{b_i}-{\overline{\widehat{s}}}_b)}{\sqrt{{\sum }_{i=1}^N(s_{a_i}-{\overline{s}}_a{)}^2}\sqrt{{\sum }_{i=1}^N({\widehat{s}}_{b_i}-{\overline{\widehat{s}}}_b{)}^2}},$$ (5)

where N is the total number of voxels in an input subtomogram. Compared to existing methods [21, 22], which utilize translation-invariant upper bound to approximate the Pearson’s correlation objective to reduce the computational cost, Gum-Net optimizes Pearson’s correlation directly for more accurate alignment.

3.4. Baseline methods

We implemented two most popular state-of-the-art subtomogram alignment methods for comparison: H-T align [21] and F&A align [22]. We performed three ablation studies with existing modules: Gum-Net Max Pooling (Gum-Net MP), Gum-Net Average Pooling (Gum-Net AP), and Gum-Net Single Correlation (Gum-Net SC). Detailed implementation can be found in Supplementary Section S2.

4. Experiments

Gum-Net was evaluated on six real and nine realistically simulated datasets at different SNR. On the simulated datasets, the accuracy of subtomogram alignment was evaluated by comparing the estimated transformation parameters ϕ^tr and ϕ^rot to the ground truth. On the real datasets, since the transformation ground truth is not available, in practice, the optimal transformation is usually obtained by parametric space exhaustive grid search to optimize the cross-correlation between s_a and ${\widehat{s}}_b$. Therefore, we compared the cross-correlation between s_a and ${\widehat{s}}_b$ computed by Gum-Net and baseline methods as an indirect indicator of the alignment accuracy. The visualization of subtomograms in different datasets can be found in Supplementary Section S3.

4.1. Datasets

4.1.1. Real datasets

GroEL/GroES dataset: this dataset contains 786 experimental subtomograms of purified GroEL and GroEL/GroES complexes from 24 tomograms [23]. Each subtomogram is rescaled to size 32³ with voxel size 0.933 nm and 25° missing wedge.

Rat neuron culture dataset: this recent dataset is a set of tomograms from rat neuron culture [76]. In total 1095 ribosome subtomograms and 1527 capped proteasome subtomograms were extracted by template matching [55] and biology expert annotation. Each subtomogram is of size 32³ with voxel size 1.368 nm and 30° missing wedge.

S. cerevisiae 80S ribosome dataset: this dataset contains 3120 subtomograms extracted from 7 tomograms of purified S. cerevisiae 80S ribosomes [77]. Each subtomogram is rescaled to size 32³ with voxel size 1.365 nm and 30° missing wedge.

TMV dataset: this dataset contains 2742 Tobacco Mosaic Virus (TMV) subtomograms, a type of helical virus [78]. Each subtomogram is binned to size 32³ with voxel size 1.080 nm and 30° missing wedge.

Aldolase dataset: this recent dataset contains 400 purified rabbit muscle aldolase subtomograms [79]. Each subtomogram is rescaled to size 32³ with voxel size 0.750 nm and 30° missing wedge.

Insulin receptor dataset: this recent dataset contains 400 purified human insulin-bounded insulin receptor subtomograms [80]. Each subtomogram is rescaled to size 32³ with voxel size 0.876 nm and 45° missing wedge.

4.1.2. Simulated datasets

The subtomogram dataset simulation utilized a standard procedure in [81, 82] which takes into account the tomographic reconstruction process with missing wedges and contrast transfer function (detailed simulation procedure in Supplementary Section S3). We chose five representative macromolecular complexes: spliceosome (PDB ID: 5LQW), RNA polymerase-rifampicin complex (1I6V), RNA polymerase II elongation complex (6A5L), ribosome (5T2C), and capped proteasome (5MPA). All five structures are asymmetric so that there exists only one alignment ground truth. We simulated five datasets, one relatively clean (SNR 100) and four with SNR close to the experimental conditions (0.1, 0.05, 0.03, and 0.01), each consists of 2100 subtomogram pairs of each structure (in total 10500 subtomogram pairs). 5000 subtomogram pairs from each dataset were used for training and 500 pairs for validation. The rest 5000 subtomogram pairs from each dataset are used for testing. For a pair of subtomograms, one structure is a randomly transformed copy of the other and the two structures were processed independently to obtain its tomographic image distortions. Each subtomogram is of size 32³ with voxel size 1.2 nm. The s_b in each pair has a typical missing wedge 30° while s_a has no missing wedge.

For subtomogram averaging, we simulated four datasets of 500 ribosomes (PDB ID: 5T2C) in the same manner of SNR 0.1, 0.05, 0.03, and 0.01.

4.2. Implementation

The deep models were implemented in Keras [83] with custom layers backend by Tensorflow [84]. All inputs have size 32³. We note that due to the flexibility of input and output size of the DCT spectral pooling & filtering layers, the input size can be arbitrary. Higher resolution can be achieved with larger input subtomogram sizes. Detailed implementation of Gum-Net and baselines can be found in Supplementary Section S2.

For each epoch, we randomly draw 5000 subtomogram pairs s_a and s_b from the training dataset regardless of their structural class information. Therefore, Gum-Net is fully unsupervised without instance-level or category-level matching information for weak supervision as in other geometric matching methods [26, 27, 28, 29, 30]. For a simulated dataset, there are 5000² possible image pairs. As a result, we did not observe any overfitting issue.

4.3. Subtomogram alignment

Given the transformation ground truth, we measure the alignment accuracy with two metrics: (1) the translation error defined as the Euclidean distance between the translation estimation and the ground truth and (2) the rotation error defined as the Euclidean distance between the flattened rotation matrix of estimation and the ground truth.

On simulated datasets:

Table 3.4 shows the alignment accuracy. Gum-Net achieved similar performance on the clean dataset (SNR 100). As max pooling achieves more local transformation invariance [85], Gum-Net MP performs worse than Gum-Net AP in all settings as expected. When the SNR is close to experimental condition (the real datasets have SNR around 0.01 to 0.1), CNN-based methods generally perform better than traditional methods. Specifically, Gum-Net outperformed all the baseline methods, demonstrating the improvement from the proposed modules.

In our experiments, the training, validation, and testing datasets are independent, which ensured no overfitting. However, since Gum-Net is fully unsupervised, even if the testing dataset is from a different domain source, such as collected under different imaging conditions, it is possible to fine-tune a trained model on the testing dataset (with no ground truth) for adaptation. In terms of speed, with a trained model, Gum-Net only takes 17.6 seconds to align 1000 subtomograms on a single GPU core. The training takes less than 10 hours. Since there is no available GPU-accelerated version of the traditional algorithms, H-T align and F & A align take 1916.4 seconds and 1251.2 seconds to align 1000 subtomograms on a CPU core, respectively. Therefore, in practice, this results in 70 to 110 times speedup over traditional methods.

On real datasets:

We split the GroEL/GroES dataset into a training dataset of 617 subtomograms, a validation dataset of 69 subtomograms, and a testing dataset of 100 subtomograms. There are 4950 pairs of subtomograms in the testing dataset. We align them pairwise by Gum-Net, H-T align, and F&A align and calculates the cross-correlation. Gum-Net achieved cross-correlations of 0.0908±0.0204, significantly better (p < 0.001) than H-T align (0.0756±0.0194) and F&A align (0.0838±0.0204).

We split the rat neuron culture dataset into a training dataset of 2270 subtomograms, a validation dataset of 252 subtomograms, and a testing dataset of 100 ribosome and 100 capped proteasome subtomograms. There are 19900 pairs of subtomograms in the testing dataset. Gum-Net achieved cross-correlations of 0.0615±0.0187, significantly better (p < 0.001) than H-T align (0.0541±0.0235) and F&A align (0.0607±0.0199). We use the pairwise correlation matrix to cluster the subtomograms by defining the pairwise distance as 1 - pairwise correlation. Applying the complete-linkage hierarchical clustering algorithm with k = 2, Gum-Net achieved an accuracy of 92%, better than F&A align (65%) and H-T align (53.5%).

4.4. Non-parametric reference-free subtomogram averaging

Structures present in multiple noisy copies (usually thousands of) in a tomogram must be averaged through geometric transformation to obtain higher resolution 3D views [19]. To eliminate potential bias, subtomogram averaging is often done without any external structural reference. One major approach of reference-free subtomogram averaging is non-parametric alignment-based averaging in which all subtomograms are iteratively aligned to their average and re-averaged for the next iteration [86]. Figure 4 illustrates such a process in which the initial average is generated by simply averaging all the subtomograms without any transformation. The structural resolution of the subtomogram average is gradually improved through the iterative process.

Figure 4:

Illustration of alignment-based subtomogram averaging using Gum-Net. On the left are five example input subtomograms at SNR 0.1 in our experiment. On the right are subtomogram averages at different iterations and the true structure. The 2D slices representations are shown in Supplementary Section S3.

The iterative alignment-based non-parametric reference-free subtomogram averaging was tested using the proposed and baseline methods. The standard resolution measurement for assessing subtomogram averaging is Fourier shell correlation (FSC) [87] (mathematical definition in Supplementary Section S3), which measures the maximal discrepant structural factors between the subtomogram average and the true structure. The smaller the value, the better the results. As shown in Table 4.4, Gum-Net achieved the overall best averaging performance and improved the resolution by around 10%.

5. Conclusion

Cryo-ET subtomogram alignment and averaging revolutionize the discovery of 3D native macromolecular structure details in single cells. Such information provides critical insights into the precise function/dysfunction of the cellular processes. However, with a rapidly increasing amount of cryo-ET data collected, there is an urgent need to drastically improve the efficiency of subtomogram alignment methods. We developed the first unsupervised deep learning approach for 3D subtomogram alignment and averaging. Using the three proposed modules, Gum-Net achieved fast and accurate alignment with end-to-end unsupervised learning. Gum-Net opens up the possibility for continued improvement of subtomogram alignment and averaging efficiency and accuracy with better model design and training. This work serves as an important step toward in situ high-throughput detection and recovery of macromolecular structures for a better understanding of the molecular machinery in cellular processes.

Gum-Net can be integrated into existing cryo-ET analysis software in several ways. For example, EMAN2 [81] performs exhaustive 3D rotational and translational search followed by local refinement for alignment-based averaging. RELION [77] maximizes the likelihood of a model with Gaussian noise assumption by exhaustively scanning the 3D rigid transformation space for integration. Gum-Net improves the accuracy and efficiency of subtomogram alignment, especially for a large amount of cryo-ET data. Therefore, integrating Gum-Net with existing software can boost the speed of their alignment step or quickly generate initial structural models for averaging refinement. Gum-Net can also be easily extended to related tasks including tomographic tilt series alignment [88] and cryo-electron microscopy single-particle reconstruction [89]. The proposed modules can be adapted to other geometric matching tasks for images of strong transformation variation such as face alignment under pose variations [90, 91], or of high noise level such as synthetic aperture radar imaging [92, 93] and sonar imaging [94, 95].

Figure 3:

Example alignment inputs and outputs at SNR 100. 2D slices representations are shown in Supplementary Section S3.

Table 1:

Subtomogram alignment accuracy on five datasets with SNR specified. In each cell, the first term is the mean and standard deviation of the rotation error and the second term, the translation error. We highlighted Gum-Net results that are significantly better (p < 0.001) than all baselines by the paired sample t-test. More detailed results and analysis can be found in Supplementary Section S3.

Method	SNR 100	SNR 0.1	SNR 0.05	SNR 0.03	SNR 0.01
H-T align	0.30±0.68, 1.82±2.69	1.22±1.07, 4.76±4.56	1.93±0.98, 7.26±4.77	2.22±0.77, 8.86±4.72	2.38±0.57, 11.33±5.02
F&A align	0.33±0.70, 1.93±2.86	1.34±1.13, 5.39±4.90	1.95±0.98, 7.54±4.94	2.22±0.77, 8.99±4.81	2.38±0.57, 11.32±4.92
Gum-Net MP	0.90±0.87, 3.34±3.41	1.30±0.79, 4.93±3.36	1.44±0.79, 5.46±3.38	1.53±0.78, 5.96±3.34	1.67±0.77, 7.28±3.38
Gum-Net AP	0.60±0.71, 2.32±2.71	1.09±0.73, 4.20±2.96	1.30±0.77, 5.00±3.15	1.45±0.77, 5.70±3.25	1.65±0.78, 7.18±3.35
Gum-Net SC	0.70±0.75, 2.63±2.86	1.16±0.77, 4.41±3.23	1.36±0.79, 5.13±3.34	1.48±0.78, 5.75±3.34	1.67±0.77, 7.24±3.46
Gum-Net	0.41±0.70, 1.59±2.63	0.62±0.69, 2.41±2.61	0.87±0.74, 3.20±2.78	1.13±0.75, 4.29±2.75	1.50±0.78, 6.78±4.22

Table 2:

Subtomogram averaging results in FSC resolution (nm). ‘0.1’ denotes simulated dataset at SNR 0.1. ‘80S’, ‘TMV’, ‘Aldolase’, and ‘Insulin’ denote the real datasets. The best resolution is highlighted.

Method	0.1	0.05	0.03	0.01	80S	TMV	Aldolase	Insulin
H-T align	2.89	3.79	4.92	4.41	3.05	2.23	2.34	1.90
F&A align	2.78	4.36	3.81	4.53	2.77	2.52	3.13	2.18
Gum-Net	2.78	2.95	4.01	4.22	2.73	2.16	1.97	1.77