Properties from mechanisms: an equivariance perspective on identifiable representation learning

1 INTRODUCTION . Modern unsupervised representation learning techniques can generate images of our world with intricate detail ( e.g . Karras et al. , 2019 ; Song et al. , 2020 ; Razavi et al. , 2019 ) , and yet , the latent representations from which these images are generated remain entangled and challenging to interpret ( Schölkopf et al. , 2021 ; Locatello et al. , 2019 ) . At the same time , the success of pre-trained transformers ( Devlin et al. , 2018 ; Brown et al. , 2020 ) shows that advances in our ability to disentangle the underlying generative factors can lead to dramatic improvements in the sample complexity of downstream supervised tasks ( Bengio & LeCun , 2007 ; Bengio et al. , 2013 ) . In order to consistently replicate this success , we need methods that can reliably invert the data generating process . This is a challenging task because unsupervised representation learning with independent and identically distributed ( IID ) data is hopelessly unidentified : even in the nice case in which observations are generated via some p ( x|z ) , there exists an infinitely large class of possible latent distributions p̃ ( z|x ) that are consistent with the observed marginal distribution , p ( x ) , and only one of them corresponds to the true latent distribution p ( z|x ) ( Locatello et al. , 2019 ; Khemakhem et al. , 2020a ) . If we want to build systems which are able to provably identify1 the true latent variables z that generated our observed data x = g ( z ) for some observation model g , then we need to exploit structural assumptions that constrain this set of possible distributions . Most of the prior work that can provide such guarantees was developed in the independent component analysis ( ICA ) literature . The key ICA assumption is that the latent factors are ( conditionally ) independent and non-Gaussian . Then , if the observation model , g : Z → X is linear and invertible , z = g−1 ( x ) is identified ( Comon , 1994 ) , and for nonlinear g there are a number of recent approaches that leverage non-stationarity in the distributions over z to identify g−1 ( Hyvarinen & Morioka , 2016 ; 2017 ; Hyvarinen et al. , 2019 ; 1A problem is identified if there exists a unique solution in the infinite data limit and no constraints on model capacity . Khemakhem et al. , 2020a ) . These results give a tantalizing demonstration that representation learning with identification guarantees is possible , but the requirement that the latent factors are statistically ( conditionally ) independent is limiting2 ( Higgins et al. , 2018 ; Schölkopf et al. , 2021 ) . In this paper we take a very different approach . We study how the mechanisms that govern an environment ’ s evolution constrain the set of possible latent representations that are consistent with the data . As a simple concrete example , consider the bouncing ball environment shown in Figure 1 . The latent state can be completely described by a vector , z , containing the position , velocity and acceleration of the balls3 , and given this latent state , the images , x , shown in Figure 1 , are produced via some rendering engine g : Z → X . Our task is to leverage sequences of observations x1 , . . . , xT and knowledge of m to recover z1 , . . . , zT . If we can show that this task has a unique solution that is consistent with the observations and mechanisms , then the problem is identified . Our main result is that when we know the true mechanism , m , the system is identified up to any equivariances of the mechanism ; or equivalently , in Section 2.2 we prove that we can identify z up to any invertible transformation a : Z → Z that commutes with m , such that the composition m◦a = a◦m . For example , in the bouncing balls environment shown in Figure 1 , the laws of physics are equivariant with respect to your choice of units of measurement—changing from representing z in meters to inches leaves the output of the mechanism unchanged up to a corresponding unit change— and hence we can only hope to identify z up to some scaling function a ( z ) which corresponds to an arbitrary choice of units of measurement . Interestingly , when the environment evolves according to multiple known mechanisms , the sources of non-identifiability are even further constrained : such a system is identified up to equivariances that are shared by all n mechanisms . Perfectly knowing the true mechanisms and when they are applied is unlikely , but in Section 2.3 we show that we can relax that assumption to a setting where we instead know a hypothesis class of possible mechanisms that could have been applied . This weaker assumption leads to an additional 2As a simple example of dependence between latent variables , assume the bouncing balls shown in Figure 1 have different masses indicated by their colors . If the initial conditions were such that all the balls have the same momentum , then mass and velocity will be inversely correlated . 3A complete generative model of these images would also need to track the colors and shapes of the elements ; we will return to this issue in the discussion of the results . source of non-identifiability : the mechanisms in our hypothesis class can imitate each other if there exists an invertible transformation a such that for any two mechanisms m1 and m2 in our hypothesis class , m2 = a−1 ◦m1 ◦ a . For example , if we are in a setting where our hypothesis class includes both a product mechanism , m1 ( z ) = ∏ i zi , and a sum mechanism , m2 ( z ) = ∑ i zi , then m1 can imitate m2 if a ( z ) = exp ( z ) , since ∑ i zi = log ( ∏ i exp ( zi ) ) . As before , this result is complete , in the sense that these two sources of non-identifiability—equivariance and imitation—are the only sources of non-identifiability in such systems . This gives us a natural way of thinking about the way knowledge of deterministic mechanisms constrains a representation learning task : with complete knowledge , the only source of non-identifiability is any equivariances inherent in the mechanism , but by allowing a hypothesis class of possible mechanisms , we introduce potential non-identifiability via imitation . That said , the relationship between the size of the hypothesis class and the size of the set of a ’ s that commute via imitation is not necessarily monotonic : for environments governed by multiple mechanisms , a larger hypothesis class can lead to fewer imitations . Section 3 , shows that we can derive analogous results for stochastic mechanisms , m ( z , U ) , where the mechanism defines a conditional distribution p ( zt+1|zt ) . This generalization gives us a way of comparing our mechanism-based perspective with existing identifiability results . We demonstrate this in Section 4 by showing that it is possible to view the distributional assumptions made in Klindt et al . ( 2020 ) as a particular choice of mechanism , and by doing so , we can leverage our theory to give alternative proofs of these results . This strategy required weaker distributional assumptions , thereby generalizing their result . Finally , we give a mechanism-based perspective on the related work in Section 5 and Section 6 concludes with a discussion of the open problems that need to be addressed in order to reliably leverage this approach in practice . 2 MECHANISM BASED IDENTIFICATION . 2.1 DATA GENERATION PROCESS . The state of the system at time t ∈ { 1 , · · · , T } is given by zt ∈ Z ⊆ Rd . At each time t , we observe g ( zt ) = xt ∈ X ⊆ Rn , which is some transformation of the latent zt . We can think of g : Z → X as a function that transforms the ( typically low dimensional ) state variables to the ( typically high dimensional ) observed variables ; for example , in the bouncing ball environment described in the introduction , g is the rendering engine that produces the images shown in Figure 1 . We assume g is injective 4 with respect to Rn—i.e . g ( z1 ) = g ( z2 ) implies z1 = z2 ; or equivalently , any change to the underlying state , z , is reflected in some pixel-level change to the observation x—and we make g bijective by restricting its inverse g−1 to any x on the data manifold which we denote X ( i.e . X is the image of g ) . The state transition from time t to t+ 1 is governed by a mechanism mt : Z → Z . There may be multiple mechanisms in a given environment . For example in Figure 1 the transition from z1 to z2 does not involve any collisions so the state evolves according to Newton ’ s first law of motion ( Newton , 1687 ) ; the transition from z2 to z3 involves a collision between two of the balls that is described by Newton ’ s third law . Together mt and g describe the data-generation from time t to t+ 1 as follows , xt ← g ( zt ) , zt+1 ← mt ( zt ) . ( 1 ) If we fix the initial conditions , z0 , this data generation process is deterministic . In Section 2.2 , we provide results when the underlying mechanism is known and in Section 2.3 , we we extend those results to the case when the mechanisms are not known . In Section 3 , we extend our results to stochastic mechanisms , m ( Zt , U ) that take samples from some distribution U ∼ Uniform ( 0 , 1 ) as input . These stochastic mechanisms can represent any conditional distribution P ( Zt+1|Zt ) over states.5 2.2 IDENTIFYING ENCODERS WHEN THE UNDERLYING MECHANISM IS KNOWN . We begin in the simplest version of the system described by equation ( 1 ) : assume that at each time t the same mechanism m : Z → Z is used , and we know this mechanism . From these assumptions we 4for ball collision example , if Z equals position and velocity , we can achieve injectivity by frame stacking . 5For example , if Zt+1|Zt ∼ N ( Zt , σ ) then m ( z , U ) = z + σΦ−1 ( U ) where Φ is the cumulative density of a standard Gaussian . Strictly , representing a joint ( Zt+1 , Zt ) as ( m ( Zt , U ) , Zt ) requires a restriction to nice spaces . See Austin ( 2015 , Lemma 3.1 ) . can derive an identity that describes how the observations xt and xt+1 are related , zt+1 = m ( zt ) , g −1 ( xt+1 ) = m ◦ g−1 ( xt ) , xt+1 = g ◦m ◦ g−1 ( xt ) . ( 2 ) It may be helpful to think of this identity as describing an autoencoder where we require that the encoder g−1 ( xt ) inverts the data generating process to produce some latent zt from xt , and that the decoder has to reproduce xt+1 from m ( zt ) ; i.e . the representation transformed by the mechanism . Importantly , this identity describes the true data generating process , rather than some model of it . Our hypothesis class over the possible encoder / decoder functions is the set , G , of all bijective functions from Z → X . By assuming that bijectivity we are essentially assuming that the reconstuction task is solved:6 G is the set of all autoencoders that perfectly reproduce the data , such that for any x on the data manifold X , and any encoder / decoder pair ( g̃−1 , g̃ ) with g̃ ∈ G , their composition is the identity function , x = g̃ ◦ g̃−1 ◦ x . Because we assumed that the true g is bijective , we know that it is in our search space , G. We can constrain this set using our knowledge of the mechanism by only considering solutions that also satisfy the identity given in equation 2 , such that for every pair of observations ( xt , xt+1 ) , an analogous identity holds , xt+1 = g̃ ◦m ◦ g̃−1 ( xt ) ( 3 ) Now suppose that we have access to observations xt from the entire set X , then the above identities have to hold for all xt , ∈ X with corresponding xt+1 from equation 2 , so we can conclude that the following functions are equal , g ◦m ◦ g−1 = g̃ ◦m ◦ g̃−1 ( 4 ) This relationship will hold for any of the possible decoders g̃ ( and corresponding encoders g̃−1 ) that are observationally equivalent given our assumptions . We denote the set of all such decoders , Gid = { g̃ | g̃ ∈ G , g ◦m ◦ g−1 = g̃ ◦m ◦ g̃−1 } . If Gid = { g } then we have shown that the problem is exactly identified , which means that if we manage to find a g̃ that satisfies equation 2 , then g̃ = g ; but if Gid , also includes other functions g̃ 6= g , then these functions are the sources of non-identifiability . Equivariant mechanisms To see an example of a setting where the problem is not exactly identified , consider a mechanism which is equivariant with respect to some bijective transformation a : Z → Z . A mechanism m is said to be equivariant w.r.t a if a ◦ m = m ◦ a . For example , a mechanism derived from Newtonian mechanics may be equivariant with respect to your choice of units of measurement , such that m ( cz ) = cm ( z ) for some scaling constant c. Similarly , if a mechanism transforms sets of items , any permutation of z would lead to a corresponding permutation of the mechanism ’ s output . If we have a mechanism that is equivariant with respect to some transformation a ( where a is not identity map ) , that implies that there exists a function g̃ 6= g in Gid , so the problem is not exactly identified . We can see this as follows , g ◦m ◦ g−1 = g ◦ a−1 ◦ a ◦m ◦ g−1 = g ◦ a−1 ◦m ◦ a ◦ g−1 = g̃ ◦m ◦ g̃−1 where the first equality uses the fact that a−1 ◦ a is the identity function , the second applies the definition of equivariance , and the final equality defines g̃ : = g ◦ a−1 and g̃−1 : = a ◦ g−1 . This shows that if the mechanism is equivariant , an encoder , g̃−1 , can output a transformed z̃ = a ( z ) , and the decoder , g̃ , inverts this transformation before producing its output , thereby leaving the observed variables , x , unchanged . If we denote the set of all equivariances of the mechanism E = { a | a is a bijection , a ◦ m = m ◦ a } , then we can define the set of all such sources of non-identification as , Geq = { g̃ | g̃ = g ◦ a−1 , a ∈ E } . This is a natural source of non-identification : given that we are relying on the mechanism to constrain the encoder g̃−1 , it is unsurprising we can not prevent transformations that are not affected by the mechanism . The more interesting observation , which we will show below , is that this set is the only source of non-identification when the mechanism is known , and hence we recover the true z up to equivariances in the mechanism . To state this theorem , we first define the notion of identifiability up to an equivalence class defined by a family of bijections A , where a ∈ A is a map a : Z → Z . Definition 1 . Identifiability up to A . If the learned encoder g̃−1 and the true encoder g−1 are related by some bijection a ∈ A , such that g̃−1 = a ◦ g−1 ( or equivalently g̃ = g ◦ a−1 ) , then g̃−1 is said to learn g−1 up to bijections in A . We denote this g̃−1 ∼A g−1 . 6This is obviously a strong assumption—learning autoencoders that perfectly reconstructed the data is not at all easy—but it focuses the discussion on the identification issues that remain after reconstruction is solved . Suppose , for example , A is a family of a permutations . Identifiability up to A implies that the true latent variables will be recovered , but that they would not necessarily be ordered in the same way as they were in the original data generation process . In this setting where the mechanism , m , is known , the following theorem shows that A is just E , the set of equivariances of m. Theorem 1 . If the data generation process follows equation 2 , the encoders that solve the observation identity in equation 4 identify true encoder g−1 up to equivariances of m ( g̃−1 ∼E g−1 ) . To prove the above theorem , we need to establish that the set of solutions to the identity in equation 4 and the set of maps derived from equivariances are equal Gid = Geq . The proof is given in Section A.1 of the appendix . From Theorem 1 , we can derive a number of observations . First , notice that if we have a standard autoencoder7 then the mechanism is the identity map , m ( z ) = z , and its set of equivariances E is any invertible function a , and hence Theorem 1 shows that the encoder is essentially unconstrained . However , if the mechanism is any non-trivial function , m ( z′ ) 6= z′ for some z′ , then the space of possible encoders is significantly reduced to just those invertible transformations that commute with m. If the system involves multiple known mechanisms , { m1 , . . . , mT } , where at each time zt+1 = mt ( zt ) , then the encoder is even further constrained . Define the set of all the mechanisms that are used at least once in the evolution of the system asM∗ = ∪Tt=1 { mt } . Suppose E i denotes the equivariances of mi ∈ M∗ . Define the equivariances that are shared across all the mechanisms as E∗ = ∩iE i ; Corollary 1 . If the data generation process follows equation 1 , then the encoders that satisfy observation identity in equation 4 for all m ∈M∗ identify true encoder g−1 up to the equivariances shared across all the mechanisms inM∗ , E∗ ( g̃−1 ∼E∗ g−1 ) . The proof of the above claim is in Section A.3 . This corollary implies a blessing that comes with more complex environments : if an object is transformed by multiple mechanisms which are diverse ( in the sense that they share few equivariances ) , then it becomes easier to identify . Given access to inputs and outputs of a mechanism , if we can not tell apart whether some transformation was applied to the input or the output , then the mechanism is equivariant with respect to the transformation . Together Theorem 1 and Corollary 1 , state that we can learn to invert the data generation process but we can not distinguish latents that were transformed by equivariances shared across the mechanisms .

Summary: **Scope.** The paper studies unsupervised/weakly-supervised representation learning from an identifiability perspective. The considered problem setting is that observations $x_t$ are generated from some bijective function $g$ applied to latent codes $z_t$, $x_t=g(z_t)$; and that we have access to successive observations $..., x_t, x_{t+1}, ...$ which are related through shared *mechanisms* $m_t$ dictating the dynamics in the latent space, $z_{t+1}=m_t(z_t)$. The paper then provides a (purely) theoretical analysis regarding the identifiability of $g$ given (partial) knowledge of the mechanisms $m_t$. In particular, the identifiability class of $g$ is derived in terms of *equivariances* of the mechanisms.  **Results.** For the case of a single known mechanism $m$, it is shown that $g$ can be recovered only up to equivariances of $m$, i.e., up to bijections $a$ that commute with $m$, meaning that any learnt function $g'$ relates to the true $g$ via $g'=g\circ a$ for some $a$. This result is strengthened when multiple distinct known mechanisms govern the dynamics, where identifiability is obtained up to shared equivariances, i.e., such $a$ that commute with all $m_t$. When the relevant mechanisms are not known, non-identifiability issues increase with the size of the hypothesis class which is attributed to the presence of so-called *imitators*, transformations that give rise to the same behaviour as some mechanism when combined with another mechanism. Finally, a result is presented for the case of stochastic mechanisms which give rise to conditional transition distributions $p(z_{t+1}|z_t)$: here, equivariance and imitation are relaxed to hold only in distribution. The insights are used to re-derive a known identifiability result by Klindt et al. (2020).