Affine Equivariant Autoencoder

1 Introduction

Deep neural networks (DNNs) have brought breakthroughs in different machine learning domains and applications. To accomplish the task at hand, DNNs learn task-specific features and discard the information that is irrelevant to the current task. For example, applying local transformation to an image example will not change the feature representation in a DNN for classification. This kind of feature is called transformation invariant feature. The transformation invariance property improves the robustness of the feature and favors the task at hand. However, in some cases where there's no specific task to achieve, especially in representation learning field [Bengio et al., 2013], the feature is required to preserve as much information as possible. To capture the transformation information which is directly discarded by traditional DNNs, the feature should change accordingly with the transformation of the input image. We call the feature with the above property transformation equivariant feature.

The transforming autoencoder [Hinton et al., 2011] is the first work to mention about the importance of the equivariant feature. But it is only an approximation of the given transformation function. Neither theoretical guarantee nor empirical validation is provided to prove the equivariance. Capsule network [Sabour et al., 2017] encodes the instantiation parameters to a group of neurons but can not control what parameters to learn. The rotation equivariant network [Li et al., 2018] is the only work that explicitly incorporates the equivariance property. But it can only achieve the equivariance to rotation transformation by defining a series of complex layers.

∗Corresponding Authors

In this paper, we propose an affine equivariant autoencoder (AEAE) to learn features that are equivariant to the affine transformation in an unsupervised manner. We first derive an objective to directly satisfy the definition of the equivariance. Then we relax the objective to avoid instability. The resulting objective consists of three parts: 1) self-reconstruction of the original example, 2) self-reconstruction of the affine transformed example, and 3) approximation of the affine transformation function. The first two parts ensure the content information is captured by the encoder network. The third part incorporates the equivariance property to the encoder. We conduct experiments to validate the equivariance of the learned features both quantitatively and qualitatively. We also perform clustering and classification on the learned features to prove that the feature is discriminative.

Our main contributions are summarized as follows.

• We propose an affine equivariant autoencoder (AEAE) to learn deep features that are equivariant to the affine transformation. To the best of our knowledge, this is the first work to explicitly learn equivariant features in an unsupervised manner.
• We design experiments to validate the equivariance of the learned features both quantitatively and qualitatively.
• Our model achieves the state-of-the-art clustering performance by simply running spectral clustering on the learned features.

2 Affine Equivariant Autoencoder

2.1 Equivariance and Invariance

Given two sets X and Z, and a mapping function f : X → Z. Formally, the mapping f is equivariant to a family of transformations T , if for any transformation T ∈ T there exists a corresponding transformation T ′ , such that

And the mapping f is invariant to T , if for all T ∈ T the following equation holds:

In this paper, we only consider image datasets and the affine transformation. Specifically, let X ⊂ Rc×w×h be an image dataset in which each element x is an image represented by a c × w × h matrix where c, w, h are channels, width and height, respectively. The Z ⊂ Rd is comprised of d-dimensional feature vectors of images in X. Let Tσ be parameterized by σ ∈ R5 where σ0 represents the degree of rotation anticlockwise, σ1 the pixels of translation rightward, σ2 the pixels of translation downward, σ3 the scaling factor, and σ4 the degree of shear transformation.

To disentangle affine parameters from feature vectors, we force the transformation T ′ to be parameterized by σ and to take a simple form:

where t linearly transforms the elements in σ to a reasonable range and expands the size of σ from 5 to d by padding zeros. Therefore, we aim to find a feature extractor f that satisfies

2.2 Model Description

To find a feature extractor f that satisfies (4), we implement f by a deep neural network. Given a dataset with n examples X = {xi}ni=1 ⊂ X ⊂ Rc×w×h . We define an encoder network fw(·) to transform each example xi to zi ∈ Z ⊂ Rd and a decoder network gu(·) to map zi back to xi where w, u are the sets of trainable weights in encoder and decoder networks, respectively.

We need to make sure that the autoencoder can reconstruct all x and Tσ(x) in order to extract their features by fw. Therefore, the ideal objective is:

where ‖ · ‖2 denotes the sum of squared errors and ‖ · ‖2 is the ℓ2 norm. The first two terms enable fw to be a valid feature extractor for both x and Tσ(x), and the last term makes fw equivariant to affine transformation T . If the above conditions are met, z = fw(x) is called affine equivariant feature.

However, from the perspective of optimization, the error signal of the last term can hardly back-propagate to the encoder's layers through two paths. The typical solution is to block one path and back-propagate through another. But we find this solution is unstable and sometimes leads to divergence of the objective. To solve this problem, we propose an alternative objective:

We prove (6) is approximately equivalent to (5) under the condition that t(σ) is in a small range near zero. When (6) reaches the minimal value 0, we have

Apply fw to both sides, we get

The first equation in (8) implies fw(gu(·)) is an identity mapping for fw(xi) ∈ Z. If we restrict t(σ) to a small range near 0, we have

which means the equivariance of fw can be approximately satisfied by optimizing (6). This completes the proof.

Another benefit of (6) is that the affine transformation Tσ(xi) is explicitly approximated by gu(fw(xi) + t(σ)). So we can validate whether the feature vector disentangles the affine factors by adding turbulence to the feature and observing the reconstructed image.

Finally, we depict the proposed model in Figure 1. There are three data flows denoted by blue, green, and red dashed lines, corresponding to three terms in (6). The blue data flow represents a standard autoencoder which takes an example x as input and tries to reconstruct x. The green data flow also defines a standard autoencoder with the input and target are the affine transformed example Tσ(x). The blue and green autoencoders share the same weights {w, u}, enabling the encoder fw to be a valid feature extractor for both x and Tσ(x). The red data flow tries to reconstruct Tσ(x) with the clean example x and affine parameter σ as inputs. The affine parameter σ ∈ R5 is scaled by element and then added to the last five elements of the feature vector fw(x). The resulting feature is fed to the decoder to construct Tσ(x). This red flow learns an approximate mapping for the affine transformation Tσ and ensures the equivariance of fw to affine transformations T .

2.3 Implementation

The autoencoder's structure can either be fully connected or convolutional in theory. However, convolutional networks, especially those including pooling layers, are known to be invariant to local transformations, which is harmful to our goal of learning an equivariant mapping. Therefore, we implement the encoder fw and decoder fu by using fully connected layers instead of convolutional layers. The detailed structure is as shown in Figure 2. There are total eight fully connected layers with number of neurons D - 500 - 500 - 2000 - 15 - 2000 - 500 - 500 - D, where D = cwh is the dimension of the vectorized input image. Given an image example x ∈ Rc×w×h which can be optionally transformed by affine transformation Tσ, we flatten it to a vector with length cwh and then feed it into the proposed model. The output is also a vector with length cwh which is finally reshaped to c×w×h in order to compute the distance ‖ · ‖ with x. Except the embedding (with 15 neurons) and output layers are followed by linear and sigmoid activations, respectively, all internal layers are activated by ReLU [Glorot et al., 2011].

We limit the parameters σ of affine transformations to a reasonable range as shown in Table 1. The function t is defined as

in order to keep t(σ) in a small range near 0 which is a necessary condition of the equivalence of (5) and (6) as described in Section 2.2.

We initialize all weights by following [He et al., 2015]. The AEAE is trained in an end-to-end manner by using Adam [Kingma and Ba, 2014] optimizer with initial learning rate 0.001, β1 = 0.9, and β2 = 0.999. The maximum number of epochs is set to 100 for large datasets (n > 10000) and 500 for small ones (n <= 10000). The mini-batch size is fixed to 256. The source code is publicly available on https://github.com/XifengGuo/AEAE.

3 Experiments

We conduct extensive experiments to validate the effectiveness of the proposed AEAE. First, we demonstrate the equivariance property quantitatively and qualitatively. Then we assess the quality of features learned by our AEAE according to the clustering and classification performances. Four image datasets are used in our experiments:

• MNIST-full: A popular handwritten digit dataset with 70,000 examples [LeCun et al., 1998].
• MNIST-test: A dataset that only contains the test set of MNIST-full, with 10,000 examples. We consider this a separate dataset because the number of examples may affect the clustering performance, by following [Yang et al., 2016; Dizaji et al., 2017].
• USPS 1 : A dataset contains 9298 gray digit images with size of 16x16 pixels divided into 10 categories.
• Fashion: A dataset of Zalando's article images [Xiao et al., 2017], sharing the same statistics with MNIST-full but more challenging for machine learning algorithms.

Table 2 shows the statistics of these datasets. All image examples in these datasets are padded two pixels of zeros in each of four directions to avoid information loss during the affine transformation. The pixel values are scaled to [0, 1].

3.1 Equivariance Validation

We empirically validate the equivariance of the encoder mapping fw on MNIST-full and Fashion datasets. We directly check if (4) is satisfied after optimizing (6). To this end, we first compute the difference between features of original and affine transformed examples:

1 http://www.cad.zju.edu.cn/home/dengcai/Data/MLData.html

The σ is restricted to have only one element active, i.e, only one of rotation, horizontal translation, vertical translation, scaling, and shear transformations is applied to all original examples. According to the definition of equivariance and invariance, if ∆i = 0 for all i, fw will be invariant to the affine transformation. However, we observe ∆i changes linearly with the affine parameter σ. As shown in the first row in Figure 3, the jth element of the average ∆ = ∑ni=1 ∆i is positively correlated to the kth element of σ where j = 10+k, k = 0, 1, . . . , 4. To be exact, ∆ is close to its expectation t(σ). On the other hand, we record the average value of ‖fw(Tσ(xi)) - (fw(xi) + t(σ))‖22 over i in the second row in Figure 3. For all transformations, this value is close to 0. To conclude, (4) is approximately satisfied when σ is limited by Table 1. That is, the empirical quantitative evidence supports that fw is equivariant to the affine transformation.

In addition to the direct validation of the equivariance property, we give an indirect evidence as follows. As shown in the last term of (6), we attempt to reconstruct the affine transformed example Tσ(x) from the original example x with affine parameters σ added to the embedding feature of x. Therefore, if (6) is successfully optimized, the output of our model will be continuously transformed when adding turbulence to the last five elements of the feature of x. To cover the range of affine parameters σ as listed in Table 1, we sample a value from {-0.4,-0.3, . . . , 0.4} and add it to one or all of the last five elements (corresponding to rotation, horizontal translation, vertical translation, scaling, and shear) of fw(x). The reconstructed images under different configurations are shown in Figure 4. The input examples are shown in the first row, denoted by "Input". The corresponding pure reconstructions are list in row "0". From row "-0.4" to "+0.4", the images are transformed continuously, even though the model only sees examples that are transformed in a range smaller than [-0.3, 0.3]. The last column "Affine" denotes adding noise to all the last five elements, which is an extreme test but still reflects the affine transformation. The result provides qualitative evidence that the encoder mapping fw is equivariant to the affine transformation.

3.2 Clustering

This subsection focuses on evaluating the equivariant features by performing unsupervised clustering task. After training the AEAE, we extract all features {zi = fw(xi)}ni=1 from the embedding layer. Then we perform k-means and spectral clustering (SC) on the extracted features, denoted by AEAE+k-means and AEAE+SC, respectively. We run each method for five times and record the average performance in terms of clustering ACCuracy (ACC) and Normalized Mutual Information (NMI). The baseline methods include conventional shallow clustering and state-of-the-art deep clustering. We report the results by excerpting from the corresponding papers or by running their released code when available.

The results on four datasets are reported in Table 3. Our AEAE+k-means is comparable to deep clustering algorithms and AEAE+SC achieves the best performance. This demonstrates the AEAE can learn more discriminative features. It is possible to further improve the performance by jointly optimizing the clustering model and encoder network, as does existing deep clustering method.

We also compare with other autoencoders including vanilla autoencoder (AE) and denoising autoencoder (DAE). They share the same network structure (number of layers, number of neurons in each layer, activation function, etc.) and optimization setting (learning rate, batch size, maximum epochs, etc.) with our model. The only difference is the objective: AE tries to reconstruct x from x and DAE attempts to reconstruct x from ˆx where ˆx is the corrupted version of x by applying Dropout [Srivastava et al., 2014] with probability 0.2. We report the SC result on the features of different autoencoders in Table 4. Our AEAE outperforms the other autoencoders by a large margin in terms of ACC and NMI. This proves that it is the proposed objective (6), not network structure, that leads to the better discriminative features.

As validated in Section 3.1, the affine factors are encoded in the last five neurons of the embedding layer. So the discriminative features are supposed to be captured by the rest neurons. To validate this point, we extract features from the first ten neurons of the embedding layer and report the SC result (appended by "-10") in Table 4. Our AEAE-10 is comparable to AEAE, but AE-10 and DAE-10 are much worse than AE and DAE, respectively. This well validates the above claim. This also provides guidance in how to select a part of features for the task at hand.

3.3 Classification

We further evaluate the quality of AEAE's features by performing classification task. Each dataset is divided into training and testing set by the ratio of 6 : 1. The features are normalized to zero mean and unit variance. The SVM classifier with RBF kernel is used. We select the best penalty parameter C of the error term from [20 , 21 , . . . , 29 ] by cross validation. We compare with the vanilla autoencoder (AE) and denoising autoencoder (DAE) which share the same network structure and optimization parameters with our AEAE.

We run all methods for five times and report the average and standard deviation of the classification accuracies.

The results are shown in Table 5. First, our AEAE outperforms AE and DAE, which validates that the AEAE can learn discriminative features. Second, by comparing AE-10 (DAE10) with AE (DAE), the classification performance drops dramatically when removing the last five dimensions of the features. Third, AEAE-10 is comparable to AEAE in terms of accuracy, which implies that the discriminative features are mainly encoded by the first ten dimensions. These conclusions are consistent with that drawn in the last subsection.

4 RelatedWork

The proposed AEAE is closely related to the transforming autoencoder [Hinton et al., 2011] which can also reconstruct the affine transformed image from the original image and affine parameters. However, there exists large differences between them as follows. The transforming autoencoder aims to solve the optimization problem of minw,u ∑ni=1 ‖Tσ(xi) - gu(fw(xi) +σ)‖2 where fw(xi) and σ have the same number of dimensions. It can only learn a mapping to approximate the transformation T and can not make fw a valid feature extractor for xi and Tσ(xi). Furthermore, the direct addition of fw(xi) and σ disables the representation fw(xi) to be discriminative, even if it is a valid feature of xi. By contrast, our AEAE can learn discriminative features by forcing the encoder fw to be equivariant to affine transformations. The network structure is also different. The transforming autoencoder only consists of one hidden layer, leading to the lack of capacity. While our multi-layer model has powerful ability to transform data.

Another related model is the capsule network [Sabour et al., 2017] which extends the transforming autoencoder by incorporating dynamic routing mechanism. The capsule network also avoids to learn local transformation invariant features by encoding the instantiation parameters into a group of neurons termed capsule. However, the capsule network requires supervision provided to accomplish this goal. While our AEAE takes advantage of the affine transformation. In other words, our model is trained in an unsupervised manner. Another difference is that our model explicitly encodes the affine factors into specified neurons while the capsule network implicitly learns the instantiation parameters which are uncontrollable.

The rotation equivariant network [Li et al., 2018] focuses on the rotation transformation which is only a specific type of affine transformation. Three new types of convolutional layers (cycle layer, isotonic layer, and decycle layer) were proposed in [Li et al., 2018] to implement rotation equivariance. The rotation equivariant network is trained in supervised manner while our model is unsupervised. It can only learn rotation equivariant features and is hard to adapt to other affine transformations like translation, scaling and shearing. Our AEAE is easy to implement other transformation equivariance if the transformation can be manipulated in a known way.

5 Conclusion

The proposed affine equivariant autoencoder (AEAE) can learn a mapping that is equivariant to the affine transformation in an unsupervised manner. The mapping can serve as a valid feature extractor for images and the resultant features are proved to be very discriminative by performing clustering and classification tasks. The affine factors are disentangled and encoded by the specified neurons of the embedding layer, making the features more interpretable. Future work includes investigating other form of the transformation T ′ σ in the feature space, incorporating the equivariance to more transformations other than the affine transformation, and extending our unsupervised model to semi-supervised scenario.

Acknowledgments

This work was supported by the National Key R & D Program of China 2018YFB1003203 and the National Natural Science Foundation of China (project no. 61773392 and 61672528).