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: The authors introduce an approach for learning identity preserving transformations from data. First the low-dimensional manifold structure of the data is learned using an autoencoding neural network, then the transformations (the Lie group operators and their coefficients of combination) that map between perceptually similar image pairs are learned. The main contribution of the paper is that the learned operators can transform the data semantically. This is a challenging since 1) semantic transformation labels are not usually available; and 2) not all semantic transformations make sense for all elements of the dataset. The authors address these challenges using two auxiliary networks. The first identifies perceptually similar image pairs in the dataset and the second identifies where particular transformations are likely to be used.