UNSUPERVISED DOMAIN ALIGNMENT BASED OPEN SET STRUCTURAL RECOGNITION OF MACROMOLECULES CAPTURED BY CRYO-ELECTRON TOMOGRAPHY

2.1. Problem Formulation

Suppose we are given a labeled source data set $D^s=\left\{{\left(x_i^s,y_i^s\right)}_{i=1}^{n_s}\right\}$ from the source domain where each data Point $x_i^s$ is a cubic 3 dimensional grey scale image ie $x_i^s\in R^{H\times H\times H\times 1}$. Furthermore, we have that each $x_i^s$ belongs to a known macromolecule class in C_known or an unknown class in C_unknown where |C_known| = k − 1 and |C_unknown| is not known. We let the labels 1, …, k − 1 represent the known classes, and the label k represent all unknown classes. Our goal is to label a new target data set $D^t=\left\{{\left(x_i^t\right)}_{i=1}^{n_t}\right\}$ from a potentially different domain by correctly assigning each $x_i^t\in C_{\mathit{\text{known}}}$ the label corresponding to its class from {1, …, k − 1} and assigning each $x_i^t\in C_{\mathit{\text{unknown}}}$ the label k.

2.2. Margin-based Discrimination Loss

Center loss [9] can boost the discriminative power of the extracted features in neural networks by learning a vector-like center from deep features and combining it with the cross-entropy loss. Let x^j be the deep feature from the last layer in the Siamese network source branch and let c^j be the center of its class. Center loss is defined as the l2 differences between the deep features and its class center in Equation 1.

We define a margin-based discriminative loss for the deep features from known classes. We set m₁ as the intra-class margin which specifies a maximum distance between deep features with its corresponding class center and m₂ as the inter-class margin which specifies a minimum distance between different classes. Let $H_k^s=\left\{{\left(h_i^s\right)}_{i=1}^{n_k^s}\right\}$ be the set of learned deep features for the data point that belong to the class k. Each learned feature representation $h_i^s$ should be within some distance m₁ from its class’s center. Furthermore, each class center should be at least some distance m₂ from all other class centers. With this in mind, we formulate the margin-based discriminative loss, Equation 2:

Intuitively, the inter-class distance should be larger than the intra-class distance, thus we require m₁ > m₂. In the first term of 2, $c_{y_i}\in R^d$ denotes the y_i-th class center of the deep feature y_i. In the second term, c_i and c_j represent the class centers for two arbitrary (randomly selected) macromolecule classes i and j. The second term of Equation 2 uses these class centers for the arbitrarily selected classes i and j to measure inter-class separability. We update the class centers iteratively with each batch. We use Equations 3 and 4 to update each class center, where y_i is the class of the ith data point in the batch, c_j is cluster j’s center, $h_i^s$ is the deep feature representation of the ith data point in the batch, and δ is an indicator function.

Equation 2 can classify the shared classes. However, for unknown classes the classifier also needs to separate the points from each of the known classes. Following the idea of margin-based loss, we specify a distance between the labeled known classes and the unlabeled unknown data. The margin-based cognition loss is defined by Equation 5 where ${\left\{\left(h_i^s\right)\right\}}_{i=1}^{n_{unk}^s}$ is the set of deep features for the unknown data from the source domain, $min\left({\Vert h_i^s-c_j\Vert }_2^2\right)$ measures the closest distance from each unknown sample to all the known class centers, and m₃ assigns the unknown class margin which provides a minimum length from each unknown macromolecule to any known class center.

2.3. Deep CORAL: Unsupervised Domain Alignment

For a cross-domain open set recognition task, the classifier also needs to decrease the domain discrepancy. Deep CORAL provides a simple but efficient method to match distributions of the middle features in the CNN by minimising the covariance of the source and target features [16]. The CORAL loss is expressed as the distance between the second-order statistics (covariances) of the source and target features as shown in Equation 6.

where H^s and H^t denote the deep features from the output of the bottleneck layer. $\Vert \cdot {\Vert }_F^2$ is the squared matrix Frobenius norm. The covariance matrices of the source and target data are given by Equation 7. h is the number of batch data and 1 is a column vector with all elements equal to 1. The calculation of the coral loss does not need the target labels as reference, so it is an unsupervised method for aligning two domains. Minimizing the correlation alignment (CORAL) can adjust shared weights and reduce the domain discrepancy.

2.4. Training Procedure

Let H_i be the shared feature in the last fully connected layer. Combining all the losses from Equations 2 5 6 with the softmax loss in 8 to reformulate our MLUDA as Equation 9. The softmax is formulated to make the posterior probability of sample x_i.

And, our MLUDA function can be trained by minimizing the weighted combination of L_softmax, L_d, L_c and L_CORAL. Where λ₁ and λ₂ are trade-off parameters used to balance the contribution of each loss. λ₂ should be lower than λ₁ to make sure margin-based loss dominate the total loss. The λ₂ is almost identical in different experiments. But the λ₁ need to use trial and error to find the best value. In Figure 1, the flowchart shows the training process where the loss will be calculated after the fc2 layer in the Siamese network.

3.1. Ablation Experiments on 3D MNIST

3D MNIST contains 3D point clouds generated from the original images of the MNIST dataset.¹ In this experiment, we use the first five classes as the source dataset and the last five classes as the target dataset. We extract the last features from our Network and used t-distributed stochastic neighbor embedding (TSNE) to display the feature space. We ablated each part of our loss and tested their performances separately. As Figure 2 shows, our method can clearly separate the class in the feature space in this cross-domain scenario. (a) means we only use cross entropy as the loss function. (b) means we use cross entropy loss and the Intra-class center as the total loss. In (c), we use cross entropy term and the margin-based discrimination loss we descibed in section 2.2. (d) is our method.

3.2. Experiments on Real Cryo-EM Tomogram