Learning Identity-Preserving Transformations on Data Manifolds

1 INTRODUCTION . A goal of many machine learning models is to accurately identify objects as they undergo natural transformations – a task that humans are adept at . According to the manifold hypothesis , natural variations in high-dimensional data lie on or near a low-dimensional , nonlinear manifold ( Fefferman et al. , 2016 ) . Additionally , the manifolds representing different classes are separated by low density regions ( Rifai et al. , 2011a ) . Natural physical laws govern the possible transformations that objects can undergo and many of the identity-preserving transformations ( e.g. , changes in viewpoint , color , and lighting ) are shared among classes of data . Sharing of transformations between classes enables increased efficiency in defining data variations – a model can represent a limited set of transformations that can describe a majority of variations in many classes . Several machine learning models incorporate specific identity-preserving transformations that are shared among a large number of classes to generalize the performance of their model to unseen data . These include equivariant models that incorporate transformations like translation and rotation into intermediate network layers ( Cohen & Welling , 2016 ; Cohen et al. , 2018 ) and data augmentation techniques that apply known identity-preserving variations to data while training ( Cubuk et al. , 2019 ; Ho et al. , 2019 ; Lim et al. , 2019 ; Sohn et al. , 2020 ; He et al. , 2020 ; Chen et al. , 2020 ) . However , many datasets have natural transformations shared among classes that are not easily prespecified from intuition , making it critical that we develop a model that can learn both 1 ) a representation for these transformations without explicit transformation labels and 2 ) the context for which locations in the data space each transformation is likely to be relevant . Manifold learning strategies estimate the low-dimensional manifold structure of data . A subset of these techniques learn to transform points on the manifold through nonlinear Lie group operators ( Rao & Ruderman , 1999 ; Miao & Rao , 2007 ; Culpepper & Olshausen , 2009 ; Sohl-Dickstein et al. , 2010 ; Cohen & Welling , 2014 ; Hauberg et al. , 2016 ; Connor & Rozell , 2020 ; Connor et al. , 2021 ) . Lie group operators represent infinitesimal transformations which can be applied to data through an exponential mapping to transform points along a manifold , and a manifold can be globally defined by a set of operators that each move in different directions along it ( Hoffman , 1966 ; Dodwell , 1983 ) . A Lie group operator model is well-suited for representing natural data variations because the operators can be learned from the data , applied to data points to transform them beyond their local neighborhoods , and used to estimate geodesic paths . While the Lie group operator models have many benefits , previous approaches demonstrate the two shortcomings noted above . First , to learn Lie group operators that represent a data manifold , pairs of training points are selected which lie within a neighborhood of one another . The training objective encourages efficient paths between these nearby points and the choice of training point pairs influences the types of manifold transformations that are learned . Recent papers incorporating Lie group operators into machine learning models have either used predefined operators that represent known transformations groups ( e.g. , the 3D rotational group SO ( 3 ) ( Falorsi et al. , 2019 ) ) , required transformation labels for selecting point pairs when training ( Connor & Rozell , 2020 ) , or randomly selected pairs of points from the same class ( Connor et al. , 2021 ) . To learn an effective model with datasets having no labeled transformation structure , we require a point pair selection strategy that identifies points that are related through the transformations the model aims to learn . Second , existing Lie group operator models have lacked a method for determining which regions of the manifold are appropriate for each operator , meaning that every operator is equally likely to be used at every point on the manifold . This is a flawed assumption because , while many transformations are shared between classes , there are also data variations that are unique to specific data classes . Additionally , in a dataset with several manifolds ( each representing a different class ) , there is a limited extent to which a transformation can be applied without moving a point onto another manifold . The main contributions of this paper are the development of methods to address the two critical shortcomings of Lie group operator-based manifold models noted above . Specifically , motivated by finding perceptually similar training samples without transformation labels , we first introduce a point pair selection strategy to learn a manifold representation of natural variations shared across multiple data classes without requiring transformations labels . Second , we develop a method that uses a pretrained classifier ( measuring identity-preservation of transformed samples ) to learn the local regions where each operator is likely to be used while preserving the identity of transformed samples . This approach enables us to analyze the local structure of the data manifold in the context of the learned operators and to describe the invariances of the classifier . We demonstrate the efficacy of these strategies in the context of the Manifold Autoencoder ( MAE ) model ( Connor & Rozell , 2020 ) to learn semantically meaningful transformations on MNIST ( LeCun et al. , 1998 ) , Fashion MNIST ( Xiao et al. , 2017 ) , and CelebA ( Liu et al. , 2015 ) . 2 BACKGROUND . Manifold Learning Manifold learning models estimate the low-dimensional structure of highdimensional data by utilizing the property that local neighborhoods on the manifold are approximately linear . Traditional techniques represent the manifold through a low-dimensional embedding of the data points ( Tenenbaum et al. , 2000 ; Roweis & Saul , 2000 ; Belkin & Niyogi , 2003 ; Maaten & Hinton , 2008 ) or through estimates of linear tangent planes that represent local directions of manifold motion ( Dollár et al. , 2007 ; Bengio & Monperrus , 2005 ; Park et al. , 2015 ) . While these traditional manifold learning approaches are useful for understanding low-dimensional data structure , in many cases the input data space is an inefficient representation of the data . For example , data in the pixel space suffers from the curse of dimensionality and can not be smoothly interpolated while maintaining identity ( Bengio et al. , 2005 ) . Many approaches have used neural networks to learn a low-dimensional latent space in which manifold models can be incorporated . The contractive autoencoder ( CAE ) estimates manifold tangents by minimizing the the Jacobian of the encoder network , encouraging invariance of latent vectors to image space perturbations ( Rifai et al. , 2011c ; b ; a ; Kumar et al. , 2017 ) . Several methods estimate geodesic paths in the latent space of a trained variational autoencoder ( VAE ) model ( Arvanitidis et al. , 2018 ; Chen et al. , 2018 ; Shao et al. , 2018 ; Arvanitidis et al. , 2019 ) and use this approach to learn VAEs with priors that are estimated using the Riemannian metrics computed in the latent space ( Arvanitidis et al. , 2021 ; Kalatzis et al. , 2020 ) . Lie Group Operators A Lie group is a group of continuous transformations which also defines a manifold by representing infinitesimal transformations that can be applied to input data ( Hoffman , 1966 ; Dodwell , 1983 ) . Several methods incorporate Lie groups into neural networks to represent data transformations that are identity-preserving within the model ( Cohen & Welling , 2014 ; Cohen et al. , 2018 ; Cosentino et al. , 2021 ) . A prevalent strategy is to learn a dictionary of Lie group operators that are mapped to a specific group element through the matrix exponential expm ( · ) ( Rao & Ruderman , 1999 ; Ham & Lee , 2006 ; Miao & Rao , 2007 ; Culpepper & Olshausen , 2009 ; SohlDickstein et al. , 2010 ; Cohen & Welling , 2014 ; Connor & Rozell , 2020 ; Connor et al. , 2021 ) . In these models , each operator Ψm , called a transport operator , describes a single direction along the manifold and is parameterized by a single coefficient cm . Given an initial data point z , the transport operators define a generative model where transformations can be derived from sampling sparse coefficients cm ∼ Laplace ( 0 , ζ ) : ẑ = expm ( M∑ m=1 Ψmcm ) z + , ( 1 ) where ∼ N ( 0 , σ2 I ) . The manifold autoencoder ( MAE ) incorporates the transport operator model into the latent space of an autoencoder to learn a dictionary of operators that represent the global , nonlinear manifold structure in the latent space ( Connor & Rozell , 2020 ) . This model has been shown to effectively learn reusable operators with transformation supervision , and it will provide the context for demonstrating the effectiveness of the methods developed in this paper . Disentanglement One class of methods that also aim to identify factors of variation in data is disentanglement methods which learn features that each vary one independent characteristic of the data ( Higgins et al. , 2018 ) . In a supervised learning scenario , there are many disentangling techniques that separate data style from content by encouraging similarity between the content features in samples from the same class ( Tenenbaum & Freeman , 2000 ; Reed et al. , 2014 ; Mathieu et al. , 2016 ; Bouchacourt et al. , 2018 ) or encouraging features to map to known class ( Cheung et al. , 2015 ) or transformation labels ( Hinton et al. , 2011 ; Kulkarni et al. , 2015 ; Yang et al. , 2015 ) . There are also techniques that can disentangle latent representations without labels like InfoGAN ( Chen et al. , 2016 ) and β-VAE ( Higgins et al. , 2017 ) . The goal of our work is distinct from disentanglement work because , rather than identifying independently varying factors of variation , we aim to learn non-linear operators that correspond to transformations on the data manifold ( which may not vary independently ) . Our approach is advantageous because it can faithfully represent longer transformation paths , learn variations while maintaining image reconstruction , and change the number of learned variations by increasing or decreasing the number of learned operators . 3 METHODS . The MAE learns a low-dimensional latent representation of the data by defining an encoder function f : X → Z that maps high-dimensional data points x ∈ RD to low-dimensional latent vectors z ∈ Rd and a decoder function g : Z → X that maps the latent vectors back into the data space ( Connor & Rozell , 2020 ) . Transport operators Ψ are incorporated into the latent space to learn manifoldbased transformations . Before learning the transport operators , the autoencoder is pretrained to extract a latent representation of the data using the traditional autoencoder reconstruction objective . After pretraining , the autoencoder weights are fixed and the operators are trained with the following objective , which encourages the learning of transport operators that generate efficient paths between the latent vectors z0 and z1 ( coinciding with f ( x0 ) and f ( x1 ) ) that are nearby on the manifold : LΨ = 1 2 ∥∥∥∥∥z1 − expm ( M∑ m=1 Ψmcm ) z0 ∥∥∥∥∥ 2 2 + γ 2 ∑ m ‖Ψm‖2F + ζ‖c‖1 , ( 2 ) where γ , ζ > 0 . Objective ( 2 ) is minimized via an alternating minimization scheme . Specifically , at each training iteration , points pairs x0 and x1 are selected in the input space and encoded into the latent space z0 and z1 . Then , coefficients c are inferred between the encoded latent vectors ( by minimizing ( 2 ) with respect to c ) to estimate the best path between z0 and z1 . After inference , the coefficients are fixed and a gradient step is taken to minimize ( 2 ) with respect to transport operator weights . Once learned , these operators represent different types of motion that traverse the manifold and they can be combined to generate natural paths on the manifold . After fitting transport operators to the latent space of a fixed autoencoder network , there is a final fine-tuning training phase that updates the network weights and transport operators simultaneously using a joint objective that combines the autoencoder reconstruction loss with LΨ . This fine-tuning step addresses the potential of a mismatch between the data manifold and the learned latent structure by adapting the latent structure to fit the transport operators learned between the selected training point pairs . We empirically find that breaking up training into these three phases increases the stability with which the transport operators can be learned . Given the context of this MAE model , in the following sections we describe our contributions .

Summary: This paper proposes to learn natural transformations in datasets, with manifold auto encoder (MAE), where the underlying identity-preserving transformations are not easily identifiable. By using Lie group operator, this problem reduces to learning paths/motion in the latent space of MAE. The main challenge in training MAE is how to choose a transformation pair, a point and its identity preserving transformed counterpart, in an unsupervised way. This paper introduces the idea of using penultimate layer  of a pretrained classifier on Imagenet to choose such point pairs and learn such MAE more effectively. The proposed method is validated on three datasets.