S ELF - SUPERVISED REPRESENTATION LEARNING ON MANIFOLDS
A BSTRACT
We explore the use of a topological manifold, represented as a collection of charts, as the target space of neural network based representation learning tasks.
This is achieved by a simple adjustment to the output of an encoder’s network architecture plus the addition of a maximal mean discrepancy based loss function for regularization.
Most algorithms in representation learning are easily adaptable to our framework and we demonstrate its effectiveness by adjusting SimCLR to have a manifold encoding space.
Our experiments show that we obtain a substantial performance boost over the baseline for low dimensional encodings.
Code for reproducing experiments is provided at https://github.com/ekorman/ neurve .
I NTRODUCTION
Representation learning algorithms typically produce encodings into a Euclidean space.
However, if the manifold hypothesis (the assumption that the data is well-approximated by a manifold) is taken seriously then a Euclidean target space is unnecessarily big for encoding this manifold: by the
Whitney embedding theorem, to embed an n -dimensional manifold into a Euclidean space one may need up to 2 n dimensions for the ambient space.
Additionally, such algorithms seldom regularize the encoding space (i.e. encourage it to look like a prior distribution) to ensure that it is well-behaved.
In this work we develop a framework for using a manifold as a target space of deep representation learning algorithms.
Following prior work (Korman, 2018) which uses manifolds as the latent space of an autoencoder, this is done by having the encoding network output chart embeddings and membership probabilities for an atlas of a manifold.
In other words, instead of learning a single encoder f : X → R d (where X is the input data and R d is the Euclidean encoding space), our technique learns n -such encoders together with a scoring function q : X → [0 , 1] n that determines which encoding output to use for a given input.
Figure 1 shows such an example for X a circle.
For matching the distribution of embeddings to a manifold prior, we introduce a maximal mean discrepancy (MMD) (Gretton et al., 2012) loss for manifolds. [Equation]
Our framework is ﬂexible enough to apply to a variety of representation learning algorithms and in this paper we do experiments with a manifold generalization of SimCLR
(Chen et al., 2020a).
Our results show that, in low dimensions, replacing the Euclidean target space of these algorithms with a manifold gives a signiﬁcant improvement in the learned embeddings.
Besides being of theoretical interest, low dimensional embeddings have the practical beneﬁts of having faster distance computations, requiring less storage, providing visualizations (in the case of dimensions 2 or 3), and avoiding issues with curse of dimensionality for downstream tasks from the embedding (such as clustering).
1. R ELATED WORK
Learning an atlas The idea of learning a manifold as an atlas was discussed in Pitelis et al.
(2013) where they learn a collection of linear charts via a generalization of principal component analysis.
Our earlier work (Korman, 2018), which we build upon, learns an atlas as the latent space of an autoencoder.
In that work the regularization of the latent space is via adversarial training instead of the MMD loss we use in this paper.
Encoding space regularization Regularization of a latent space is typically only done in algorithms that have a decoder/generator, such as variational autoencoders (Kingma & Welling, 2013) and adversarial autoencoders (Makhzani et al., 2015).
The effectiveness of using an MMD loss in particular on such an encoding space is demonstrated in (Tolstikhin et al., 2017) and inspired our
MMD loss function.
The work of Grattarola et al.
(2019) is of similar spirit to ours, as they extend the latent spaces of adversarial autoencoders to more geometrically interesting ones, namely constant curvature Riemannian manifolds.
M ANIFOLDS AS ENCODING SPACES 2. R EPRESENTING A DATASET AS A MANIFOLD
For formally modeling a distribution of data, X , as a manifold, we use the same approach as in our earlier work (Korman, 2018).
Namely we posit the existence of a latent space Z = [0 , 1] d × { 1 , . . . , n } of n , d -dimensional charts with coordinate maps ψ i : [0 , 1] d → X that forms an atlas of a manifold.
We use the uniform distribution as the prior on Z and we let X, Z, J denote the random variables on X , [0 , 1] d , and { 1 , . . . , n } , respectively.
We will denote by 1 y the distribution supported at a single point y .
In our decoder-free setup, we wish to learn:
The inverse mappings of the ψ i , which we denote by φ i : X → [0 , 1] d and which satisfy
[Equation]
The chart membership function q = ( q 1 , . .
. , q n ) :
X → [0 , 1] n deﬁned by
[Equation]
We can then compute the posterior in terms of q and the φ i as [Equation] which gives the prior on Z as
[Equation]
which we wish to be uniform.
An additional desire is that we have an efﬁcient atlas in the sense that any point x should be in as few charts as possible.
Thus while p ( J ) should be uniform, we want the conditional distributions p ( J | X = x ) to have low entropy: if the distributions p ( J | X = x ) are mostly deterministic then encoding x requires us to only keep the coordinates and chart number for the chart with highest probability for x .
In other words, for the representation of x at inference time we take x (cid:55)→ ( φ i ( x ) , i ) ∈
R d × { 1 , . . . , n } , where i = argmax j
[Equation]
, (3) which has just log n -more bits of information than the Euclidean case.
A M AXIMAL M EAN D ISCREPANCY L OSS FOR M ANIFOLDS
For a given embedding task, we propose to parameterize the functions { φ 1 , . .
. , φ n , q } using neural networks and optimize the parameters via gradient descent for a loss function consisting of a task-speciﬁc term (e.g. a contrastive loss function) plus a regularization term that encourages the distribution p ( z, j ) given by (2) to be close to uniform and the discrete distribution q ( x ) in (1) to be close to a deterministic distribution for each x ∈ X .
In the Euclidean case (i.e. when n = 1 ), there are two popular ways for regularizing the latent space to match a prior distribution: using adversarial training (Makhzani et al., 2015;
Tolstikhin et al., 2017) or via an MMD (Gretton et al., 2012) loss (Tolstikhin et al., 2017).
In Korman (2018) we used an adversarial loss for a manifold latent space but in this work we use an MMD loss due to better training stability and less hyperparameters to tune.
If P and Q are two distributions on a common space and k is a reproducing kernel, then the MMD gives a measure of the difference between the distributions, and is deﬁned by
[Equation]
Let p and q be given by (2) and (1), respectively, U Z denote the uniform distribution on Z , and U J denote the uniform distribution on J = { 1 , . . . , n } .
For the loss function encouraging p to be close to U Z we will take an approximation of MMD k Z ( p, U Z ) 2 and for the loss function encouraging q ( x ) to be far from U J , we take − E x MMD k J ( q ( x ) , U J ) .
This will achieve the goal of encouraging p to be uniform and for q ( x ) to be deterministic.
To get a kernel k Z on Z we can start with a kernel k 0 on [0 , 1] d and then deﬁne
[Equation]
where δ ii = 1 , δ ij =
if i (cid:54) =
j .
For k 0 we take, as in Tolstikhin et al.
(2017), the inverse multiquadratics kernel but pulled back via the sigmoid function 1 : k 0 ( x, y ) =
d/ 6 d/ 6+ | σ − 1 ( x ) − σ − 1 ( y ) | 2 .
We approximate
MMD k Z ( p, U Z ) 2 using the U-statistic estimator in Gretton et al.
(2012), adjusted to our manifold setting.
Explicitly:
Proposition 1.
Let p be the distribution on Z deﬁned by (2) and let U Z denote the uniform distribution on Z .
Given a random sample { x 1 , . . .
, x N } of X and a random sample { w 1 , . .
. , w N } drawn uniformly from [0 , 1] d , an estimator for MMD k Z ( p, U Z ) 2 is:
[Equation]
See section A.3 for the proof.
For E x MMD k J ( q ( x ) , U J ) we take the kernel k J : { 1 , . . . , n } × { 1 , . . . , n } → R , ( i, j ) (cid:55)→ δ ij and deﬁne
[Equation]
The total regularization loss is then
[Equation]
for some hyperparameters λ 1 , λ 2 ∈ [0 , ∞ ) .
2. S UMMARY OF OUR FRAMEWORK
We now describe our general technique of turning a deep representation learning algorithm to one that has a manifold encoding space.
The typical setup for such an algorithm, A , is to learn an encoder f mapping the dataset X to the space R d via optimizing some loss function (cid:96) A deﬁned on embeddings of a mini-batch { x 1 , . . .
, x N } ⊂ X .
To adjust the algorithm to one that has a manifold as the encoding space, we see from the discussion in 2.1 that f should be replaced by a collection of maps φ 1 , . .
. , φ n , q where φ i : X → [0 , 1] d is the i th coordinate map and q : X → [0 , 1] n is the chart membership function.
In practice, for the functions φ 1 , . . . , φ n , q , we take a backbone network and attach n + 1 linear heads followed by a sigmoid activation on the ﬁrst n (which have output in R d ) and a softmax activation on the head deﬁning q (which has output in R n ).
In many contrastive representation learning algorithms (such as SimCLR (Chen et al., 2020a), MoCo v2 (Chen et al., 2020b), and BYOL (Grill et al., 2020)), the loss (cid:96) A is a function of an auxiliary projection head h : R d → R ˜ d applied to the embedding vectors.
In these cases, for the manifold version we use n -many projection heads (one for each coordinate chart), and then produce a single projection vector for every data point by taking a sum of these individual projection vectors weighted by q .
We supplement the resulting loss (cid:96) MA with the regularization loss (6).
At inference and evaluation time, we use the compressed representation (3).
This ensures that the representation is indeed d -dimensional and gives a fair evaluation comparison to A , which embeds into a d -dimensional Euclidean space.
For example, if q i ( x ) were the uniform distribution then using the full-representation ( φ 1 ( x ) , . . . , φ n ( x )) instead of the compressed one would essentially be “cheating” into an nd -dimensional Euclidean representation.
We summarize this procedure in Table 1.
We also note that the case of d = 2 yields a powerful data visualization by plotting the input data at their embedding coordinates for their most probable chart.
We show such visualizations in the appendix A.2.
TYPICAL SETUP
[Table]
MS IM CLR
We recall that SimCLR trains a neural network encoder f : X → R d using a projection head h : R d → R ˜ d .
A mini-batch is formed by choosing N -images and augmenting in two different ways, producing examples { x 1 , . . .
, x 2 N } .
The loss function over this batch is a function of the projection head output of the embeddings of these images:
[Equation]
where c is a contrastive loss.
We follow our meta-procedure from the previous section to adjust
SimCLR (Chen et al., 2020a) to have a manifold as encoding space (with resulting algorithm denoted by MSimCLR ) and run experiments on MNIST (LeCun et al., 2010), FashionMNIST (Xiao et al., 2017), and CIFAR10 (Krizhevsky et al., 2009) for n ∈ { 1 , 4 , 16 , 32 } and d ∈ { 2 , 4 , 8 } .
In Chen et al.
(2020a) representations of SimCLR are evaluated based on the accuracy of a linear classiﬁer trained on top of the representation.
In our case, since we have a collection of charts, we evaluate our manifold representation by putting a linear classiﬁer on each chart.
We report the mean and standard deviation of the accuracy of our models on the hold out test sets in Table 2, from which we see that our method provides a signiﬁcant performance boost over vanilla SimCLR, especially in dimensions two and four.
Section A.1 outlines the details of our experiments, including hyperparameter selection
[Table]
[Table]  C ONCLUSION We presented a method for adjusting representation learning algorithms to learn an atlas of a manifold instead of mapping into a Euclidean space. This allows, for a given encoding dimension, more interesting geometries of the embedding space.
Our experiments with SimCLR show that for low encoding dimensions our approach gives much more powerful representations over the baseline.
We hope that this work leads to further research into atlas-based manifold learning techniques
A A PPENDIX A. E XPERIMENTS SETUP
We ran our experiments using PyTorch Paszke et al.
In all cases we use a batch size ofand the Adam optimizer with learning rate 10 − 4 and β 1 = 0 .
9 , β 2 = 0 .
999 (the PyTorch defaults).
For MNIST and FashionMNIST we use a ResNet
He et al.
(2016) backbone and train forepochs while for CIFAR10 we use a ResNet50 backbone and train for 1,000 epochs.
We use the same data augmentation used in the CIFAR10 experiments in the original work Chen et al.
(2020a): color jitter with strength 0.5 (and a probability of 0.8 of applying), random grayscale with probability 0.2, and a random resized crop.
Each of the datasets comes with a standard train/test split.
For hyperparameter selection of λ 1 , λ 2 , and τ (the temperature used in the contrastive loss function) we do a grid search (with d = 2 , n = 16 ) over λ 1 ∈ { 0 . 1 , 1 , 5 , 10 , 20 } , λ 2 ∈ { 0 . 05 , 0 . 1 , 0 . 2 } , and τ ∈ { 0 . 1 , 0 .
5 , 1 } by training on a random selection of 80% of the training data and then computing the piece-wise linear evaluation on the remaining 20%.
For MNIST and FashionMNIST we chose the hyperparameters that give the best average rank across the two datasets.
This yielded λ 1 = 20 , λ 2 = 0 .
1 , and τ = 1 for
MNIST and FashionMNIST and λ 1 = 20 , λ 2 = 0 . 1 , and τ = 0 .
5 for CIFAR10.
Using these parameters we train on the entire train dataset across d ∈ { 2 , 4 , 8 } (we noticed for higher d there is not much improvement over the baseline) and n ∈ { 1 , 4 , 16 , 32 } .
We also train baseline SimCLR for comparison (but were unable to get convergence in dimension two for CIFAR10).
We repeat each training conﬁguration three times and report the mean ± std of accuracy on the holdout test set in Table 2.
A. V ISUALIZATIONS
In this section we display visualizations for each dataset using a model with d = 2 , n = 16 .
Each image corresponds to a chart and every image in the test set gets plotted at its ( x, y ) coordinate in the chart with highest probability.
The images are best viewed by zooming in.
A.2. CIFAR
A.2. F ASHION MNIIST
We note that in this case only 15 charts are displayed since the network did not end up assigning any points to one of them.
A.2.
MNIST
A. P ROOF OF PROPOSITION
Proof.
We breakdown the three terms deﬁning MMD (4).
For the ﬁrst term, we have:
[Equation]
where the ﬁrst equality comes from (applying twice) the fact that
E z
[Equation]
and the approximation in the last line uses the U-statistic estimate
[Equation]
which holds for any distribution Q and random sample { y 1 , .
y N } drawn from Q .
For the second term we have, letting U [0 , 1] d denote the uniform distribution on [0 , 1] d ,
[Equation]
where the ﬁrst equality again follows from (7) and the approximation in the last line uses the standard expected value estimator applied to each factor.
Finally for the third term we again use the estimate (8) to obtain
[Equation]
Combining these gives the right hand side of (5).