Disentanglement via Mechanism Sparsity Regularization: A New Principle for Nonlinear ICA
: Bernhard Sch ¨ olkopf, Caroline Uhler and Kun Zhang
Abstract
This work introduces a novel principle we call disentanglement via mechanism sparsity regularization , which can be applied when the latent factors of interest depend sparsely on past latent factors and/or observed auxiliary variables.
We propose a representation learning method that induces disentanglement by simultaneously learning the latent factors and the sparse causal graphical model that relates them.
We develop a rigorous identiﬁability theory, building on recent nonlinear independent component analysis (ICA) results, that formalizes this principle and shows how the latent variables can be recovered up to permutation if one regularizes the latent mechanisms to be sparse and if some graph connectivity criterion is satisﬁed by the data generating process.
As a special case of our framework, we show how one can leverage unknown-target interventions on the latent factors to disentangle them, thereby drawing further connections between ICA and causality.
We propose a VAE-based method in which the latent mechanisms are learned and regularized via binary masks, and validate our theory by showing it learns disentangled representations in simulations.
1. Introduction
It has been proposed that causal reasoning will be central to move modern machine learning algorithms beyond their current shortcomings, such as their lack of robustness , transferability and interpretability (Pearl, 2019; Sch¨olkopf, 2019; Sch¨olkopf et al., 2021; Goyal and Bengio, 2021).
However, it is still unclear how to reconcile the causal graphical model (CGM) formalism (Pearl, 2009;
Peters et al., 2017), which operates on semantically meaningful high-level variables, with deep neural networks (Goodfellow et al., 2016), which excel on unstructured, low-level, high-dimensional observations, e.g. images.
One way forward would be a two-step approach in which we ﬁrst disentangle the high-level variables from low-level observations (Bengio et al., 2013; Locatello et al., 2019), then learn a CGM that relates them.
Instead, this work proposes a method to do both steps simultaneously , and provides a rigorous theory that shows how doing so can induce disentanglement when the CGM is regularized to be sparse.
Our contribution is based on recent theoretical results in the nonlinear ICA literature (Hyv¨arinen et al., 2019; Khemakhem et al., 2020a,b) that assume the data is explained by unobserved and meaningful latent variables, or factors, Z , which are transformed by a decoder , or mixing function , f , to produce the observation X .
The problem of disentanglement can then be formulated as recovering, or reconstructing, the latent variables from the observation.
This problem is plagued by the difﬁcult question of identiﬁability .
Indeed, Hyv¨arinen and Pajunen (1999) showed that this task is impossible with general nonlinear mixing under the standard assumption of independent latent factors.
Nevertheless, recent theoretical developments have shown identiﬁability of the latent factors is possible in the nonlinear setting, assuming the latent variables are conditionally independent given an observed auxiliary variable A (Hyv¨arinen et al., 2019; Khemakhem et al., 2020a,b).
This auxiliary variable can be, for instance, a time or an environment index, an action in an interactive environment, or even a previous observation if the data has temporal structure, as long as its effect on the latent factors is “sufﬁciently strong”.
The present paper introduces mechanism sparsity regularization as a new path to disentanglement.
By building on the recent theoretical developments in ICA, we show that if the high-level variables have a sparse temporal structure and/or an action is observed and affects the high-level variables sparsely , then the latent variables can be recovered by regularizing the inferred graphical model to have sparse dependencies (Thm. 5).
In estimating the latent variables, the presented methodology estimates the causal graph describing them and their relation to the action
A (when available).
A very similar disentanglement method based on graph sparsity was proposed independently by Volodin (2021), but this concurrent work does not analyze identiﬁability formally (Sec. 3).
In contrast, our theory provides precise conditions, e.g. on the ground-truth graph, to ensure identiﬁability, thus extending the domain of known cases where latent variables can be recovered.
The hypothesis that high-level concepts can be described by a sparse dependency graph has been described and leveraged for out-of-distribution generalization originally by Bengio (2019) and Goyal et al.
(2021b), which were early sources of inspiration for this work.
To the best of our knowledge, our theory is the ﬁrst to show formally that this inductive bias can sometimes be enough to recover the latent factors.
As a special case, it also shows formally how unknown-target interventions on the latent factors can be leveraged to disentangle them (Sec. 2.5), which is closely related to the sparse mechanism shift hypothesis described by Sch¨olkopf et al.
Fig.
1 shows a minimal motivating example in which our approach could be used to extract the high-level variables (such as the x -position of the three objects) and learn their dynamics (how the objects move and affect one another) from a time series of images, X t , and agent actions, A t .
Thm. 5 shows how the sparse dependencies between the objects can be leveraged to estimate the latent variables as well as the graph describing their dynamics.
The learned CGM could subsequently be used to simulate interventions on semantic variables (Pearl, 2009; Peters et al., 2017), such as changing the torque of the robot or the weight of the ball.
Interventions allow an agent to imagine situations it has never seen before, which would not be possible without a disentangled representation (Sch¨olkopf, 2019).
Moreover, disentanglement could be useful for interpretability by allowing for the extraction of a causal graph of the agent actions (Pearl, 2019).
Contributions:
A new principle to achieve disentanglement based on mechanism sparsity regularization motivated by a rigorous and novel identiﬁability theory (Thm. 5).
A formal connection to the sparse mechanism shift hypothesis introduced by Sch¨olkopf et al.
(2021) via the notion of unknown-target interventions on the latent factors (Sec. 2.5).
An estimation procedure which relies on variational autoencoders (VAEs) (Kingma and Welling, 2014) and learned causal mechanisms regularized for sparsity via binary masks. 4.
An illustration of our theoretical predictions being satisﬁed in practice by our estimation procedure on synthetic datasets.
The paper is structured as follows.
Sec.
2.1 introduces the model under consideration.
Sec. 2.
deﬁnes the notions of linear and permutation equivalence between representations.
Sec.
2.3 provides conditions to identify the model up to linear equivalence.
Sec.
2.4 shows how mechanism sparsity regularization can induce permutation-identiﬁability, i.e. disentanglement.
Sec.
2.6 proposes a
VAE-based approach to model estimation, as well as a practical way to induce sparsity based on binary masks.
Sec.
3 situates this paper in the related literature.
Sec. 4 illustrates the proposed identiﬁability theory and learning methods on synthetic data.
2. Disentanglement via Mechanism Sparsity Regularization 2.1. An identiﬁable latent causal model
We now specify the setting under consideration.
Assume we observe the realization of a sequence of d x -dimensional random vectors { X t } Tt =1 and a sequence of d a -dimensional auxiliary vectors { A t } T − 1 t =0 .
The coordinates of A t are either discrete or continuous and can potentially represent, for example, an action taken by an agent, or the index of the environment the corresponding observation was taken from.
Going forward, we will refer to A t as the action vector.
The observations { X t } are assumed to be explained by a sequence of hidden d z -dimensional continuous random vectors { Z t } Tt =1 via the equation X t
= f ( Z t )
+ N t where N t ∼ N (0 , σ 2 I ) are mutually independent across time and independent of all Z t and A t .
Throughout, we assume d z ≤ d x and that f : Z → X is a diffeomorphism 1 where Z is the support of Z t for all t , and X := f ( Z ) , i.e. the image of
Z under f .
App.
A.5 discusses the implications of the diffeomorphism assumption.
We suppose that each factor Z ti contains semantic information about the observation, e.g. for high-dimensional images, the coordinates Z ti might be the position of an object, its color, or its orientation in space.
We denote Z ≤ t :=
[ Z 1 · · · Z t ] ∈ R d z × t and analogously for Z <t and other random vectors.
Similar to previous work on nonlinear ICA (Hyv¨arinen et al., 2019; Khemakhem et al., 2020a), we assume the variables Z ti are mutually independent given Z <t and A <t
[Equation]
Our theory holds for a rich family of conditional densities p ( z ti | z <t , a <t ) called the exponential family (Wainwright and Jordan, 2008), which has the following form:
[Equation]
Well-known distributions which belong to this family include the Gaussian and beta distribution.
In the Gaussian case, the sufﬁcient statistic is T i ( z ) :
=
( z, z 2 ) and the base measure is h
i ( z ) :
= 1 √ 2 π
[Equation]
The function λ i ( G zi (cid:12) z <t , G ai (cid:12) a <t ) outputs the natural parameter vector for the conditional distribution and can be itself parametrized, for instance, by a multi-layer perceptron (MLP) or a recurrent neural network (RNN).
We will refer to the functions
λ i as the mechanisms or the transition functions .
In the Gaussian case, the natural parameter is two-dimensional and is related to the usual parameters µ and σ 2 via the equation ( λ 1 , λ 2 ) =
( µσ 2 , − 1 σ 2 ) .
We will denote by k the dimensionality of the natural parameter and that of the sufﬁcient statistic (which are equal).
Thus, k = 2 in the Gaussian case.
It will be useful to keep this example in mind throughout the paper.
The remaining term ψ i ( z <t , a <t ) acts as a normalization constant.
The binary vectors G zi ∈ { 0 , 1 } d z and G ai ∈ { 0 , 1 } d a act as masks selecting the direct parents of z ti .
The Hadamard product (cid:12) is applied element-wise and broadcasted along the time dimension.
2 We deﬁne
[Equation]
as well as G :=
[ G z G a ] which is the adjacency matrix of the causal graph.
Indeed, (1) & (2) describe a CGM over the unobserved variables Z ≤ T conditioned on the auxiliary variables
A <T .
We deﬁne λ ( z <t , a <t ) ∈
R kd z to be the concatenation of all λ i ( G zi (cid:12) z <t , G ai (cid:12) a <t ) and similarly for T ( z t ) ∈
R kd z .
Note that λ
( z <t , a <t ) depends on G , implicitly to simplify the notation.
The learnable parameters are θ :=
( f , λ , G ) , which induce a conditional probability distribution
P X ≤ T | a ; θ over X ≤ T , given A <T = a .
Let A ⊂ R d
a be the set of possible values
A t can take.
We assume p ( a <T ) has probability mass over all A T .
This could arise, for instance, when A t is sampled from a policy π ( a t | z t ) distribution with probability mass everywhere in A .
A motivating example.
Fig.
1 represents a minimal example where our theory applies.
The environment consists of three objects: a tree, a robot and a ball with x -positions T t , R t and B t , respectively.
Together, they form the vector Z t of high-level latent variables, i.e. Z t =
( T t , R t , B t ) .
A remote controls the direction in which the wheels of the robot turn.
The vector
A t records these actions, which might be taken by a human or an artiﬁcial agent trained to accomplish some goal.
The only observations are the actions A t and the images X t representing the scene which is given by X t
= f ( Z t )
+ N t .
The dynamics of the environment is governed by the transition function λ .
Assuming a Gaussian model with ﬁxed variance for the latent factors Z ≤ T , λ would output the expected position of every object given their previous positions.
Plausible connectivity graphs G z and G a are given in Fig.
1 showing how the latent factors are related, and how the controller affects them.
For every object, its position at time step t depends on its position at t − 1 .
The position of the tree, T t , is not affected by anything, since neither the robot nor the ball can change its position.
The robot, R t , changes its position based on both the action, A t − 1 and the position of the tree, T t −(in case of collision).
The ball position, B t , is affected by both the robot, which can kick it around by running into it, and the tree, on which it can bounce.
The key observations here are that (i) the different objects interact sparsely with one another and (ii) the action A t affects very few objects (in this case, only one).
Thm. 5 will show how one can leverage this sparsity for disentanglement.
2.2. Identiﬁability and model equivalence
To formalize the problem of disentanglement, we will rely on the notion of identiﬁability , which is a property a model has when its parameters can be uniquely determined by the distribution that it represents.
Formally, given some distribution P θ parameterized by θ , this means
[Equation]
For a model as ﬂexible as the one described in the previous section, identifying the exact parameter θ is too strong a demand.
Instead, we will be interested in identifying the parameter θ up to an equivalence class , which amounts to substituting some equivalence relation θ ∼ ˜ θ for θ =
˜ θ in (4).
We now present two equivalence relations for the model presented in Sec.
2.1 adapted from Khemakhem et al.
(2020a): linear and permutation equivalence.
The latter will help us formalize disentanglement.
In what follows, we overload the notation by deﬁning f − 1 ( z <t ) :=
[ f − 1 ( z 1 ) · · · f − 1 ( z t − 1 )] .
Deﬁnition 1 (Linear equivalence)
Let X :
= f ( Z ) and ˜ X :=
˜ f
( Z ) , i.e., the image of the support of Z t under f and ˜ f , respectively.
We say θ is linearly equivalent to ˜ θ if and only if X =
˜ X
and there exists an invertible matrix
L ∈ R
kd z × kd z as well as vectors b, c ∈ R kd z such that
T ( f − 1 ( x ))
=
L T ( ˜ f − 1 ( x ))
+ b , ∀ x ∈ X ; and 2.
L (cid:62) λ ( f − 1 ( x <t ) , a <t )
+ c =
˜ λ
(˜ f − 1 ( x <t ) , a <t ) , ∀ t ∈ { 1 , ..., T } , x <t ∈ X t − 1 , a <t ∈
A t .
In this case, we write θ ∼ L ˜ θ .
Hence, two models are linearly equivalent if they entail the same data manifold X and their respective representations f − 1 ( x ) and ˜ f − 1 ( x ) transformed through the element-wise sufﬁcient statistic T are the same everywhere on X up to an afﬁne transformation .
In the Gaussian case, with variance ﬁxed to one, T ( z ) :
=
z and λ outputs the usual mean parameter µ (here, k = 1 ).
The ﬁrst condition therefore means we can go from one representation to another via an afﬁne transformation.
Suppose θ corresponds to the data generating process, while ˆ θ is some learned model.
Both being linearly equivalent is not enough to declare the learned representation disentangled, since the matrix L might still “mix up” the variables i.e. one component of f − 1 corresponding to multiple components of ˆ f − 1 .
However, if L happens to have a (block-)permutation structure, we have a oneto-one correspondence between the ground truth latent factors of the data and the coordinates of the learned representation.
Deﬁnition 2 (Permutation equivalence)
We say θ is permutation-equivalent to ˜ θ if and only if θ ∼ P ˜ θ
(Def. 1) where P has a block-permutation structure respecting T , i.e. there are d z invertible k × k matrices L 1 , ..., L d z and a d z -permutation π such that for all y =
[ y 1 . . .
y d z ] (cid:62) ∈
R kd z ,
P
[Equation]
In the Gaussian case with a ﬁxed variance, permutation equivalence implies that each coordinate i of one representation is equal to the scaled and shifted coordinate π ( i ) of the other, for some permutation π .
Inspired by previous works on nonlinear ICA, we deﬁne disentanglement as follows.
Deﬁnition 3 (Disentanglement)
Given a ground-truth model θ , we say a learned model ˆ θ is disentangled when θ and ˆ θ are permutation-equivalent.
2.3. Conditions for linear identiﬁability
From now on, it will be useful to think of θ as the ground-truth parameter and ˆ θ as a learned parameter .
The following theorem provides conditions that ensure linear identiﬁability , which is deﬁned as
[Equation]
where L is some invertible matrix (not necessarily with a block-permutation structure).
This theorem is an adaptation and minor extension of Thm.
1 from Khemakhem et al.
(2020a), which we elaborate upon in Sec. 3, and is central to the stronger permutation-identiﬁability (disentanglement)
theorems of the following section.
A proof can be found in App. A.
Theorem 4 (Conditions for linear identiﬁability - Extended from Khemakhem et al.
(2020a))
Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) for a ﬁxed sequence length T .
Suppose the following assumptions hold:
For all i ∈ { 1 , ..., d z } , the sufﬁcient statistic T i is minimal (see next paragraph below).
[Sufﬁcient variability] There exist ( z ( p ) , a ( p ) )
kd z p =0 in their respective supports such that the kd z -dimensional vectors ( λ ( z ( p ) , a ( p ) )
− λ
( z (0) , a (0) ))
kd z p =1 are linearly independent.
Then, we have linear identiﬁability: P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈
A T implies θ ∼ L ˆ θ .
The ﬁrst assumption is a standard one saying that T i is deﬁned appropriately to ensure that the parameters of the exponential family are identiﬁable (see e.g. Wainwright and Jordan (2008, p.
See Def. 9 for a formal deﬁnition of minimality.
The second assumption is sometimes called the assumption of variability (Hyv¨arinen et al., 2019), and requires that the conditional distribution of Z t depends “sufﬁciently strongly” on Z <t and/or A <t .
We stress the fact that this assumption concerns the ground-truth data generating model θ .
Notice that the z ( p ) represent values of Z <t for potentially different values of t and can thus have different dimensions.
In the Gaussian case with variance ﬁxed to one, the sufﬁcient variability assumption requires that the values of the conditional mean E [ Z t | z <t , a <t ] are not all contained in a proper 4 afﬁne subspace of R d z .
This can be interpreted as having a sufﬁciently complex transition model.
2.4. Permutation-identiﬁability via mechanism sparsity regularization
We are now ready to present the core contribution of this work, i.e. a novel permutation-identiﬁability result based on mechanism sparsity regularization (Thm. 5).
The intuition for this result is that, under appropriate assumptions (that are satisﬁed in the motivating example of Fig. 1), models that have an entangled representation also have a denser adjacency matrix ˆ G .
Thus, by regularizing ˆ G to be sparse, we exclude entangled models, leaving us with only the disentangled ones.
Thm. 5 gives precise conditions about the data-generating model θ under which ﬁtting the model ˆ θ and regularizing the graph ˆ G =
[ ˆ G z ˆ G a ] to be sparse will be sufﬁcient to obtain a disentangled model (Def. 3).
Recall that ˆ G z controls the connectivity between the latent variables from one time step to another and that ˆ G a controls the connectivity between the action A <t and the latent variable Z t .
Sec. 3 will contrast these results with those introduced in the recent literature on nonlinear ICA.
Theorem 5 (Disentanglement via mechanism sparsity) Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) for a ﬁxed T representing the same distribution, i.e. P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈A T .
Suppose the assumptions of Thm.
4 hold and:
The sufﬁcient statistic T is d z -dimensional ( k = 1 ) and is a diffeomorphism from
Z to T ( Z ) .
[Sufﬁcient time-variability]
The Jacobian of the ground-truth transition function λ with respect to z varies “sufﬁciently”, as formalized in Assumption 1 in the next section.
Then, there exists a permutation matrix P such that P G z P (cid:62) ⊂
ˆ G z .
5 Further assume that 3.
[Sufﬁcient action-variability]
The ground-truth transition function λ is affected “sufﬁciently strongly” by each individual action a (cid:96) , as formalized in Assumption 2 in the next section.
Then P G a ⊂ ˆ G a .
Further assume that
[Sparsity] ||
ˆ G || 0 ≤
|| G ||
[Equation]
Then, P G z P (cid:62) =
ˆ G z and P G
a =
ˆ G a .
Further assume that
[Graphical criterion]
For all p ∈ { 1 , ..., d z } , there exist sets I , J ⊂ { 1 , ..., d z } and L ⊂ { 1 , ..., d a } such that
[Equation]
where Pa zi and Ch zi are the sets of parents and children of node z i in G z , respectively, while
Ch a(cid:96) is the set of children of a (cid:96) in G a .
Then θ and ˆ θ are permutation-equivalent (Def. 2), i.e. the model ˆ θ is disentangled.
The ﬁrst assumption is satisﬁed for example by the Gaussian case with variance ﬁxed to one since T ( z ) =
z is a diffeomorphism.
In contrast, it is not satisﬁed in the Gaussian case with ﬁxed mean since T ( z ) =
z (cid:12) z is not invertible.
Sufﬁcient variability.
Thm. 5 has also two sufﬁcient variability assumptions, one for G z and one for G a .
Rigorous statements of these are delayed until Sec. 2.4.1.
Intuitively, both assumptions require that the ground-truth transition function λ is complex enough.
Sparsity.
The ﬁrst three assumptions imply that the learned graph ˆ G is a supergraph of some permutation of the ground-truth graph G .
By adding the sparsity assumption, we have that the learned graph ˆ G is exactly a permutation of the ground-truth graph G .
This assumption is satisﬁed if ˆ G is a minimal graph among all graphs that allow the model to exactly match the ground-truth generative distribution.
In Sec. 2.6, we suggest achieving this by regularizing ˆ G to be sparse.
a ta ta tG a z tz tz t
Ch a
Ch a where actions a are enough to satisfy the graphical criterion of Thm. 5.
Indeed, we have { z 1 } =
Ch a 1 ∩ Ch a 2 , { z 2 } =
Ch a 1 ∩ Ch a 3 and { z 3 } =
Ch a 2 ∩
Ch a 3 .
Graphical criterion.
The very last assumption is a graphical criterion that guarantees disentanglement.
This criterion is trivially satisﬁed when G z is diagonal, since { i } =
Pa zi for all i (actions are not necessary here).
This simple case amounts to having mutual independence between the sequences Z ≤ T i , which is a standard assumption in the ICA literature (Tong et al., 1990; Hyvarinen and Morioka, 2017; Klindt et al., 2021).
The illustrative example we introduced in Fig.
1 has a more interesting “non-diagonal” graph satisfying our criterion.
Indeed, we have that { T } =
Pa zT , { R } =
Ch zR ∩ Pa zR and { B } =
Ch zB .
This example is actually part of an interesting family of graphs that satisfy our criterion:
Proposition 6 (Sufﬁcient condition for the graphical criterion)
If G zi,i = 1 for all i (all nodes have a self-loop) and G z has no 2-cycles, then G satisﬁes the graphical criterion of Thm. 5.
Proof Self-loops guarantee i ∈
Pa zi ∩
Ch zi for all i .
Suppose j ∈ Pa zi ∩
Ch zi for some i and j (cid:54) = i .
This implies i and j form a 2-cycle, which is a contradiction.
Thus { i } =
Pa zi ∩
Ch zi for all i .
Prop. 6 implies that actions are not necessary for permutation-identiﬁability.
Conversely, it is also possible to satisfy the graphical criterion without temporal dependencies, e.g., see Fig. 2, which highlights the fact that the condition of Prop. 6 is not necessary.
In App. A.3.2 & A.3.3, we present two theorems which are versions of Thm.
5 specialized to time-sparsity and action-sparsity, respectively.
Note that they are not special cases of Thm. 5.
2.4.1. S UFFICIENT VARIABILITY ASSUMPTIONS
We now present the two technical variability assumptions of Thm. 5.
Intuitively, both assumptions require that the data generating model has a “sufﬁciently complex” transition function λ .
Notation.
Let R d z × d z G z be the set of matrices
M ∈ R d z × d
z such that M i,j = 0 whenever G zi,j = 0 .
Similarly, let R d z Ch a(cid:96) be the subspace of R d z where all coordinates outside Ch a(cid:96) are zero.
Assumption 1 (Sufﬁcient time-variability)
There exist { ( z ( p ) , a ( p ) , τ ( p ) ) }
|| G z
|| 0 p
= belonging to their respective support such that
[Equation]
where D τ ( p ) z λ and D z T are Jacobians with respect to z τ ( p ) and z , respectively.
Notice the Jacobian D τ ( p ) z λ is always in R d z ×
d
z G z because of how G z masks the input of λ i in (2).
The sufﬁcient time-variability assumption further requires that the Jacobian varies “enough” so that it cannot be contained in a proper subspace of R d z × d z G z .
The following sufﬁcient action-variability assumption has an analogous interpretation.
Assumption 2 (Sufﬁcient action-variability) For all (cid:96) ∈ { 1 , ..., d a } , there exist { ( z ( p ) , a ( p ) , (cid:15) ( p ) , τ ( p ) ) } |
[Equation]
a(cid:96) | p
=1 belonging to their respective support such that
[Equation]
where ∆ τ(cid:96) λ ( z <t , a <t , (cid:15) ) is a partial difference deﬁned by
[Equation]
(6) where (cid:15) ∈ R and E (cid:96),τ ∈ R d a × t is the one-hot matrix with the entry ( (cid:96), τ ) set to one.
Thus, (6) is the discrete analog of a partial derivative w.r.t.
a τ (cid:96) .
In App.
A.7, we provide a plausible transition function λ based on the illustrative example of
Fig. 1 and show that it has sufﬁcient variability.
Given the complex interactions which abound in the real world, we conjecture that “realistic” transition functions are “complex enough” to satisfy both assumptions.
2.5. Actions as interventions with unknown targets and sparse mechanism shifts
An important special case of Thm. 5 is when A t − 1 corresponds to a one-hot vector indexing an intervention with unknown targets on the latent variables Z t .
This speciﬁc kind of intervention has been explored previously in the context of causal discovery where the intervention occurs on observed variables instead of latent variables like in our case (Eaton and Murphy, 2007; Mooij et al., 2020;
Squires et al., 2020; Jaber et al., 2020; Brouillard et al., 2020; Ke et al., 2019).
Speciﬁcally, assume
A t − 1 ∈ { 0 , e 1 , ..., e d a } , where each e (cid:96) is a one-hot vector.
The action A t − 1 = 0 corresponds to the observational setting , i.e. when no intervention occurred, while A t − 1 =
e (cid:96) corresponds the (cid:96) th intervention.
In that context, the unknown graph G a describes which latents are targeted by the intervention, i.e. G ai,(cid:96)
= 1 if and only if z i is targeted by the (cid:96) th intervention.
Here, the partial difference ∆ t − 1 (cid:96) λ
( z <t , 0 , 1) measures the difference of natural parameters between the observational setting and the (cid:96) th intervention.
In this context, the assumption that G a is sparse corresponds precisely to the sparse mechanism shift hypothesis from Sch¨olkopf et al.
(2021), i.e. that only a few mechanisms change at a time.
Thm. 5 thus provides precise conditions for when sparse mechanism shifts induce disentanglement.
2.6. Regularized model estimation
In order to estimate from data the model presented in previous sections, we propose to use a maximum likelihood approach based on the well-known framework of variational autoencoders (VAEs) (Kingma and Welling, 2014) in which the decoder neural network corresponds to the mixing function f .
We consider an approximate posterior of the form
[Equation]
[Equation]
where q ( z t | x t ) is a Gaussian distribution with mean and diagonal covariance outputted by a neural network encoder ( x t ) .
In our experiments, the transition functions λ i are parameterized by fully connected neural networks that look only at a ﬁxed window of s lagged latent variables.
In all experiments, ˆ p ( z ti | z <t , a <t ) is Gaussian with a learned variance that does not depend on ( z <t , a <t ) (see App.
B.2 for details).
This variational inference model induces the following evidence lower bound (ELBO) on log ˆ p ( x ≤ T | a <T ) :
[Equation]
We derive this fact in App.
A.8.
The learned distribution will exactly match the ground truth distribution if (i) the model has enough capacity to express the ground-truth generative process, (ii) the approximate posterior has enough capacity to express the ground-truth posterior p ( z t | x ≤ T , a <T ) , (iii) the dataset is sufﬁciently large and (iv) the optimization ﬁnds the global optimum.
If, in addition, the ground truth generative process satisﬁes the assumptions of Thm. 4, we can guarantee that the learned model ˆ θ will be linearly equivalent to the ground truth model θ .
To go from linear identiﬁability (Def. 1) to permutation-identiﬁability (Def. 2), Thm. 5 suggests we should not only ﬁt the data, but also choose the model such that ˆ G is sparse (or minimal).
To achieve this in practice, we add regularization terms − α z
|| ˆ G z
||
and − α a ||
ˆ G a || 0 to the ELBO objective, where α z ≥ 0 and α a ≥ 0 are hyperparameters.
To make the objective amenable to gradient-based optimization, we treat ˆ G zi,j and ˆ G ai,(cid:96) as independent Bernoulli random variables with probabilities of success sigmoid ( γ zi,j ) and sigmoid ( γ ai,(cid:96) ) and optimize the continuous parameters γ z and γ a using the Gumbel-Softmax gradient estimator (Jang et al., 2017; Maddison et al., 2017).
This strategy has been used successfully in previous work to enable gradient-based causal discovery (Ng et al., 2019; Brouillard et al., 2020).
A regularization that is too weak or too strong will result in graphs that are too dense or too sparse, respectively.
In Sec. 4, we select α z and α a using an adaptation of the unsupervised model selection criterion proposed by Duan et al.
3. Related work
Recent theoretical results have shown that nonlinear ICA is possible when leveraging additional assumptions, e.g., nonstationarity (Hyvarinen and Morioka, 2016) and temporal dependencies (Hyvarinen and Morioka, 2017).
Hyv¨arinen et al.
(2019) generalized these works by introducing the notion of auxiliary variables (which correspond to A in our work).
All of these methods rely on noise contrastive estimation (NCE) (Gutmann and Hyv¨arinen, 2012), which underlies the state-ofthe-art in self-supervised representation learning (Oord et al., 2018; Chen et al., 2020), in which identiﬁability has been used as an analysis tool (Roeder et al., 2021; Zimmermann et al., 2021;
Von K¨ugelgen et al., 2021).
Subsequent works have shown similar results using VAEs (Khemakhem et al., 2020a;
Locatello et al., 2020; Klindt et al., 2021), normalizing ﬂows (Sorrenson et al., 2020) and energy-based models (Khemakhem et al., 2020b).
Khemakhem et al.
(2020a), which introduced iVAE, is likely the closest to the present work.
Thm. 4 is quite similar to Thm.
1 from Khemakhem et al.
(2020a), but
iVAE’s notion of linear equivalence is different in that it does not characterize the relationship between λ and ˆ λ , which is crucial for our proof of Thm. 5.
The most signiﬁcant distinction between the theory of (Khemakhem et al., 2020a)
and ours is how permutation-identiﬁability is obtained: Thm. 2 & 3 from iVAE shows that if the assumptions of their Thm. 1 are satisﬁed and T i has dimension k > 1 or is non-monotonic, then the model is not just linearly, but permutation-identiﬁable.
In contrast, our theory covers the case where k = 1 and T i is monotonic, like in the Gaussian case with ﬁxed variance.
Interestingly, Khemakhem et al.
(2020a) mentioned this speciﬁc case as a counterexample to their theory in their Prop. 3.
The extra power of our theory comes from the extra structure in the dependencies of the latent factors coupled with sparsity regularization.
In App.
A.6, we argue that the assumptions of iVAE for disentanglement are less plausible in an environment like the one depicted in Fig.
1, thus highlighting the importance of the case k
= 1 with monotonic T i of Thm. 5.
Similar to our theory, PCL (Hyvarinen and Morioka, 2017) and SlowVAE (Klindt et al., 2021) leverage temporal dependence, but always assume mutual independence of the sequences { Z ti } .
In our notation, this amounts to assuming the graph G z is diagonal.
Our theory allows for more ﬂexibility by accounting for a variety of dependency structures, like a triangular graph G z .
However, we do not claim our theory is a strict generalization of these works, since, for instance, the latent
Laplacian transition model assumed by SlowVAE is not in the exponential family.
Locatello et al.
(2020) also leverages temporal dependence, but assumes each pair ( Z t , Z t +1 ) shares a random subset S of its components.
Our theory allows for every latent factor to change constantly and, thus, does not make this assumption.
Interestingly, they assume that, for all i ,
P ( S ∩ S (cid:48) =
{ i } ) > 0 (for i.i.d S and S (cid:48) ), which resembles our graphical criterion.
While there is a signiﬁcant amount of interest in learning probabilistic or causal graphs between high-level latent variables extracted from low-level observations (Bengio, 2019; Sch¨olkopf, 2019;
Sch¨olkopf et al., 2021; Goyal and Bengio, 2021; Ke et al., 2021), there have been comparatively few practical solutions contributed to the literature.
Of works which learn CGMs, a number assume the causal graph structure is known (Kocaoglu et al., 2018;
Shen et al., 2021; Nair et al., 2019).
The concurrent work of Volodin (2021) independently proposed a very similar approach to jointly disentangle the latent factors by learning a sparse causal graph relating them using binary masks, but focuses more on exploring various algorithm-speciﬁc decisions than on formal identiﬁability proofs and does not use a VAE-based approach to estimate their model.
Bengio et al.
(2020) suggests using adaptation speed as a heuristic objective to disentangle latent factors and their causal relationship in the bivariate case.
Yang et al.
(2021) learns the causal graph by incorporating a “causal model layer” into iVAE, but does not rely on mechanism sparsity to disentangle, does not apply to time-series and is limited to linear CGMs.
The assumption that high-level variables are sparsely related to one another has been leveraged also by Goyal et al.
(2021b,a); Madan et al.
(2021) via attention mechanisms.
Although these works are, in part, motivated by the same core assumption as ours, their focus is more on empirically verifying out-of-distribution generalization than it is on disentanglement (Def. 3) and formal identiﬁability theory.
The assumption that individual actions often affect only one factor of variation has been leveraged for disentanglement by Thomas et al.
Loosely speaking, the theory we developed in the present work can be seen as a formal justiﬁcation for such an approach.
4. Experiments
To illustrate Thm. 5 and the beneﬁt of mechanism sparsity regularization for disentanglement, we apply the regularized VAE method of Section 2.6 on synthetic datasets that both satisfy (Fig. 3) and
[Figure]
a
[Figure]
s u p e
r
v
i
[Figure]
d r a n d o m
[Figure]
violate (Fig. 8, in App.
B.4)
the assumptions of Thm. 5.
Details about the implementation of our approach are provided in App.
B.2 and the code used to run these experiments can be found here: https://github.com/slachapelle/disentanglement_via_mechanism_sparsity .
Synthetic datasets.
The datasets we considered are separated in two groups: time-sparsity and action-sparsity datasets.
The former group has only temporal dependence without actions, we thus ﬁx ˆ G a = 0 , while the latter has only actions without temporal dependence, we thus ﬁx ˆ G z = 0 .
In each dataset, the ground-truth mixing function f is a randomly initialized neural network.
The dimensionality of Z and X are d z = 10 and d x = 20 , respectively.
In the action-sparsity datasets, the dimensionality of A is d a = 10 .
The ground-truth transition model p ( z t | z <t , a <t ) is always a Gaussian with covariance σ 2 z
I and a mean outputted by some function µ G ( z t − 1 , a t − 1 ) (the data is Markovian ).
Hence, each dataset has a 1d sufﬁcient statistic ( k = 1 ) that is also monotonic and, thus, is not covered by the theory of Khemakhem et al.
(2020a).
App.
B.1 provides a more detailed descriptions of the datasets including the explicit form of µ and G in each case.
Note that the learned transition model ˆ p ( z t | z t − 1 , a t − 1 ) is also Gaussian where the mean is outputted by a MLP.
Performance metrics.
To evaluate disentanglement, we will use encoder ( x ) as a proxy for the learned ˆ f − 1 ( x ) .
To assess linear identiﬁability , we perform linear regression to predict the ground-truth latent factors from the inferred ones, and report the coefﬁcient of determination R 2 .
To assess permutation-identiﬁability , i.e. disentanglement, we report the mean correlation coefﬁcient (MCC), which has been used in similar contexts, e.g. Khemakhem et al.
(2020a).
This metric is obtained by
ﬁrst computing the Pearson correlation matrix C ∈ R
d z × d z between the learned representation and the ground truth latent variables.
Then, MCC = max π ∈ permutations 1 d z (cid:80) d z i =1 | C i,π ( i ) | .
For our method, which is the only one learning a graph, we also report the structural hamming distance (SHD) between the ground-truth graph and the learned graph permuted by π ∗ , the optimal permutation found when computing MCC.
We normalize SHD by the maximal number of edges to ensure it is always between 0 and 1.
The normalized SHD is thus the proportion of incorrectly estimated edges in the graph.
Baselines.
On the temporal-sparsity datasets, we compare our approach with TCVAE (Chen et al., 2018), PCL (Hyvarinen and Morioka, 2017) and SlowVAE (Klindt et al., 2021).
On the action-sparsity datasets, we compare with TCVAE and iVAE (Khemakhem et al., 2020a).
We also report the performance of a randomly initialized encoder (Random) and one trained via least-square regression directly on the ground-truth latent factors (Supervised).
See App.
B.6 for details.
Unsupervised hyperparameter selection.
The hyperparameters of the baselines were selected via unsupervised disentanglement ranking (UDR) (Duan et al., 2020).
For our approach, Fig. 3 &show performance for a range of regularization coefﬁcients α z and α a .
We suggest selecting it using
UDR and excluding coefﬁcients that yield graphs with less than d z = 10 edges, as the graphical criterion cannot be achieved in that case.
Fig.
3 & 8 show this unsupervised procedure selects a reasonable regularization coefﬁcient (as indicated by the black star).
See App.
B.7 for details.
Effect of regularization.
Fig.
3 shows that the right amount of mechanism sparsity regularization leads to improved disentanglement (as measured by MCC), which is in line with Thm. 5.
Improvements in SHD indicate that regularization allows for estimation of the causal graph G
(see
App. B.5 for visualizations).
When α z and α a are selected with the ﬁltered UDR procedure, our approach outperforms the baselines by a signiﬁcant margin, while without regularization, it performs similarly.
Most baselines obtain a high R 2 but a low MCC, indicating linear identiﬁability without disentanglement.
These observations are consistent across all four datasets.
Similar observations also hold for randomly sampled graphs (Fig. 5), different noise-levels on the latent variables (Fig. 6) and different noise-levels on the observations (Fig. 7).
See App.
B.3 for details and caveats.
Violating assumptions.
In App.
B.4, Fig.
8 shows experiments on data violating either the sufﬁcient variability assumption or the graphical criterion.
Additionally, Fig.
9 shows data with a sufﬁcient statistic T i of dimension k = 2 , thus violating the ﬁrst assumption of Thm. 5.
On all these datasets, except the time-sparsity data with insufﬁcient variability, regularization improved
MCC, although by a smaller margin than when assumptions are met.
This suggests some of these assumptions might be relaxed.
5. Conclusion
This work proposed a novel principle for disentanglement based on mechanism sparsity regularization .
The idea is based on the assumption that the mechanisms that govern the dynamics of high-level concepts are often sparse: objects usually interact sparsely with each other and actions usually affect only a few entities.
Building on recent developments in nonlinear ICA, we constructed a rigorous theory which provides precise conditions, e.g. on the structure of the ground-truth dependency graph, for when regularizing the mechanisms to be sparse will result in disentanglement.
A special case of our framework shows how one can leverage unknown-target interventions on the latent factors, or sparse mechanism shifts , for disentanglement.
We proposed a regularized
VAE-based approach and demonstrated that it can improve disentanglement in controlled synthetic settings, thereby preparing the stage for more realistic scenarios, e.g. interactive environments.
We believe this work opens up new possibilities at the intersection of causality and disentanglement that leverage structural assumptions .
For instance, we posit that contextual sparsity, i.e., the assumption that objects only interact with each other in particular situations, could be formalized and leveraged for disentanglement using the tools developed in this work.
Acknowledgments
This research was partially supported by the Canada CIFAR AI Chair Program, by an IVADO excellence PhD scholarship, by a Google Focused Research award, the German Federal Ministry of
Education and Research (BMBF): T¨ubingen AI Center, FKZ: 01IS18039A, and by Mitacs through the Mitacs Accelerate program.
The experiments were in part enabled by computational resources provided by Calcul Quebec and Compute Canada.
The authors would like to thank Yoshua Bengio for inspiring mechanism sparsity regularization through various talks and discussions.
The authors would also like to thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Yash Sharma.
Simon Lacoste-Julien is a CIFAR Associate Fellow in the
Learning in Machines & Brains program.
Conclusion
A Theory
A.1 Technical Lemmas and deﬁnitions . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Proof of linear identiﬁability (Thm. 4) . . . . . . . . . . . . . . . . . . . . . . . .
A.3 Theory for disentanglement via mechanism sparsity . . . . . . . . . . . . . . . . .
A.3.
Central Lemmas for Thm. 5, 21 & 22 . . . . . . . . . . . . . . . . . . . .
A.3.
Proof of the specialized time-sparsity theorem (Thm. 21) . . . . . . . . . .
A.3.
Proof of the specialized action-sparsity theorem (Thm. 22) . . . . . . . . .
A.3.
Proof of the combined theorem (Thm. 5) . . . . . . . . . . . . . . . . . .
A.4 Minor extensions of the theory . .
A.4.
Identiﬁability when σ 2 is learned . .
A.4.
Identiﬁability when some learned parameters are independent of the past .
A.4.
Identiﬁability when λ
i at t = 1 differs from λ
i at t >. . . . . . . . . .
A.
On the invertibility of the mixing function f . . . . . . . . . . . . . . . . . . . . .
A.6 Contrasting with the assumptions of iVAE . . . . . . . . . . . . . . . . . . . . . .
A.7 Illustrating the sufﬁcient variability assumptions . . . . . . . . . . . . . . . . . . .
A.8 Derivation of the ELBO . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B Experiments
B.1 Synthetic datasets . .
B.2 Implementation details of our regularized VAE approach . . . . . . . . . . . . . .
B.3 Additional experiments . .
B.4 Experiments that violate assumptions . . . . . . . . . . . . . . . . . . . . . . . .
B.5 Visualizing learned graphs . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.6 Baselines . .
B.7 Unsupervised hyperparameter selection . .
C Author contributions
Appendix A. Theory A.1. Technical Lemmas and deﬁnitions
In this section, we prove a technical lemma which will be useful for central results of this work.
However, this section can be safely skipped at ﬁrst read.
Deﬁnition 7 (Support of a random variable) Let X be a random variable with values in R n with measure P X .
Let O n be the standard topology of R n (i.e. the set of open sets of R n ).
The support of X is deﬁned as
[Equation]
Lemma 8 Let Z be a random variable with values in R m with distribution P Z and Y = f
( Z ) where f : supp ( Z ) ⊂
R m → Y ⊂ R n is a homeomorphism.
Then
[Equation]
Proof.
We ﬁrst prove that f ( supp ( Z ))
⊂ supp ( Y ) .
Let y ∈
f
( supp ( Z ))
⊂ Y and N be an open neighborhood of y , i.e. y ∈ N ∈
O m .
Note that there exists z ∈ supp ( Z ) such that f ( z ) =
y .
Note that z ∈ f − 1 ( { y } ) ⊂ f − 1 ( N ∩ Y ) and that, by continuity of f , f − 1 ( N ∩ Y ) is an open neighborhood of z .
Since z ∈ supp ( Z ) , we have
[Equation]
[Equation]
Hence y ∈ supp ( Y ) , which concludes the “ ⊂ ” part.
To prove the other direction, we notice that Z = f − 1 ( Y ) with f − 1 continuous.
We can thus apply the same argument as above to show f − 1 ( supp ( Y ))
⊂ supp ( Z ) , which implies that supp ( Y ) ⊂
f
( supp ( Z )) .
(cid:4)
We recall the deﬁnition of a minimal sufﬁcient statistic in an exponential family, which can be found in Wainwright and Jordan (2008, p. 40).
Deﬁnition 9 (Minimal sufﬁcient statistic) Given a parameterized distribution in the exponential family, as in (2) , we say its sufﬁcient statistic T i is minimal when there is no v (cid:54) =
such that v (cid:62) T i ( z ) is constant for all z ∈ Z .
The following Lemma gives a characterization of minimality which will be useful in the proof of Thm. 4.
Lemma 10 (Characterization of minimal T )
A sufﬁcient statistic T : Z → R k is minimal if and only if there exists z (0) , z (1) , ..., z ( k ) belonging to the support Z such that the following kdimensional vectors are linearly independent:
[Equation] [Equation]
Proof.
We start by showing the “if” part of the statement.
Suppose there exist z (0) , ..., z ( k ) in Z such that the vectors of (15) are linearly independent.
By contradiction, suppose that T is not minimal, i.e. there exist a nonzero vector v and a scalar b such that v (cid:62) T ( z )
=
b for all z ∈ Z .
Notice that b = v (cid:62) T ( z (0) ) .
Hence, v (cid:62) ( T ( z ( i ) ) − T ( z (0) ))
=
for all i = 1 , ..., k .
This can be rewritten in matrix form as
[Equation]
which implies that the matrix in the above equation is not invertible.
This is a contradiction.
We now show the “only if ” part of the statement.
Suppose that there is no z (0) , ..., z ( k ) such that the vectors of (15) are linearly independent.
Choose an arbitrary z (0) ∈
Z .
We thus have that
U :=
span { T ( z ) − T ( z (0) )
| z ∈ Z} is a proper subspace of R k .
This means the orthogonal complement of U , U ⊥ , has dimension 1 or greater.
We can thus pick a nonzero vector v ∈ U ⊥ such that v (cid:62) ( T ( z ) − T ( z 0 ))
=
for all z ∈ Z , which is to say that v (cid:62) T ( z ) is constant for all z ∈ Z , and thus, T is not minimal.
(cid:4)
A.2. Proof of linear identiﬁability (Thm. 4)
The following theorem and its proof are a minor extension of that of Khemakhem et al.
(2020a).
The key differences are (i) the fact that the sufﬁcient statistics T i do not have to be differentiable, which allows us to cover discrete latent variables (even though this is not highlighted in the main text), (ii) the notion of linear equivalence does not say anything about the link between λ and ˆ λ , which is crucial for the proof of Thm. 5, and (iii) allowing λ to depend on z <t .
Strictly speaking, point (iii) was not covered by previous nonlinear ICA frameworks since z <t is not observed and, thus, cannot be treated as an auxiliary variable (which must be observed).
Theorem 4 (Conditions for linear identiﬁability - Extended from Khemakhem et al.
(2020a))
Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) for a ﬁxed sequence length T .
Suppose the following assumptions hold:
For all i ∈ { 1 , ..., d z } , the sufﬁcient statistic T i is minimal (Def. 9).
[Sufﬁcient variability] There exist ( z ( p ) , a ( p ) )
kd z p =0 in their respective supports such that the kd z -dimensional vectors ( λ ( z ( p ) , a ( p ) )
− λ
( z (0) , a (0) ))
kd z p =1 are linearly independent.
Then, we have linear identiﬁability: P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈
A T implies θ ∼ L ˆ θ .
Proof.
Equality of Denoised Distributions.
Deﬁne Y t :
= f
( Z t ) .
Given an arbitrary a ∈ A T and a parameter θ =
( f , λ ) , let P Y ≤ T | a ; θ be the conditional probability distribution of Y ≤ T , let P Z ≤ T | a ; θ be the conditional probability distribution of Z ≤ T and let P N ≤ T be the probability distribution of
N ≤ T (the Gaussian noises added on f ( Z ≤ T ) , deﬁned in Sec. 2.1).
First, notice that
[Equation]
(17) where ∗ is the convolution operator between two measures.
We now show that if two models agree on the observations, i.e. P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ , then P Y ≤ T | a ; θ = P Y ≤ T | a ;ˆ θ .
The argument makes use of the Fourier transform F generalized to arbitrary probability measures.
This tool is necessary to deal with measures which do not have a density w.r.t either the Lebesgue or the counting measure, as is the case of P Y ≤ T | a ; θ (all its mass is concentrated on the set f ( R d z ) ).
See Pollard (2001, Chapter 8) for an introduction and useful properties.
[Equation]
where (20) & (23) use the fact the Fourier transform is invertible, (21) is an application of the fact that the Fourier transform of a convolution is the product of their Fourier transforms and (22) holds because the Fourier transform of a Normal distribution is nonzero everywhere.
Note that the latter argument holds because we assume σ 2 , the variance of the Gaussian noise added to Y t , is the same for both models.
For an argument that takes into account the fact that σ 2 is learned and assumes d z < d x , see Appendix A.4.1.
We can further derive that, by Lemma 8, we have that
[Equation]
where we overloaded the notation by deﬁning f ( z ≤ T ) :
= ( f ( z 1 ) , ..., f ( z T )) and analogously for ˆ f ( z ≤ T ) .
Equation (24) shows that that the data manifolds are the same for both models, which is part of the linear equivalence deﬁnition we want to show (Def. 1).
Equality of densities.
Continuing with (23),
[Equation] [Equation]
[Equation] [Equation]
where v :=
ˆ f − 1 ◦ f , with v : Z → ˆ Z .
Note that this composition is well deﬁned because f ( Z )
=
ˆ f ( ˆ Z ) .
We chose to work directly with measures (functions on sets), as opposed to manifold integrals in (Khemakhem et al., 2020a), because it simpliﬁes the derivation of (28) and avoids having to deﬁne densities w.r.t.
measures concentrated on a manifold.
We now derive the density of P Z ≤ T | a
[Equation]
ˆ θ ◦ v w.r.t.
to the Lebesgue measure m .
Let E ⊂ Z T be an event, we then have
[Equation]
where D v is the Jacobian matrix of v and ˆ p refers to the conditional density of the model with parameter ˆ θ .
If Z t are discrete random variables, we can do the same except replacing m by the counting measure and forgetting about the Jacobian of v .
We will present the rest of the argument in the case where Z t is continuous.
The reader should keep in mind that the discrete case is exactly the same, except without the Jacobian of v appearing.
Since P Z ≤ T | a ; θ and P Z ≤ T | a ;ˆ θ ◦ v are equal, they must have the same density:
[Equation]
) where p refers to the conditional density of the model with parameter θ .
For a given t 0 , we have t 0 (cid:89)
[Equation]
[Equation]
by integrating ﬁrst z T , then z t (cid:57) 1 , then ..., up to z t 0 +1 .
Note that we can integrate z t 0 and get t 0 − 1 (cid:89)
[Equation]
[Equation]
By dividing (34) by (35), we get
[Equation]
Linear relationship between T ( f − 1 ( x )) and T (ˆ f − 1 ( x )) .
Recall that we gave an explicit form to these densities in Sec.
2.1 Equations (1) & (2).
By taking the logarithm on each sides of (36) we get d
[Equation]
[Equation]
Note that (37) holds for all z <t and a <t .
In particular, we evaluate it at the points given in the assumption of sufﬁcient variability of Thm. 4.
We evaluate the equation at ( z t , z ( p ) , a ( p ) ) and ( z t , z (0) , a (0) ) and take the difference which yields
[Equation]
=
[Equation]
We regroup all normalization constants ψ into a term d
( z ( p ) , z (
[Equation]
) , a ( p ) , a (0) ) and write
[Equation] [Equation]
) where we used T and λ as deﬁned in Sec. 2.1.
Deﬁne
[Equation]
which yields
[Equation]
We can regroup the w ( p ) into a matrix and the d ( p ) into a vector:
[Equation]
Since (43) holds for all 1 ≤ p ≤ kd z
[Equation]
we can write
[Equation]
Note that W is invertible by the assumption of variability, hence
[Equation]
Let b :=
( dW − 1 ) (cid:62) and L :
= ( ˆ W W − 1 ) (cid:62) .
We can thus rewrite as
[Equation]
Invertibility of L .
We now show that L is invertible.
By Lemma 10, the fact that the T i are minimal (Assumption 1) is equivalent to, for all i ∈ { 1 , ..., d z } , having elements z (0) i , ..., z ( k ) i in Z such that the family of vectors
[Equation]
is independent.
Deﬁne
[Equation]
For all i ∈ { 1 , ..., d z } and all p ∈ { 1 , ..., k } , deﬁne the vectors
[Equation]
For a speciﬁc 1 ≤ p ≤ k and i ∈ { 1 , ..., d z } , we can take the following difference based on (49)
[Equation]
53) where the left hand side is a vector ﬁlled with zeros except for the block corresponding to T i ( z ( p,i ) i ) −
T i ( z (0) i ) .
Let us deﬁne
[Equation]
Note that the columns of ∆ T ( i ) are linearly independent and all rows are ﬁlled with zeros except for the block of rows { ( i − 1) k + 1 , ..., ik } .
We can thus rewrite (53) in matrix form
[Equation]
We can regroup these equations for every i by doing
[Equation]
[Equation]
Notice that the newly formed matrix on the left hand side has size kd z × kd z and is block diagonal.
Since every block is invertible, the left hand side of (55) is an invertible matrix, which in turn implies that L is invertible.
We can rewrite (49) as
[Equation]
56) which completes the proof of the ﬁrst part of θ ∼ L ˆ θ .
Linear identiﬁability of sufﬁcient statistics.
We now want to show the second part of the equivalence which links λ and ˆ λ .
We start from (37) and rewrite it using T and λ
T (
[Equation]
t
[Equation] [Equation]
<
t
[Equation]
a <t ) =
[Equation]
( v ( z t )
[Equation]
(cid:62) ˆ λ ( v ( z <t )
[Equation]
a <t )
+ d ( z <t , a <t )
+ c ( z t )
[Equation]
where all terms depending only on z t are absorbed in c ( z t ) and all terms depending only on z <t and a <t are absorbed in d ( z <t , a <t ) .
Using (49)
[Equation]
we can rewrite (57) as
[Equation]
where ¯ d ( z <t , a <t ) absorbs all terms depending only on z <t and a <t .
Simplifying further we get
[Equation]
where the second equality is obtained by making the change of variable z t ← v − 1 ( z t ) .
Again, one can take the difference for two distinct values of z t , say z t and ¯ z t
[Equation]
while keeping z <t and a <t constant which yields
[Equation]
Using an approach analogous to what we did earlier in the “Invertible L ” step
[Equation]
we can construct an invertible matrix [∆ T (1) ...
∆ T ( d z ) ] and get
[Equation]
where the ∆ c ( i ) are deﬁned analogously to ∆ T ( i ) .
Since [∆ T (1) ...
∆ T ( d z ) ] is invertible we can write
[Equation]
where
[Equation]
which can be rewritten as
[Equation]
This completes the proof.
(cid:4)
A.3. Theory for disentanglement via mechanism sparsity
To understand the proof of Thm.
5, it will be useful to see it as the combination of two other theorems, Thm. 21 & 22.
These theorems can be understood as specialized versions of Thm.
5 for time-sparsity and action-sparsity , respectively.
This section is organised as follows.
App. A.3.
present the lemmas central to all three theorems.
App. A.3.2 & A.3.3 present proofs of the specialized time-sparsity theorem (Thm. 21) and the action-sparsity theorem (Thm. 22), respectively.
App.
A.3.4 ﬁnally demonstrates the theorem presented in the main text, Thm. 5.
All permutationidentiﬁability theorems are based on the linear identiﬁability theorem (Thm. 4).
The structure of the proof is summarized in Fig.
Lemma Lemma
Lemma
Time-sparsity (Thm 21)
Action-sparsity (Thm 22)
Combined (Thm 5)
Linear identifiability (Thm 4)
Permutation identifiability
A.3.1. C ENTRAL L EMMAS FOR T HM . 5, 21 &
Throughout this section, we abstract away the details of our problem setting by working with an arbitrary function of the form
[Equation]
67) where Γ is some arbitrary set.
This function will be replaced by the Jacobian of the transition function for temporal sparsity, or by the discrete derivative of the action function for action sparsity.
We will study how some notion of sparsity of this function behaves when composed with an invertible linear transformation.
The j th column of Λ( γ ) and its i th row will be denoted as Λ · ,j ( γ ) and Λ i, · ( γ ) , respectively.
We use analogous notation for subset of indices S ⊂ { 1 , ..., m } × { 1 , ..., n } :
[Equation] [Equation]
Hence, the above sets correspond to horizontal and vertical slices of S , respectively.
We introduce further notations in the following deﬁnition.
Deﬁnition 11 (Aligned subspaces of R m )
Given a subset S ⊂ { 1 , ..., m } , we deﬁne
[Equation]
Deﬁnition 12 (Aligned subspaces of R m × n ) Given a subset S ⊂ { 1 , ..., m }×{ 1 , ..., n } , we deﬁne
[Equation]
Analogously, given a binary matrix B ∈ { 0 , 1 } m × n , we deﬁne
[Equation]
We can now deﬁne what we mean by sparsity:
Deﬁnition 13 (Sparsity pattern of Λ ) A sparsity pattern of Λ is the smallest subset S of { 1 , ..., m } × { 1 , ..., n } such that Λ(Γ)
⊂
R m × n S .
Thus, the sparsity pattern of Λ describes which entries of the matrix Λ( γ ) are not zero for some γ ∈ Γ .
The following two simple lemmas will be key for the following results.
Lemma 14 Let S, S (cid:48) ⊂ { 1 , ..., m } 2 and let ( B i,j ) ( i,j ) ∈ S be a basis of R m × m S .
Let L be a real m × m matrix.
Then
[Equation]
Proof.
Choose ( i 0 , j 0 ) ∈ S .
We can write the one-hot matrix E i 0 ,j 0 as (cid:80) ( i,j
) ∈ S α i,j B i,j for some coefﬁcients
α i,j .
Thus
[Equation]
where the ﬁnal ” ∈ ” holds because each element of the sum is in R m × m S (cid:48)
[Equation]
Lemma 15 Let S, S (cid:48)
⊂ { 1 , ..., m } and ( b i ) i ∈ S be a basis of R mS .
Let L be a real m × m matrix.
Then
[Equation]
Proof.
Choose i 0 ∈ S .
We can write the one-hot vector e i 0 as (cid:80) i ∈
S
α i b i for some coefﬁcients α i (since ( b i ) i ∈
S forms a basis).
Thus
[Equation]
where the ﬁnal ” ∈ ” holds because each element of the sum is in R mS (cid:48) .
(cid:4)
We also need to deﬁne what we are looking for
Deﬁnition 16 (Permutation-Scaling Matrix)
A matrix is said to be permutation-scaling if every row or column contains exactly one non-zero element.
Alternatively, permutation-scaling matrices are deﬁned as the matrices that can be written as P D where P is a permutation matrix and D is a full rank diagonal matrix.
We are now ready to show the lemma central to Thm. 21.
Lemma 17 (Sparsest L (cid:62) Λ( · )
L implies L is a permutation) Let Λ : Γ → R m × m with sparsity pattern S .
Let L ∈ R m × m be an invertible matrix and ˆ S be the sparsity pattern of ˆΛ( . ) :
= L (cid:62) Λ( · ) L .
Assume that
[Equation]
Then there exists an m -permutation σ such that σ ( S ) ⊂ ˆ S , where σ ( S ) :
= { ( σ ( i ) , σ ( j )) | ( i, j ) ∈ S } .
Further assume that
[Equation]
Then σ ( S )
=
ˆ S .
Further assume that 3.
[Graphical Criterion]
For all p ∈ { 1 , ..., m } , there exist index sets I , J ⊂ { 1 , ..., m } such that
[Equation]
Then L is a permutation-scaling matrix.
Proof .
We separate the proof in four steps.
The ﬁrst step leverages the Assumption 1 and Lemmato show that L must contain ”many” zeros.
The second step leverages the invertibility of L to show that σ ( S ) ⊂ ˆ S .
The thirst step uses Assumption 2 to establish σ ( S )
=
ˆ S .
The fourth step uses
Assumption 3 to show that L must have a permutation structure.
We will denote by N the sparsity pattern of L (Deﬁnition 13).
N is thus the set of index couples corresponding to nonzero entries of L .
Step 1: By Assumption 1, there exists ( γ i,j ) ( i,j
) ∈ S such that (Λ( γ i,j ))
( i,j
) ∈
S spans R m × m S .
Moreover, by the deﬁnition of ˆ S as sparsity pattern of L (cid:62) Λ( . )
L (Deﬁnition 13), we have for all ( i, j ) ∈ S
[Equation]
Then, by Lemma 14, we must have
[Equation]
We can rewrite our ﬁnding as
[Equation]
Step 2: Since the matrix L is invertible, its determinant is non-zero, i.e.
[Equation]
where S m is the set of m -permutations.
This equation implies that at least one term of the sum is non-zero, meaning
[Equation]
In other words, this m -permutation σ is included in the sparsity pattern of L – i.e. it is such that for all i ∈ { 1 , ..., m } , σ ( i ) ∈ N · ,i .
We thus have for all ( i, j ) ∈ S that
[Equation]
which implies that
[Equation]
where σ ( S ) :
= { ( σ ( i ) , σ ( j ))
| ( i, j ) ∈ S } .
This proves the ﬁrst claim of the Thm..
Step 3:
By Assumption 2, | ˆ S | ≤ | S | = | σ ( S ) | , we must have that
[Equation]
[Equation]
which proves the second statement of the Thm..
Step 4:
To show that L is a permutation-scaling matrix we show that any two rows cannot have nonzero entries on the same column.
We will proceed by contradiction.
Suppose there are distinct rows i 1 and i 2 such that N i 1 , · ∩ N i 2 , · (cid:54) = ∅ .
Choose i 3 such that
[Equation]
Notice that we must have i 3 (cid:54)
=
i 1 or i 3 (cid:54) =
i 2 .
Without loss of generality, assume the former holds.
By Assumption 3, there are sets of indices
I and J such that (cid:0)(cid:84)
i ∈I S i, · (cid:1) ∩ (cid:16)(cid:84) j ∈J S · ,j (cid:17) =
{ i 1 } .
Since i 3 (cid:54) =
i 1 , one of the two following statement holds:
[Equation]
Case 1.
Suppose ∃ i 0 ∈
I s.t.
i 3 (cid:54)∈ S i 0 , · .
We must have i 1 ∈ S i 0 , · , which is equivalent to having ( i 0 , i 1 ) ∈ S .
Equations (82) & (91) imply that
N i 0 , ·
[Equation]
N i 1 , ·
[Equation]
S , (92) and since ( σ ( i 0 ) , σ ( i 3 ))
∈ N i 0 , · × N i 1 , · (by (88)), we have that ( σ ( i 0 ) , σ ( i 3 ))
∈ ˆ S .
But this implies, by (87), that ( i 0 , i 3 ) ∈ S , which contradicts (89).
Case 2.
Suppose ∃ j 0 ∈ J s.t.
i 3 (cid:54)∈ S · ,j 0 .
We must have that i 1 ∈ S · ,j 0 which is equivalent to having (
[Equation]
Equations (82) & (94) imply that
[Equation]
and since ( σ ( i 3 ) , σ ( j 0 )) ∈ N i 1 , · × N j 0 , · (by (88)), we have that
[Equation]
σ ( i
[Equation]
, σ
[Equation]
But this implies, by (87), that ( i 3 , j 0 ) ∈ S , which violates the fact that i 3 (cid:54)∈ S · ,j 0 .
(cid:4)
We now present the lemma central to Thm. 22.
The statement and its proof are very similar to
Lemma 17.
Lemma 18 (Sparsest L T Λ( . ) implies L is a permutation)
Let Λ : Γ → R m × n with sparsity pattern S .
Let L ∈ R m × m be an invertible matrix and ˆ S be the sparsity pattern of ˆΛ := L Λ .
Assume that
[Sufﬁcient Variability] For all j ∈ { 1 , ..., n } , span (Λ · ,j (Γ))
=
R mS · ,j .
Then there exists an m -permutation σ such that σ ( S ) ⊂ ˆ S where σ ( S ) :
= { ( σ ( i ) , j ) | ( i, j ) ∈ S } .
Further assume that
[Sparsity] | ˆ S | ≤ | S | .
Then σ ( S )
=
ˆ S .
Further assume that
[Graphical Criterion]
For all p ∈ { 1 , ..., m } , there exists an index set J ⊂ { 1 , ..., n } such that (cid:84) j ∈J S · ,j = { p } .
Then, L is a permutation-scaling matrix.
Proof .
We separate the proof in four steps.
The ﬁrst step leverages the Assumption 1 and Lemmato show that L must contain “many” zeros.
The second step leverages the invertibility of L to show that σ ( S ) ⊂ ˆ S .
The third step uses Assumption 2 to show this inclusion is in fact an equality.
Finally the fourth step Assumption 3 to show that L must have a permutation structure.
We will denote by N the sparsity pattern of L (Deﬁnition 13).
N is thus the set of index couples corresponding to nonzero entries of L .
Step 1: Fix j ∈ { 1 , ..., n } .
By Assumption 1, there exists ( γ i ) i ∈ S · ,j such that (Λ · ,j ( γ i ))
i ∈ S · ,j spans R mS · ,j .
Moreover, by the deﬁnition of ˆ S as sparsity pattern of L Λ( . )
(Deﬁnition 13), we have for all i ∈ S · ,j
[Equation] [Equation]
By Lemma 15, we must have
[Equation]
Since j was arbitrary, this holds for all j .
We can thus rewrite our ﬁnding with N the sparsity pattern of L
[Equation] [Equation]
Step 2: We
know there exists
an m -permutation σ such that for all i ∈ { 1 , ..., m } , σ ( i ) ∈ N · ,i (see
Step in Lemma 17)
We thus have for all ( i, j ) ∈ S that
[Equation]
which implies that
[Equation]
where σ ( S ) :
= { ( σ ( i ) , j ) | ( i, j ) ∈ S } .
This proves the ﬁrst statement of the Thm..
Step 3:
By Assumption 2, | ˆ S | ≤ | S | = | σ ( S ) | , so the inclusion (101) is actually an equality
[Equation]
which proves the second statement.
Step 4:
To show that L is a permutation-scaling matrix, we must show that any two columns cannot have nonzero entries on the same row.
We will proceed by contradiction.
Suppose there are two distinct columns i
and i 2 such that N · ,i 1 ∩ N · ,i 2 (cid:54) = ∅ .
Choose i 3 such that
[Equation]
Notice that we must have i 3 (cid:54)
=
i 1 or i 3 (cid:54) =
i 2 .
Without loss of generality, assume the former holds.
By Assumption 3, there is a set of indices J such that (cid:84) j ∈J S · ,j = { i 1 } .
Since i 3 (cid:54) =
i 1 , there must exist some j 0 ∈ J such that
[Equation]
Moreover, we must have i 1 ∈ S · ,j 0 , which is equivalent to ( i 1 , j 0 ) ∈ S .
Equations (99) & (105) imply that
[Equation]
(106) and since ( σ ( i 3 ) , j 0 ) ∈ N · ,i 1 × { j 0 } (by (103)), we have that (
[Equation]
But this implies, by (
), that ( i 3 ,
j
0 ) ∈ S ,
which contradicts
(cid:4)
In Sec.
A.3.4, we will present Thm. 5 and its proof which is, in some sense, the combination of
Thm. 21 & 22.
Its proof relies on the following lemma, which is, analogously, the combination of
Lemmas 17 & 18.
Lemma 19 (Combining Lemmas 17 & 18) Let Λ (1) : Γ (1)
→ R m × m and Λ (2) :
Γ (2)
→ R m × n with sparsity pattern S (1) and S (2) , respectively.
Let L ∈ R m × m be an invertible matrix and ˆ S (1) and ˆ S (2) be the sparsity patterns of ˆΛ (1) :
= L (cid:62) Λ (1) ( · ) L and ˆΛ (2) :
= L (cid:62) Λ (2) , respectively.
Assume that
[Sufﬁcient Variability 1] span (Λ (1) (Γ (1) ))
=
R m × m S (1) .
Then there exists an m -permutation σ such that σ ( S (1) )
⊂ ˆ S (1) where σ ( S (1) ) :
= { ( σ ( i ) , σ ( j )) | ( i, j ) ∈ S (1) } .
Further assume that 2.
[Sufﬁcient Variability 2]
For all j ∈ { 1 , ..., n } , span (Λ (2) · ,j (Γ (2) ))
=
R m S (2) · ,j .
Then σ ( S (2) )
⊂ ˆ S (2) where σ ( S (2) ) :
= { ( σ ( i ) , j ) | ( i, j ) ∈ S (2) } .
Further assume that 3.
[Sparsity] | ˆ S (1) | + | ˆ S (
[Equation]
| ≤ | S (
[Equation]
| + | S (2) | .
Then σ ( S (1) )
=
ˆ S (1) and σ ( S (2) )
=
ˆ S (2) .
Further assume that 4.
[Graphical Criterion]
For all p ∈ { 1 , ..., m } , there exists an index sets I (1) , J (1) ⊂ { 1 , ..., m } and J (2) ⊂ { 1 , ..., n } such that
[Equation]
Then, L is a permutation-scaling matrix.
Proof .
Following proofs of Lemma 17 and 18 we separate the proof in four steps.
The only real difference with these previous proofs is in step 3 where we leverage Assumption 3.
Step1 leverages
Assumptions 1 & 2 and Lemmas 15 & 14 to show that L must contain “many” zeros.
Step 2 uses the invertibility of L to show σ ( S (1) )
⊂ ˆ S (1) and σ ( S (2) )
⊂ ˆ S (2) .
Step 3 uses Assumption 3 to conclude that σ ( S (1) )
=
ˆ S (1) and σ ( S (2) )
=
ˆ S (2) .
Step 4 uses the Assumption 4 to show that L must have a permutation structure.
We will denote by N the sparsity pattern of L (Deﬁnition 13).
N is thus the set of index couples corresponding to nonzero entries of L .
Step 1: Here, Assumption 1 is the same as Assumption 1 of Lemma 17.
The reasoning of Step 1 of its proof gives equation (82)
[Equation]
Similarly, by Assumption 2, the exact same reasoning as Step 1 of Lemma 18, we reach equation (99)
[Equation]
Step 2:
Following step 2 of Lemma 17, we obtain
[Equation]
where σ ( S (1) ) :
= { ( σ ( i ) , σ ( j )) | ( i, j ) ∈ S (1) } , thus proving the ﬁrst statement.
Similarly, following step 2 of the proof of Lemma 18, we reach
[Equation]
) where σ ( S (2) ) :
= { ( σ ( i ) , j ) | ( i, j ) ∈ S (2) } , thus proving the second statement.
Be aware that σ ( S (2) ) and σ ( S (1) ) carry different meanings.
Step 3:
By the Assumption 3, we have
[Equation] [Equation]
[using (112)] ≤ | ˆ S (1) | + | σ ( S (2) )
|
[Equation]
115) [using (111)] ≤ | ˆ S (1) | + | ˆ S (2) |
[Equation]
116) which implies that all the inequalities we used are actually equalities.
In particular | σ ( S (1) ) | = | ˆ S (1) | and | σ ( S (2) )
| =
| ˆ S (2) | .
Thanks to the inclusions (112) and (111) , we ﬁnally reach σ ( S (1) )
=
ˆ S (1) and σ ( S (2) )
=
ˆ S (2) , which proves the third statement of the Thm..
Step 4:
To show that L is a permutation-scaling matrix, we must show that every row has exactly one nonzero entry.
Since L is invertible, it will have at least one nonzero entry per row, so we only need to make sure it does not have more than one nonzero entry.
We will proceed by contradiction.
Suppose there is a column such that the i 1 th
and i 2 th elements are nonzero.
In other words, there are two distinct rows
i
and i 2 such that N i 1 , ·
∩ N i 2 , · (cid:54) = ∅ .
Choose i 3 such that
[Equation]
Notice σ ( i 3 ) is the problematic column.
Since i 1 (cid:54) =
i 2 , we must have i 3 (cid:54) =
i 1 or i 3 (cid:54) =
i 2 .
Without loss of generality, assume i 3 (cid:54) =
i 1 .
By the Assumption 4, there are sets of indices I (1) , J (1) and J (2) such that
[Equation]
Since i 3 (cid:54) =
i 1 , there are three possibilities:
∃ i 0 ∈
I (1) , i 3 (cid:54)∈ S (1) i 0 , · which is the same as Step 4 Case 1 of Lemma 17, 2.
∃ j 0 ∈ J (1) , i 3 (cid:54)∈ S (1) · ,j 0 which is the same as Step 4 Case 2 of Lemma 17, 3.
∃ j 0 ∈ J (2) , i 3 (cid:54)∈ S (2) · ,j 0 which is the same as Step 4 of Lemma 18.
From the proof of previous lemmas, we know each of these possibilities lead to a contradiction.
We conclude that L must be a permutation-scaling matrix.
(cid:4)
A.3.2. P ROOF OF THE SPECIALIZED TIME - SPARSITY THEOREM (T HM . 21)
The following deﬁnition will be useful.
Deﬁnition 20 (Inclusion of matrices)
For two matrices A and B of same size, we write A ⊂ B to say that ∀ i, j, A i,j (cid:54) =
= ⇒ B
i,j (cid:54) = 0 .
We are now ready to state and prove the specialized time-sparsity theorem.
Theorem 21 (Permutation-identiﬁability from time-sparsity) Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) representing the same distribution, i.e. P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈ A T .
Suppose the assumptions of Thm.
4 hold and that
The sufﬁcient statistic T is d z -dimensional ( k = 1 ) and is a diffeomorphism from Z to T ( Z ) .
[Sufﬁcient Variability] There exist { ( z ( p ) , a ( p ) , τ ( p ) ) }
|| G z
|| 0 p
= belonging to their respective support such that
[Equation]
where D τ ( p ) z λ and D z T are
the Jacobian matrices with respect to z τ ( p ) and z , respectively
Then, there exists a permutation matrix P such that P G z P (cid:62) ⊂ ˆ G z .
Further assume that
[Sparsity] ||
ˆ G z
||
≤
|| G z
|| 0 .
Then, P G z P (cid:62) =
ˆ G z .
Further assume that
[Graphical criterion]
For all p ∈ { 1 , ..., d z } , there exist I , J ⊂ { 1 , ..., d z }
such that
[Equation]
Then θ and ˆ θ are permutation equivalent, θ ∼ P ˆ θ , i.e. the model ˆ θ is disentangled.
Proof.
In what follows, we drop the superscript z on G z and ˆ G z to lighten the notation.
First of all, since the assumptions of Thm. 4 holds and the two models represent the same model, the following relations hold
[Equation]
We can rearrange (119) to obtain
[Equation]
where we deﬁned v :
=
ˆ f − 1 ◦ f .
Taking the derivative of (123) w.r.t.
z
[Equation]
we obtain
[Equation]
We can rewrite (120) as
[Equation]
By taking the derivative of the above equation w.r.t.
z τ for some τ ∈ { 1 , ..., t − 1 } , we obtain
[Equation]
where we use D τz to make explicit the fact that we are taking the derivative with respect to z τ .
By plugging (126) in the above equation and rearranging the terms, we get the master equation
[Equation]
It is now time to make the connection between the mathematical objects of the present theorem and the more abstract ones of Lemma
[Equation]
where the argument γ ∈ Γ of the abstract function Λ( γ ) corresponds to ( z <t , a <t , τ ) .
Deﬁne S and ˆ S the sparsity patterns (Def. 13) of Λ and ˆΛ respectively.
The key point of this proof is to notice the correspondence between sparsity patterns of Jacobians, and dependency graphs.
Assumption 2 of the present theorem guarantees that D τz λ ( z <t , a <t ) D T ( z τ ) − 1 spans
R m × m G , which implies that the sparsity pattern of this Jacobian is equal to G
[Equation]
(131) in the sense that ( i, j ) ∈ S ⇐⇒
G i,j = 1 .
By assumption 1, D T ( · ) − 1 is diagonal and full rank, thus the sparsity pattern of D
τz ˆ λ ( v ( z <t ) , a <t ) D T ( v ( z τ ))
− 1 is the same as D
τz ˆ λ ( v ( z <t ) , a <t ) .
Notice that if ˆ G i,j = 0 , then D τz ˆ λ ( v ( z <t ) , a <t ) i,j = 0 everywhere, and thus ( i, j ) (cid:54)∈ ˆ S .
Taking the contraposition, we get
[Equation]
in the sense that ( i, j ) ∈ ˆ S = ⇒
ˆ
G
i,j = 1 .
We now proceed to demonstrate that all three statements of the present theorem holds.
This is done by showing that all three assumptions of Lemma 17 are satisﬁed by exploiting the correspondence between them and Assumptions 2, 3 & 4 of the present theorem.
Statement 1: Assumption 1 of Lemma 17 directly holds for Λ( γ ) =
D τz λ
( z <t , a <t ) D T ( z τ ) −by Assumption 2 of the present theorem.
This implies that there exists a permutation σ such that σ ( S ) ⊂ ˆ S .
Using (132) & (131), we have that
[Equation]
where P is the permutation matrix associated with σ .
This proves the ﬁrst statement.
Statement 2: We will now show that Assumption 2 of Lemma 17 holds.
Since S = G and ˆ S ⊂ ˆ G , we have
[Equation] [Equation]
By Assumption 3, ||
ˆ G || 0 ≤
|| G || 0 .
Thus, we have
[Equation]
The above equation is precisely Assumption 2 of Lemma 17.
Using || ˆ G || 0 ≤ || G || 0 and (133), we can easily see that
[Equation]
137) which proves the second statement.
Statement 3: We ﬁnally show that Assumption 3 of Lemma 17 holds.
By Assumption 4 of the present theorem, we have that for all p ∈ { 1 , ...,
d z } , there are subsets I , J ⊂ { 1 , ..., d a } such that
[Equation]
where the equivalence holds because S = G .
We can thus apply Lemma 17 to conclude that L is a permutation-scaling matrix, which is the third and ﬁnal statement (cid:4)
A.3.3. P ROOF OF THE SPECIALIZED ACTION - SPARSITY THEOREM (T HM . 22)
The proof of Thm.
22 is very similar in spirit to the proof of Thm. 21, except we use Lemmainstead of Lemma 17.
Theorem 22 (Permutation-identiﬁability from action-sparsity) Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G a ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G a ) representing the same distribution, i.e. P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈ A T .
Suppose the assumptions of Thm.
4 hold and that
Each Z ti has a 1-dimensional sufﬁcient statistic,
i
e. k = 1 . 2.
[Sufﬁcient variability] For all (cid:96) ∈ { 1 , ..., d a } , there exist { ( z ( p ) , a ( p ) , (cid:15) ( p ) , τ ( p ) ) } | Ch a(cid:96)
| p
=1 belonging to their respective support such that
[Equation]
Then, there exists a permutation matrix P such that P G a ⊂ ˆ G a .
Further assume that
[Sparsity] ||
ˆ G a || 0 ≤ || G a || 0 .
Then, P G a =
ˆ G a .
Further assume that
[Graphical criterion]
For all p ∈ { 1 , ..., d z } , there exist L ⊂ { 1 , ..., d a } such that
[Equation]
Then, θ and ˆ θ are permutation-equivalent.
Proof.
In what follows
[Equation]
we drop the superscript a on the the graphs G a
and ˆ G a to lighten notation.
First of all
[Equation]
since the assumptions of Thm. 4 holds, we must have that the following relations hold
[Equation]
We can rewrite (140) as
[Equation]
We can take a partial differences w.r
[Equation]
t. a τ(cid:96) (deﬁned in (6)) on both sides of the equation to obtain
[Equation]
where (cid:15) is some real number.
We can regroup the partial difference for every (cid:96) ∈ { 1 , ..., d a } and get
[Equation]
This allows us to rewrite (142) and obtain the master equation
[Equation]
We now make the connection between the mathematical objects of the present theorem and the more abstract ones of Lemma
L
[Equation]
where the argument γ ∈ Γ of the abstract function Λ( γ ) corresponds to ( z <t , a <t , (cid:15) , τ ) .
Deﬁne S and ˆ S the sparsity patterns of Λ and ˆΛ respectively.
The key point of this proof is to notice the correspondence between sparsity patterns of ﬁnite difference matrices, and dependency graphs.
Assumption 2 of the present theorem guarantees that, for all (cid:96) ≤ d a , ∆ τ(cid:96) λ
( z <t , a <t , (cid:15) )
spans
R d z Ch (cid:96) , which implies that the sparsity pattern of this ﬁnite difference is equal to [ Ch 1 , . . .
, Ch d a ] =
G
[Equation]
(145) in the sense that ( i, j ) ∈ S ⇐⇒
G i,j = 1 .
Recall that ˆ S is the sparsity pattern of ∆ τ ˆ λ ( v ( z <t ) , a <t , (cid:15) ) .
Note that if ˆ G i,(cid:96)
= 0 , then ∆ τ ˆ λ ( v ( z <t ) , a <t , (cid:15) ) i,(cid:96) =
0 everywhere, and thus ( i, j ) (cid:54)∈ ˆ S .
Taking the contraposition, we get
[Equation]
(146) in the sense that ( i, j ) ∈ ˆ S =
⇒
ˆ G i,j = 1 .
We now proceed to demonstrate that all three statements of the present theorem holds.
This is done by showing that all three assumptions of Lemma 18 are satisﬁed by exploiting the correspondence between them and Assumptions 2, 3 & 4 of the present theorem.
Statement 1: Assumption 1 of Lemma 18 directly holds for Λ · ,(cid:96) ( γ ) =
∆ τ(cid:96) λ
( z <t , a <t , (cid:15) ) , for all (cid:96) , by Assumption 2 of the present theorem.
This implies that there exists a permutation σ such that σ ( S ) ⊂ ˆ S .
Using (146) & (145), we have that
[Equation]
where P is the permutation matrix associated with σ .
This proves the ﬁrst statement.
Statement 2: We will now show that Assumption 2 of Lemma 18 holds.
Since S = G and ˆ S ⊂ ˆ G , we have
[Equation]
By Assumption 3, ||
ˆ G || 0 ≤
|| G || 0 .
Thus, we have
[Equation]
The above equation is precisely Assumption 2 of Lemma 18.
Using || ˆ G || 0 ≤ || G || 0 and (147), we can easily see that
[Equation]
which proves the second statement.
Statement 3: We ﬁnally show that Assumption 3 of Lemma 18 holds.
By Assumption 4 of the present theorem, we have that for all p ∈ { 1 , ...,
d z } , there is a subset L ⊂ { 1 , ..., d a } such that
[Equation]
where the equivalence holds because S = G .
We can thus apply Lemma 18 to conclude that L is a permutation-scaling matrix, which is the third and ﬁnal statement.
(cid:4)
A.3.4. P ROOF OF THE COMBINED THEOREM (T HM . 5)
Finally, we can prove Thm. 5, which was presented in the main text.
Theorem 5 (Disentanglement via mechanism sparsity) Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) representing the same distribution, i.e.
P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈ A T .
Suppose the assumptions of Thm.
4 hold and that
The sufﬁcient statistic T is d z -dimensional ( k = 1 ) and is a diffeomorphism from Z to T ( Z ) .
[Sufﬁcient time-variability] There exist { ( z ( p ) , a ( p ) , τ ( p ) ) }
|| G z
|| 0 p
= belonging to their respective support
such
that
[Equation]
where D τ ( p ) z and D z are the Jacobian operators with respect to z τ ( p ) and z , respectively.
Then, there exists a permutation matrix P such that P G z P (cid:62) ⊂
ˆ G z .
8 Further assume that
[Sufﬁcient action-variability] For all (cid:96) ∈ { 1 , ..., d a }
, there
exist { ( z ( p ) , a ( p ) , (cid:15) ( p ) , τ ( p ) ) } | Ch a(cid:96) | p =
belonging to their respective support such that
[Equation]
Then P G a ⊂ ˆ G a .
Further assume that
[Sparsity] ||
ˆ G || 0 ≤
|| G || 0 .
Then, P G z P (cid:62) =
ˆ G z and P G
a =
ˆ G a .
Further assume that
[Graphical criterion]
For all p ∈
., d z
} , there exist sets I , J ⊂ { 1 , ..., d z } and L ⊂ { d a }
such that
[Equation]
Then θ and ˆ θ are permutation-equivalent, i.e. the model ˆ θ is disentangled.
Proof.
This theorem is a combination of 21 & 22 and as such we are going to re-use most of the proofs content.
The idea is to apply Lemma 19, showing the correspondence of assumptions.
Correspondence of parameters.
First of all, since the assumptions of Thm.
4 hold along with assumption 1 & 2, we can get the master equation of the proof of theorem 21 and map to the corresponding abstract functions
[Equation]
Similarly with assumption 3, we get the master equation of the proof of theorem
[Equation]
Let us introduce S (1) , ˆ S (1) , S (2) , ˆ S (2) the sparsity patterns of respectively Λ (1) , ˆΛ (1) , Λ (2) , ˆΛ (2) .
The same mapping between sparsity patterns and dependency matrices as in theorem 21 & 22 applies
[Equation]
[Equation] [Equation]
[Equation] [Equation]
Correspondence of assumptions.
Now that we identiﬁed the relevant parameters, we ready to show the correspondence between the assumptions of the present theorem and the four assumptions of Lemma 19.
The assumptions of sufﬁcient time- and action-variability respectively map to the assumptions of sufﬁcient variability 1 and 2 of Lemma 19.
We will now show that the sparsity assumption of Lemma 19 holds.
By the sparsity assumption of this theorem ||
ˆ G z
||
+ ||
ˆ G a || 0 ≤
|| G
z
||
+ ||
G a || 0 .
But we know from the previous identiﬁcation between sparsity patterns and dependency graphs that | ˆ S (1) | ≤ ||
ˆ G z
|| 0 , | ˆ S (2) | ≤
|| ˆ G a || 0 , | S (1) | =
|| G z || 0 , and | S (2)
|
=
|| G a || 0 , thus
[Equation]
The above equation is precisely the sparsity assumption of Lemma 19.
Finally, the equality between graphs and sparsity patterns mean that the graphical criterion is the same between this theorem and Lemma 19.
We can thus apply Lemma 19 to conclude that L is a permutation-scaling matrix.
(cid:4)
A.4. Minor extensions of the theory
The experiments presented in Sec. 4 differed in minor ways from the theory presented in the main paper.
In what follows, we explain how our theory can be extended to cover a wider range of models, including the one used in our experiments.
A.4.1. I DENTIFIABILITY WHEN σ 2 IS LEARNED
In this section, we show how, by adding the extra assumption that d z < d x , we can adapt the argument in Equations (18) to (23) to work when the variance σ 2 of the additive noise is learned, i.e., θ :=
( f , λ , G, σ 2 ) instead of just θ :=
( f , λ , G ) .
Recall Y t :
= f
( Z t ) and that given an arbitrary a ∈ A T and a parameter θ =
( f , λ , G, σ 2 ) ,
P Y ≤ T | a ; θ is the conditional probability distribution of Y ≤ T and P Z ≤ T | a ; θ is the conditional probability distribution of Z ≤ T .
Since now σ 2 is learned, we denote by P N ≤ T ; σ 2 the probability distribution of N ≤ T , the Gaussian noises with covariance σ 2 I .
We now present the modiﬁed argument:
[Equation] [Equation]
where (166) leverages the formula for the Fourier transform of a Gaussian measure.
We now want to show that σ 2 = ˆ σ 2 by contradiction.
Assuming without loss of generality that ˆ σ 2 > σ 2 we have that e − ˆ σ 2 − σ 2t (cid:62) t is the Fourier transform of a Gaussian distribution with mean zero and covariance (ˆ σ 2 − σ 2 ) I .
Thus, the left hand side is the Fourier transform of a probability measure concentrated on a T d z -manifold while the right hand side is the Fourier transform of a convolution between two distributions, one of which has probability mass over all R Td x .
Since d z < d x , the left measure is not absolutely continuous with respect to the Lebesgue measure (on R Td x ) while the right one is.
This is a contradiction since both measures should be equal.
Thus, σ 2 = ˆ σ 2 .
The rest of the proof of Thm. 4 follows through without modiﬁcation.
A.4.2. I DENTIFIABILITY WHEN SOME LEARNED PARAMETERS ARE INDEPENDENT OF THE PAST
Suppose we slightly modify the model of (2) to be
[Equation]
(168) where T
i is the piece of the sufﬁcient statistic with learnable natural parameters λ 0 i which does not depend on ( z <t , a <t ) .
In that case, the learnable parameters of the model are θ =
( f , λ , λ 0 , G ) .
We now show that the proof of Thm.
4 can be adapted to allow for this slightly more general model.
In the proof of Thm. 4, all steps until equation (36) do not depend on the speciﬁc form of p ( z t | z <t , a <t ) , thus they also apply to the more general model (168).
We can now adapt (37) to get: d z (cid:88) i =
[Equation]
In the following step of the proof
[Equation]
we evaluate the above equation at ( z t , z ( p ) , a ( p ) ) and ( z t , z (0) , a (0) ) and take their difference which gives
[Equation]
where the terms T
i ( z ti ) (cid:62) λ 0 i and T
i ( v i ( z t )) (cid:62) ˆ λ
i disappear since they do not depend on ( z <t , a <t ) , just like log
h
i ( z ti ) , log h
i ( v i ( z t )) and log | det D v ( z t ) | .
Notice that (170) is identical to (38) from the proof of Thm. 4.
All following steps are derived from this equation except for (57) which starts from (37), but it can be easily seen that the terms T 0 i ( z ti ) (cid:62) λ
i and T
i ( v i ( z t )) (cid:62) ˆ λ
i will be absorbed in c ( z t ) like the other terms depending only on z t .
We, thus, conclude that Thm.
4 holds also for the slightly more general model speciﬁed in (168).
Since the proofs of Thm. 5, 21 & 22 rely on Thm.
4, this slight generalization applies to them as well.
In our experiments, we parameterize p ( z t | z <t , a <t ) with the standard ( µ, σ 2 ) and not with the natural parameters.
Precisely, µ ( z <t , a <t ) is modeled using a neural network while σ 2 is learned but does not depend on the past.
The natural parameters ( λ 1 , λ 2 ) of a Normal distribution can be written as a function of µ and σ 2 :
[Equation]
We can see that λ 2 does not depend on ( z <t , a <t ) and, thus, the model we learn in practice is an instance of the slightly more general model (168).
A.4.3. I DENTIFIABILITY WHEN λ 0 i AT t = 1 DIFFERS FROM λ 0 i AT t >
The extension of Sec.
A.4.2 implicitly assumed that λ
i is the same for every time steps t .
Indeed, when taking the difference in (170), if ( z (0) , a (0) ) and ( z ( p ) , a ( p ) ) come from different time steps for which λ
i have different values, they would not cancel each other.
This assumption is somewhat problematic in the case where λ 0 i represents the variance, since we might expect the variance of the very ﬁrst latent V [ z 0 i ] to be much larger than the subsequent conditional variances V [ z ti | z <t , a <t ] .
For example, this would be the case in an environment that is initialized randomly, but that follows nearly deterministic transitions.
To solve this issue, we allow only the very ﬁrst λ
i to be different from the subsequent ones, which are assumed to be all equal.
In that case, we must modify the assumptions of sufﬁcient variability so that all ( z ( p ) , a ( p ) ) cannot be selected from the very ﬁrst time step.
A.5. On the invertibility of the mixing function f
Throughout this work as well as many others (Hyvarinen and Morioka, 2016, 2017;
Hyv¨arinen et al., 2019; Khemakhem et al., 2020a;
Locatello et al., 2020; Klindt et al., 2021), it is assumed that the mixing function mapping the latent factors to the observation is a diffeomorphism from Z to X .
In this section, we brieﬂy discuss the practical implications of this assumption.
Recall that a diffeomorphism is a differentiable bijective function with a differentiable inverse.
We start by adressing the bijective part of the assumption.
To understand it, we consider a plausible situation where the mapping f is not invertible.
Consider the minimal example of Fig.
1 consisting of a tree, a robot and a ball.
Assume that the ball can be hidden behind either the tree or the robot.
Then, the mixing function f is not invertible because, given only the image, it is impossible to know whether the ball is behind the tree or the robot.
Thus, this situation is not covered by our theory.
Intuitvely, one could infer, at least approximately, where the ball is hidden based on previous time frames.
Allowing for this form of occlusion is left as future work.
We believe the differentiable part of this assumption is only a technicality that could probably be relaxed to being piecewise differentiable.
Our experiments were performed with data generated with a piecewise linear f , which in not differentiable only on a set of (Lebesgue) measure zero, but this was not an issue in practice.
A.6. Contrasting with the assumptions of iVAE
Recall from Sec.
3 that the most signiﬁcant distinction between the theory of (Khemakhem et al., 2020a)
and ours is how permutation-identiﬁability is obtained: Thm. 2 & 3 from iVAE shows that if the assumptions of their Thm. 1 (which is almost the same as our Thm. 4) are satisﬁed and T i has dimension k > 1 or is non-monotonic, then the model is not just linearly, but permutationidentiﬁable.
In contrast, our theory covers the case where k = 1 and T i is monotonic, like in the
Gaussian case with ﬁxed variance.
Interestingly, Khemakhem et al.
(2020a) mentioned this speciﬁc case as a counterexample to their theory in their Prop. 3.
The extra power of our theory comes from the extra structure in the dependencies of the latent factors coupled with sparsity regularization.
We now argue that the assumptions of iVAE for disentanglement are less plausible in an environment such as the one of Fig.
Assuming the latent factors are Gaussian, the variability assumption of Thm.
4 combined with k > 1 requires the variance to vary sufﬁciently, which is implausible in such a nearly deterministic environment.
Assuming k = 1 with non-monotonic T i implies the conditional mean of Z t does not depend on the past (since the sufﬁcient statistic corresponding to the mean of a Gaussian is monotonous), which is also implausible in this environment.
On the other hand, the case k = 1 with monotonic T i of Thm. 5 is well suited for the situation.
Indeed, again in the Gaussian case, this would amount to predicting only the mean of the future positions of object.
That being said, we also believe practical applications of these ideas will most likely require a combination of different identiﬁability results.
How to formally combined these results is left as future work.
A.7. Illustrating the sufﬁcient variability assumptions
In this section, we construct a simple example based on the situation depicted in Fig.
1 that illustrates a simple case where the assumption of variability is satisﬁed.
Suppose we have a tree and a robot with position T t and R t , respectively (there is no ball).
Suppose the transition model p (( T t , R t ) | ( T <t , R <t ) , A <t ) is Gaussian with variance ﬁxed to a very small value (nearly deterministic) and a mean given by functions µ T ( T t − 1 ) and µ R ( T t − 1 , R t − 1 , R t − 2 , A t − 1 ) speciﬁed by
[Equation] [Equation]
 
[Equation]
if R t
[Equation]
1 < T t − 1 and T t − 1 − δ < R t − 1 + ∆ t −
T t
[Equation]
if T t − 1 < R t − 1 and R t − 1 + ∆ t − 1 < T t − 1 + δ
R t
[Equation]
∆ t
[Equation]
1 , otherwise (173) where ∆ t − 1 :=
R t − 1 − R t − 2 + A t − 1 is the expected change of position of the robot given it does not hit the tree.
Note that the ﬁrst and second case in (173) correspond to when the robot hit the tree from the left and from the right, respectively, and δ is the distance between the center of the tree and the center of the robot when they both touch each other.
The action A t − 1 thus controls the speed and direction of the robot.
Notice that, here, the graph G z and G a are given by
[Equation] [Equation]
We now show that the assumptions of sufﬁcient variability of Thm.
5 hold in this case (Assumptions 1 & 2).
First, recall that for a Gaussian distribution with ﬁxed variance, the natural parameter is given by λ
= µσ .
We thus have
[Equation]
Without loss of generality, we ﬁx σ z = 1 .
We start by showing time-sufﬁcient variability.
Note that taking the Jacobian of (176) with respect to ( T t − 1 , R t − 1 ) we obtain
D t − 1 z
[Equation]
if R t − 1 < T t − 1 and T t − 1 − δ < R t − 1 + ∆ t −(cid:34)
[Equation] [Equation]
(cid:35) , if T t − 1 < R t − 1 and R t − 1 + ∆ t − 1 < T t −
+ δ
[Equation]
, otherwise.
We can do the same thing but differentiating with respect to ( T t − 2 , R t − 2 ) :
[Equation]
t − 2 z
[Equation]

[Equation]

[Equation]
, if R t − 1 < T t − 1 and T t − 1 − δ < R t − 1 + ∆ t −(cid:34)
[Equation]
, if T t
[Equation]
1 < R t
[Equation]
1 and R t −
[Equation]
∆ t
[Equation]
< T t
[Equation] [Equation]
otherwise.
We can see that the ﬁrst and third matrix of (177) together with the third matrix of (178) form a basis of the space R 2 × 2 G z , which proves sufﬁcient time-variability.
We now prove sufﬁcient action-variability.
We can compute the following partial difference with respect to A t − 1 (for a sufﬁciently small step (cid:15) ∈ R ):
[Equation] [Equation]
t
[Equation] [Equation]
t
[Equation]
1 , R t
[Equation]
2 , A t
[Equation]
1 ; (cid:15)
[Equation]
=   (cid:34) (cid:35) , if R t − 1 < T t − 1 and T t − 1 − δ <
R t − 1 + ∆ t −(cid:34) (cid:35) , if T t − 1 < R t − 1 and R t − 1 + ∆ t − 1 < T t − 1 + δ (cid:34)(cid:15) (cid:35) , otherwise.
Clearly, the last vector in (179) spans R 2 Ch a 1 since, recall, Ch a 1 = { 2 } (the robot position R is the second coordinate).
A.8. Derivation of the ELBO
In this section, we derive the evidence lower bound presented in Sec. 2.6.
[Equation]
where the inequality holds because the term at (181) is a Kullback-Leibler divergence
[Equation]
which is greater or equal to 0.
Notice that
[Equation]
Recall that we are considering a variational posterior of the following form:
[Equation]
Equations (186) & (187) allow us to rewrite the term in (184) as
[Equation]
Notice further that
[Equation]
Using (187) & (189), the KL term (185) can be broken down as a sum of KL as:
[Equation]
Putting all together yields the desired ELBO:
[Equation]
Appendix B. Experiments B.1. Synthetic datasets
We now provide a detailed description of the synthetic datasets used in experiments of Sec. 4.
For all experiments, the dimensionality of X t is d x = 20 and the ground-truth f is a random neural network with three hidden layers of 20 units with Leaky-ReLU activations with negative slope of 0.2.
The weight matrices are sampled according to a 0-1 Gaussian distribution and, to make sure f is injective as assumed in all theorems of this paper, we orthogonalize its columns.
Inspired by typical weight initialization in NN (Glorot and Bengio, 2010), we rescale the weight matrices by (cid:113) 2 1+0 . 2(cid:113)
2 d in + d out .
The standard deviation of the Gaussian noise added to f ( z t ) is set to σ = 10 − 2 throughout.
All datasets consist of 1 million examples.
We now present the different choices of ground-truth p ( z t | z <t , a <t ) we explored in our experiments.
In all cases considered(except the experiment with k = 2 of Fig. 9), it is a Gaussian with covariance 0 .
I independent of ( z <t , a <t ) and a mean given by some function µ ( z t − 1 , a t − 1 ) carefully chosen to satisfy the assumptions of Thm. 5.
Notice that we hence are in the case where k = 1 which is not covered by the theory of Khemakhem et al.
(2020a).
We suppose throughout that d z = d a = 10 .
In all time-sparsity experiments, sequences have length T = 2 .
In action-sparsity experiments, the value of T has no consequence since we assume there is no time dependence.
Temporal sparsity with diagonal dependencies (Fig. 3).
In this dataset, each Z ti has only Z t
[Equation]
1 i as parent.
This trivially satisﬁes the graphical criterion of Thm. 5.
The mean function is given by
[Equation]
where the sin function is applied element-wise.
Notice that no auxiliary variables are required.
Temporal sparsity with triangular dependencies (Fig. 3).
We consider a case where the graphical criterion of Thm.
5 is satisﬁed non-trivially.
Let
[Equation]
be the adjacency matrix between Z t and Z t (cid:57) 1 .
The i th row of G z , denoted by G zi , corresponds to the parents of Z ti .
Notice that this connectivity matrix has no 2-cycles and all self-loops are present.
Thus, by Prop. 6, it satisﬁes the graphical criterion.
The mean function in this case is given by
[Equation]
where, the sin function is applied element-wise, the · is the dot product between two vectors and the summation in the sin function is broadcasted.
Once again, the various frequencies and phases in the sin functions ensures the sufﬁcient time-variability assumption of Thm.
5 is satisﬁed.
Temporal sparsity with triangular dependencies and insufﬁcient variability (Fig 8).
This dataset has the same ground truth adjacency matrix as in (192), but a different transition function that does not satisfy the assumption of sufﬁcient time-variability.
We sampled a transition matrix
W with independent Normal 0-1 entries.
The transition function is thus
[Equation] [Equation]
Temporal sparsity with graphical criterion violation (Fig.
In this dataset, the mean function is the same as the one given in (193) except for the adjacency matrix which does not satisfy the graphical criterion and is given by
[Equation]
where I 12 d z × 12 d z is the 12 d z × 12 d z matrix ﬁlled with ones.
Temporal sparsity with k = 2 (Fig. 9).
This dataset has the lower triangular adjacency matrix of (192) and the same mean function of (193), but the variance of z t (we assume diagonal covariance) depends on z t − 1 via
[Equation]
Action sparsity with diagonal dependencies (Fig. 3).
In this setting, d a = d x and the connectivity matrix between A t (cid:57) 1 and Z t is diagonal, which trivially implies that the graphical criterion of
Thm. 5 is satisﬁed.
The mean function is given by
[Equation]
where sin is applied element-wise.
Moreover, the components of the action vector a t (cid:57) 1 are sampled independently and uniformly between − 2 and 2 .
The same sampling scheme is used for all following datasets.
Action sparsity with double diagonal dependencies (Fig. 3).
We consider a case where the graphical criterion of Thm.
5 is satisﬁed non-trivially.
Let
[Equation]
be the adjacency matrix between A t (cid:57) 1 and Z t .
The i th row, denoted by G ai , corresponds to parents of Z ti in A t (cid:57) 1 .
Note that it is analogous to graph depicted in Fig. 2, which satisﬁes the graphical criterion.
The mean function is given by
[Equation] [Equation]
which is analogous to (193).
Action sparsity with double diagonal dependencies and insufﬁcient variability (Fig.
This dataset has the same ground truth adjacency matrix as the above dataset (197), but a different transition function which does not satisfy the assumption of sufﬁcient variability.
We sampled a matrix
W with independent Normal 0-1 entries.
The mean function is thus
[Equation]
Action sparsity with graphical criterion violation (Fig 8).
This dataset does not satisfy the graphical criterion.
The mean function is the same as (198), but its ground-truth graph G a is given by
[Equation]
where I 2 × 2 is the 2 × 2 matrix ﬁlled with ones.
Action sparsity with k = 2 (Fig. 9).
This dataset has the “double diagonal” adjacency matrix of (197) and the same mean function of (198), but the variance of z t (we assume diagonal covariance) depends on a t − 1 via
[Equation]
B.2. Implementation details of our regularized VAE approach
Learned mechanisms.
Every coordinate z i of the latent vector has its own mechanism ˆ p ( z ti | z <t , a <t ) that is Gaussian with mean outputted by ˆ µ i ( z t − 1 , a t − 1 ) (a multilayer perceptron withlayers of 512) and a learned variance which does not depend on the previous time steps.
Strictly speaking, having a learned variance that does not depend on the pas is not covered by the theory presented in the main paper, but
App. A.4.2 extends it to this slightly more general case.
For learning, we use the typical parameterization of the Gaussian distribution with µ and σ 2 and not its exponential family parameterization.
Details about how this interacts with our theory can be found in Sec.
A.4.2.
Throughout, the dimensionality of Z t in the learned model always match the dimensionality of the ground-truth (same for baselines).
Learning the dimensionality of Z t is left for future work.
Prior of Z 1 in time-sparsity experiments.
In time-sparsity experiments, the prior of the ﬁrst latent ˆ p ( Z 1 ) (when t = 1 ) is modelled separately as a Gaussian with learned mean and learned diagonal covariance.
Note that this learned covariance at time t = 1 is different from the subsequent learned conditional covariance at time t > 1 .
How this subtle point interacts with our theory is discussed in App.
A.4.3.
Learned graphs ˆ G z
and ˆ G a .
As explained in Sec. 2.6, to allow for gradient-based optimization, each edge
ˆ G i,j is viewed as a Bernoulli random variable with probability of success sigmoid ( γ i,j ) , where γ i,j is a learned parameter.
The gradient of the loss with respect to the parameter γ
i,j is estimated using the Gumbel-Softmax Gradient estimator (Jang et al., 2017; Maddison et al., 2017).
We found that initializing the parameters γ i,j to a large value such that the probability of sampling all edge is almost one improved performance.
In time-sparsity experiments, there is no action
so ˆ G a is ﬁxed to 0 , i.e. it is not learned.
Analogously, in action-sparsity experiments, there is no temporal dependence so ˆ G z is ﬁxed to 0 .
In all ﬁgures, whenever the regularization coefﬁcient is set to zero, the corresponding adjacency matrix is frozen so that all edges remain active.
Encoder/Decoder.
In all experiments, including baselines, both the encoder and the decoder is modelled given by a neural network with 6 fully connected hidden layers of 512 units with
LeakyReLU activation with negative slope 0 .
For all VAE-based methods, the encoder outputs the mean and a diagonal covariance.
Moreover, p ( x | z ) has a learned isotropic covariance σ 2 I .
Note that σ 2 I corresponds to the covariance of the independent noise N t in the equation X t
= f ( Z t )
+ N t .
The theory presented in the main paper assumes σ 2 ﬁxed, but Sec.
A.4.1 shows how our theory can be adapted to deal with a learned σ 2 , assuming d z < d x .
[Figure]
R^
[Figure]
shd
[Figure]
B.3. Additional experiments
This section presents additional experiments with (i) diverse randomly sampled graphs (Fig. 5), (ii) different levels of noise on the latents (Fig. 6) and (iii) different levels of noise on the observations (Fig. 7).
B.4. Experiments that violate assumptions
Fig. 8 & 9 show experiments on datasets that do not satisfy the assumptions of our theory.
Fig.shows data violating either the sufﬁcient variability assumption or the graphical criterion.
Fig.shows data with a sufﬁcient statistic T i of dimension k
= 2 , thus violating the ﬁrst assumption of Thm. 5.
The only dataset that does not show an improved performance with regularization is the time-sparsity data that has insufﬁcient variability.
In all other datasets, regularization improves
MCC, although by a smaller margin than when assumptions are met.
As suggested by Thm. 5, when sufﬁcient variability holds we can learn the graph, even if the graphical criterion does not hold.
[Figure]
shd
[Figure] [Figure]
[Figure] [Figure]
s e d
[Figure]
a
[Figure]
m
[Figure]
[Figure] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
B.5. Visualizing learned graphs
Fig.
10 & 11 shows examples of learned adjacency matrix ˆ G z and ˆ G a together with their Pearson correlation matrices between ground-truth and learned representations.
We can see how adding mechanism sparsity regularization improves disentanglement.
The learned adjacency matrix is not is not exactly equal to the ground-truth, but is reasonably close.
B.6. Baselines
In synthetic experiments of Sec. 4, all methods used a minibatch size of 1024 and the same encoder and decoder architecture: A MLP with 6 layers of 512 units with LeakyReLU activations (negative slope of 0.2).
We tuned manually the learning rate of each method to ensure proper convergence.
For
VAE-based methods,
i
[Figure]
e. TCVAE, SlowVAE and iVAE, we are always choosing p ( x | z ) Gaussian with a covariance σ 2 I and learn σ 2 .
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Figure]
G
[Figure]
o un d -
[Figure]
r u t h
Z
[Figure]
[Table]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
β -TCVAE.
We used the implementation provided in the original paper by Chen et al.
(2018) which is available at https://github.com/rtqichen/beta-tcvae .
We used a learning rate of 1e-4.
iVAE.
We used the implementation available at https://github.com/ilkhem/icebeem from Khemakhem et al.
(2020a).
In it, the mean of the prior p ( z | a ) is ﬁxed to zero while its diagonal covariance is allowed to depend on a through an MLP.
We change this to allow the mean to also depend on a through the neural network (with 5 layers and width 512).
We also lower bounded its variance as well as the variance of q ( z | x, a ) to improve the stability of learning.
In the original implementation, the covariance of p ( x | z ) was not learned.
We found that learning it (analogously to what we do in our method) improved performance.
We used a learning rate of 1e-4.
SlowVAE.
We used the implementation provided in https://github.com/bethgelab/ slow_disentanglement (Klindt et al., 2021).
Like for other VAE-based methods, we modelled p ( x |
z ) as a Gaussian with covariance σ 2 I and learned σ 2 .
PCL.
We used the implementation provided here: https://github.com/bethgelab/ slow_disentanglement/tree/baselines .
PCL (Hyvarinen and Morioka, 2017) stands for “permutation contrastive learning” and works as follows:
Given sequential data { X t } Tt =1 , PCL trains a regression function r (( x (cid:48) , x )) to discriminate between pairs of adjacent observations (positive pairs) and randomly matched pairs (negative pairs).
The regression function has the form r (( x, x (cid:48)
[Equation]
(202) where h : R d x → R d z is the encoder and B i : R 2 → R are learned functions.
In our implementation, the B i functions are fully connected neural networks with 5 layers and 512 hidden units.
We experimented with the less expressive function suggested in the original work, but found that the extra capacity improved performance across all datasets we considered.
B.7. Unsupervised hyperparameter selection
In practice, one cannot measure MCC since the ground-truth latent variables are not observed.
Unlike in standard machine learning setting, hyperparameter selection for disentanglement cannot be performed simply by evaluating goodness of ﬁt on a validation set and selecting the highest scoring model since there is usually a trade-off between goodness of ﬁt and disentanglement (Locatello et al., 2019, Sec. 5.4).
To circumvent this problem, Duan et al.
(2020) introduced unsupervised disentanglement ranking (UDR) which, for every hyperparameter combinations, measures how consistent are different random intializations of the algorithm.
The authors argue that hyperparameters yielding disentangled representation typically yields consistent representations.
In our experiments, the consistency of a given hyperparameter combination is measured as follows: for every pair of models, we compute the MCC between their representations.
Then, we report the median of all pairwise MCC.
This gives a UDR score for every hyperparameter values considered.
Fig.
12 report the ELBO (normalized between zero and one), the MCC and the UDR score for the experiments of Fig.
We can visualize the trade-off between ELBO and MCC.
However, MCC and UDR correlates nicely except for the non-diagonal time-sparsity dataset, where, for larger regularization values, UDR indicates highly consistent representations despite the bad MCC.
We noticed that these
[Figure]
Baselines.
Two of the baselines considered had hyperparameters to tune, SlowVAE (Klindt et al., 2021) and TCVAE (Chen et al., 2018).
For SlowVAE, we did a grid search on the following values, γ ∈ { 1 . 0 , 2 . 0 , 4 . 0 , 8 . 0 , 16 . 0 } and α ∈ { 1 , 3 , 6 , 10 } .
For TCVAE, we explored β ∈ { 1 , 2 , 3 , 4 , 5 } but the optimal value in terms of disentanglement was almost always 1.
Values of β larger than 5 led to instabilities during training.
The hyperparameters were selected using UDR, as described in the paragraph above.
Appendix C. Author contributions
S´ebastien Lachapelle developed the idea, the theory and proofs behind mechanism sparsity regularization for disentanglement, wrote the ﬁrst draft of the paper, and designed and implemented the regularized VAE-based method.
Pau Rodr´ıguez L´opez ran all experiments appearing in the paper, produced associated ﬁgures and ran experiments with image data that are still work in progress.
Yash Sharma contributed to the research process, the experimental design in particular, implemented and started running experiments on image data that are still work in progress, and contributed to the writing and the literature review.
Katie Everett implemented and started running experiments on image data that are still work in progress and contributed to the writing and ﬁgures.
R´emi Le Priol reviewed the proofs of main theorems, simpliﬁed some arguments and the overall proof presentation and contributed to the writing and ﬁgures.
Alexandre Lacoste produced image datasets that are still being investigated and provided supervision.
Simon Lacoste-Julien helped with overall paper presentation, clariﬁed the conceptual framework and the motivation and provided supervision.