Partial Disentanglement via Mechanism Sparsity
Abstract
Disentanglement via mechanism sparsity was introduced recently as a principled approach to extract latent factors without supervision when the causal graph relating them in time is sparse, and/or when actions are observed and affect them sparsely.
However, this theory applies only to ground-truth graphs satisfying a specific criterion.
In this work, we introduce a generalization of this theory which applies to any ground-truth graph and specifies qualitatively how disentangled the learned representation is expected to be, via a new equivalence relation over models we call consistency .
This equivalence captures which factors are expected to remain entangled and which are not based on the specific form of the ground-truth graph.
We call this weaker form of identifiability partial disentanglement .
The graphical criterion that allows complete disentanglement, proposed in an earlier work, can be derived as a special case of our theory.
Finally, we enforce graph sparsity with constrained optimization and illustrate our theory and algorithm in simulations.
INTRODUCTION
The need for robustness , transferability and explainability in machine learning is motivating recent efforts to develop systems that capture some form of causal understanding [Pearl, 2019, Schölkopf, 2019, Goyal and Bengio, 2021].
Driven by this goal, the emerging field of causal representation learning [Schölkopf et al., 2021] proposes methods that attempt to reconcile the strengths of deep representation learning, which excels on high-dimensional low-level observations like images, with the framework of causality, which offers a formal language to describe and reason about causal relationships between high-level variables, e g. object positions.
The notion of identifiability plays a special role in this quest to more interpretability and robustness, since models that aim at both extracting the causal variables and learning their causal relationships can easily be overdetermined, thus loosing all hope of being interpretable.
The name of the game is thus to come up with inductive biases that sufficiently restrict the model class to be identifiable, while remaining sufficiently expressive to model a complex environment.
Building from the identifiability analyses of the recent literature on nonlinear ICA [Hyvarinen and Morioka, 2016, 2017,
Hyvärinen et al., 2019, Khemakhem et al., 2020a,b], the work of Lachapelle et al.
[2022] proposed mechanism sparsity regularization as an inductive bias to identify the causal latent factors.
The authors showed how learning without supervision simultaneously both the latent factors and the sparse causal graph relating them can induce disentanglement, as long as technical conditions are satisfied, including a novel criterion on the ground-truth causal graph.
A key distinction between other works that also learn a dependency graph over latent variables [Yang et al., 2021, Yao et al., 2022] and “disentanglement via mechanism sparsity" is that, in the latter, disentanglement is driven by sparsity regularization , which allows to identify model classes which are usually not identifiable without this regularization.
Contributions:
In this work, we extend the theory of disentanglement via mechanism sparsity introduced by Lachapelle et al.
Instead of requiring a graphical criterion to guarantee complete disentanglement, our theory applies to arbitrary ground-truth graphs and specifies qualitatively how disentangled the learned representation is expected to be, via a new equivalence relation over models we call consistency (Def. 7).
This equivalence relation captures which variables are expected to remain entangled and which are not, hence the term partial disentanglement .
This allows, for example, to precisely express the fact that we cannot typically identify the basis in which the position of an object is expressed, but can typically disentangle it from the other objects nonetheless.
The graphical criterion of Lachapelle et al.
[2022], which allows complete disentan- glement, can be derived as a special case of our theory.
We also propose to enforce sparsity via constrained optimization instead of regularization, following Gallego-Posada et al.
We finally illustrate our theory in simulations.
Our contribution fits nicely into the framework of Ahuja et al.
[2022a], which shows how, in general, the equivariances of the transition mechanisms characterize how identifiable the representation is.
Lippe et al.
[2022] and Ahuja et al. [2022b] consider settings similar to ours (by interpreting actions as interventions), but the former assumes the intervention targets are known and the notion of sparse perturbation of the latter is closer to Locatello et al.
Also, Lippe et al.
[2022], Von Kügelgen et al.
[2021] and Ahuja et al.
[2022b] allow for a form of block disentanglement similar to our notion of partial disentanglement.
We refer the reader to Lachapelle et al.
[2022] for a more extensive review of the recent literature on disentanglement and nonlinear ICA.
BACKGROUND 2.1 A LATENT CAUSAL MODEL
This subsection is an almost exact transcription of the model exposition of Lachapelle et al.
[2022] which introduced it.
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 [Hyvärinen et al., 2019]
{ 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 a one-hot vector indexing which intervention the corresponding observation was taken from.
From now on, we will refer to A t as the action vector.
We assume the observations { X t } are generated from a sequence of latent d z -dimensional continuous random vectors { Z t } Tt = 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 where Z is the support of
Z t for all t , and X := f ( Z ) , i.e. the image of Z under f .
We suppose that each factor Z ti represents interpretable 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.
Following previous work on nonlinear ICA [Hyvärinen et al., 2019, Khemakhem et al., 2020a], we assume
[Equation]
where each p ( z ti | z <t , a <t ) is in the exponential family
[Wainwright and Jordan, 2008], i.e. p ( z ti | z <t , a <t )
∝
[Equation]
Note that this family includes many well-known distributions such as the Gaussian and beta distributions.
In the
Gaussian case, the sufficient statistic is T i ( z ) :
=
( z, z 2 ) and the base measure is h
i ( z ) :
= 1 √ 2 π .
The function λ i ( G zi ⊙ z <t , G ai ⊙ 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).
Lachapelle et al.
[2022] refers 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 sufficient statistic (which are equal).
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 ⊙ is applied element-wise and broadcasted along the time dimension.
Let G z :=
[ G z 1 · · · G zd z ] ⊤ ∈ R d z × d z , G a := [ G a 1 · · · G ad a ] ⊤ ∈ R d
z × d a , G :=
[ G z G a ] which is the adjacency matrix of the causal graph.
Indeed, (1) & (2) describes a causal graphical model over the unobserved variables Z ≤ T conditioned on the auxiliary variables
A <T .
Let λ ( z <t , a <t )
∈
R kd z be the concatenation of all λ i ( G zi ⊙ z <t , G ai ⊙ 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 .
2.2 MODEL EQUIVALENCE AND COMPLETE DISENTANGLEMENT
Given how expressive the model of Sec. 2.1 is, there is no hope of fully identifying the model from observations.
Fortunately, we will see that it is unnecessary to do so to maintain interpretability.
We now recall notions of model equivalence from Khemakhem et al.
[2020a] & Lachapelle et al.
In what follows, we overload the notation by defining f − 1 ( z <t ) :=
[ f − 1 ( z 1 ) · · · f − 1 ( z t − 1 )] .
Definition 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 ˜ θ , denoted θ ∼ lin ˜
θ , if and only if X = ˜
X
and there exists an invertible matrix
L ∈ R
kd z × kd z and vectors b, c ∈ R
kd z such that
for all x ∈ X ,
[Equation]
and, for all t ∈ { 1 , ..., T } , x <t ∈ X t − 1 , a <t ∈
A t ,
L
[Equation]
To interpret this definition, we consider the special case where p ( z t | z <t , a <t ) follows a Gaussian distribution with variance fixed to one.
In that case, T ( z ) :
=
z and λ outputs the usual mean parameter µ (here, k = 1 ), and thus, the first condition above requires that one can go from the representation ˜ f − 1 ( x ) to the other representation f − 1 ( x ) via an invertible affine transformation.
The second condition on λ and ˜ λ is analogous.
To make sure the latent factors of two different models can be interpreted in the same way, we need something stronger than linear equivalence, since the matrix L can still “mix up” different latent factors.
The following equivalence relation, adapted from Lachapelle et al.
[2022], does not allow for mixing.
Here we assume k = 1 to lighten the notation.
Definition 2 (“Up to permutation” equivalence, k = 1 ) .
We say two models θ :=
( f , λ , G ) and ˜
θ :=
(˜ f , ˜ λ , ˜ G ) are equivalent up to permutation , denoted θ ∼ perm ˜
θ , if and only if there exists a permutation matrix P such that
G z =
P
⊤ ˜ G
z
P
and
G a
=
P ⊤ ˜ G a , and 2.
θ ∼ lin ˜ θ (Def. 1) with L = DP ⊤ , where the matrix D is invertible and diagonal.
Coming back to the Gaussian case with a fixed variance, equivalence up to permutation means that there exists a permutation π such that each coordinate i of one representation is equal to the scaled and shifted coordinate π ( i ) of the other.
Lachapelle et al.
[2022] defines disentanglement as follows (we specify “complete” to contrast with “partial” later on).
Definition 3 (Complete disentanglement) .
Given a groundtruth model θ , we say a learned model ˆ θ is completely disentangled when θ ∼ perm ˆ θ .
We will see later how complete disentanglement can be relaxed to something which falls between linear equivalence and permutation equivalence.
2.3 LINEAR IDENTIFIABILITY
Starting now, the reader should think of θ as the ground-truth parameter and ˆ θ as a learned parameter .
The following theorem is an adaptation and minor extension of Thm.from Khemakhem et al.
[2020a] by Lachapelle et al.
A proof can be found in the latter.
Theorem 4 (Conditions for linear identifiability - Khemakhem et al.
[2020a], Lachapelle et al.
Suppose we have two models as described in Sec.
2.1 with parameters θ =
( f , λ , G ) and ˆ θ =
(ˆ f , ˆ λ , ˆ G ) for a fixed sequence length T .
Suppose the following assumptions hold:
For all i ∈ { 1 , ..., d z } , the sufficient statistic T i is minimal (Def. 9).
[Sufficient 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 identifiability: P X ≤ T | a ; θ = P X ≤ T | a ;ˆ θ for all a ∈
A T implies θ ∼ lin ˆ θ .
The most important assumption is sufficient variability, which states that the ground-truth transition function λ should be “sufficiently complex”.
PARTIAL DISENTANGLEMENT VIA MECHANISM SPARSITY 3.1 PARTIAL DISENTANGLEMENT AND CONSISTENT MODELS
We now give a very simple definition of partial disentanglement, as something which lives strictly between linear equivalence and equivalence up to permutation:
Definition 5 (Partial disentanglement) .
Given a groundtruth model θ , we say a learned model ˆ θ is partially disentangled when θ ∼ lin
ˆ θ with L having at least one zero component and θ ̸∼ perm ˆ θ .
This definition of partial disentanglement ranges from models that are almost completely entangled, i.e. those with a very dense L , to ones that are very close to being completely disentangled, i.e. those with a very sparse L .
Where a learned model falls on this continuum will depend on the ground-truth graph G underlying the data generating process.
To specify precisely where the zero entries of L will be, we will introduce a new equivalence relation over models we call consistency .
In order to do so, we first need to define the property of S -consistency for matrices.
Definition 6 ( S -consistency) .
Given a binary matrix S ∈ { [Equation] 1 } m × n , a matrix C ∈ R m × m is S -consistent when
[Equation]
where [ · ] + := max { 0 , ·} and 1 is a matrix filled with ones (assuming implicitly its correct size).
We will interpret this definition later on in Sec.
For now, it is enough to understand that an S -consistent matrix has zeros where the binary matrix [ 1 − S ( 1 − S ) ⊤ ] + has zeros.
We can now define the novel consistency equivalence relation over models:
Definition 7 (Consistency equivalence, k = 1 ) .
We say two models θ :=
( f , λ , G ) and ˜
θ :=
( ˜ f , ˜ λ , ˜ G ) are consistent , denoted θ ∼ con ˜
θ , if and only if there exists a permutation matrix P such that
G z = P ⊤ ˜
G z P and G a = P ⊤ ˜ G a , and 2.
θ ∼ lin ˜ θ (Def. 1) with L = CP ⊤ , where the matrix C is G z -consistent, ( G z )
⊤ -consistent and G a -consistent (Def. 6).
We demonstrate in App.
A.2.4 that the consistency relation over models is indeed an equivalence relation, as claimed in the the above definition.
This follows from the perhaps surprising fact that the set of invertible S -consistent matrices forms a group under matrix multiplication (see Thm. 20).
The equivalence ∼ perm is stronger than ∼ con , since a diagonal matrix is always S -consistent, for any S .
To see this, notice that [ 1 − S ( 1 − S )
⊤ ]
+ i,i = 1 for all S and i .
3.2 IDENTIFYING THE EQUIVALENCE CLASS OF CONSISTENT MODELS
We now present the main theorem of this work which can be seen as a generalization of Thm.
5 from Lachapelle et al.
It states that, under some conditions, a perfectly fitted and maximally sparse model ˆ θ will be consistent to the ground-truth distribution θ , i.e. θ ∼ con
ˆ θ (Def. 7).
It means we know qualitatively how disentangled the learned representation is expected to be, based on the graph G . See
App.
A.2.5 for a proof.
Theorem 8 (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 sufficient statistic T is d z -dimensional ( k = 1 ) and is a diffeomorphism from Z to T ( Z ) .
[Sufficient time-variability] The Jacobian of the ground-truth transition function λ with respect to z varies “sufficiently”, as formalized in App.
A.2.5.
[Sufficient action-variability]
The ground-truth transition function λ is affected “sufficiently strongly” by each individual action a ℓ , as formalized in App.
A.2.5.
[Sparsity] ||
ˆ G || 0 ≤
|| G || 0 .
Then, ˆ θ is consistent with θ , i.e. θ ∼ con
ˆ θ (Def. 7).
The conclusion that θ ∼ con ˆ θ means that the learned graph ˆ
G is a permutation of the ground-truth graph G and that the learned representation is either completely entangled, partially disentangled or completely disentangled, depending on the ground-truth graph G , as formalized by Def. 7.
The first assumption is satisfied for example by the Gaussian case with variance fixed to one since T ( z ) =
z is a diffeomorphism.
Rigorous statements of the two sufficient variability assumptions , initially introduced by Lachapelle et al.
[2022], are relayed to App.
A.2.5.
Intuitively, they both require that the ground-truth transition function λ is complex enough.
We note that these sufficient variability assumptions play a role similar to the usual faithfulness assumption in causal discovery [Peters et al., 2017, Section 6.5]. See App A.2.6 for more.
The sparsity assumption requires that the learned graph is at least as sparse as the ground-truth graph.
In Sec. 3.3, we suggest achieving this by enforcing a sparsity constraint on ˆ G .
The graphical criterion of Lachapelle et al.
Thm. 8 can be seen as a generalization of Thm.
5 from
Lachapelle et al.
The latter requires that the groundtruth graph G satisfies this criterion: 1 ∀ 1 ≤ i ≤ d z , 
[Equation]
Ch zj 
 ∩   (cid:92) ℓ ∈ Pa ai
Ch aℓ 
 =
{ i } , where Pa zi and Ch zi are the sets of parents and children of node
z i in G z , respectively, while Ch aℓ is the set of children of a ℓ in G a .
This assumption allows Lachapelle et al.
[2022] to identify θ up to ∼ perm (complete disentanglement) instead of up to ∼ con (possibly partial disentanglement).
It turns out that, when G satisfies the above criterion, the set of models that are ∼ con -equivalent to θ is equal to the set of models that are ∼ perm -equivalent to θ .
Therefore, applying Thm.to a ground-truth model that satisfies the graphical criterion will guarantee complete disentanglement (see Prop. 25).
3.2.1 An example & interpretation
We now attempt to build intuition about the equivalence ∼ con
(Def. 7) and Thm. 8 by considering an example where the ground-truth G is given by G z =
(no temporal dependencies) and G a is given by the bottom left of Fig.
1b.
In that case, what does it mean for a model ˆ θ to be consistent with the ground-truth θ ?
Following Def. 7, we first have that the learned graph ˆ G is the same as G , up to a permutation.
Secondly, we have that their representations are linked via a linear transformation L = CP ⊤ where
C is G z -consistent, ( G z )
⊤ -consistent and G a -consistent (Def. 6).
Since G z = 0 , the first two consistency properties are vacuous, i.e. they do not impose anything on C .
However, G a -consistency forces C to have the same zeros as the binary matrix [ 1 − G a ( 1 − G a ) ⊤ ] + .
This binary matrix is represented at the bottom right of Fig.
1b and captures qualitatively how disentangled the learned representation is expected to be (by Thm.
What does Thm.
8 mean in this context?
Assuming the permutation P from Def.
7 is the identity for simplicity, App.
A.2.8 derives the following interpretation: the ground-truth factor z
i is not a function of the learned factor ˆ z j ( C i,j = 0 ) whenever there exists an action a ℓ that targets z i , but not z j .
A similar exercise can be done with different graphs G .
For instance, consider the case where G a = 0 (no action) and
G z is given by the bottom left of Fig.
1a.
In that case, C will have the same zeros as the bottom right of Fig.
1a.
3.3 SPARSE MODEL ESTIMATION
In order to estimate from data the model presented in previous sections, we use almost the same approach as Lachapelle et al.
[2022], except for how sparsity is encouraged.
To estimate the various parameters of the model, we use 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 the same approximate posterior as Lachapelle et al.
[2022], that is q ( z ≤ T | x ≤ T , a <T ) :
= (cid:81)
Tt =1 q ( z t | x t ) , 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 fixed 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
[Equation]
x ≤ T
[Equation] [Equation]
See [Lachapelle et al., 2022] for a derivation of the above.
In order to obtain θ ∼ con ˆ θ .
Thm. 8 suggests that, while fitting the model, we should restrict ˆ G to have at most the same number of edges as G .
To achieve this in practice, Lachapelle et al.
[2022] introduced additional regularizing terms to the ELBO objective: − α z
|| ˆ G z
||
and − α a ||
ˆ G a || 0 .
Moreover, to make the objective amenable to gradient-based optimization, they treat ˆ G zi,j and ˆ G ai,ℓ as independent Bernoulli random variables with probabilities of success sigmoid ( γ zi,j ) and sigmoid ( γ ai,ℓ ) , respectively, and optimize the continuous parameters γ z and γ a using the Gumbel-Softmax gradient estimator [Jang et al., 2017,
Maddison et al., 2017].
We employ a similar strategy, but instead of adding regularization terms, we add a sparsity constraint of the form E
||
ˆ G || 0 ≤ β and solve it using a variant of gradient descent-ascent on the associated Lagrangian function, as originally suggested by Gallego-Posada et al.
[2021] to learn sparse neural networks.
We use the python library Cooper [Gallego-Posada and Ramirez, 2022] which implements this algorithm for PyTorch.
The main advantage of the constrained approach is that the hyperparameter β , the upper bound of the constraint, is easier to interpret than the regularizer coefficients α z and α a , which results in easier value selection, e.g. via cross-validation.
Moreover, this interpretability allowed us to design a very simple schedule for the value of β : We start training with β = max G
|| G
||
and linearly decrease its value until the desired number edges is reached.
See App.
B.2 for optimization details.
EXPERIMENTS
The goal of this section is to demonstrate empirically that
Thm. 8 holds in practice, i.e. that we can identify the equivalence class of models that are consistent (Def. 7) to the ground-truth model.
Our experimental setting is largely based on the one of Lachapelle et al.
and our implementation is also built on their publicly available code.
Synthetic datasets.
We used the same synthetic datasets as Lachapelle et al.
[2022], but with different ground-truth graphs to highlight partially identifiable cases where com-
[Table]
In these cases, our theory can predict qualitatively how disentangled the learned representation is expected to be, via the ∼ con -equivalence
(Def. 7).
We consider two types of datasets, those with temporal dependencies, and those with actions.
In both types of datasets, the ground-truth decoder f is a neural network initialized randomly.
The latent variable Z and observation X have dimensionality d z = 10 and d x = 20 , respectively.
For datasets with actions, d a = 5 .
Just like in Lachapelle et al.
[2022], the ground-truth p ( z t | z <t , a <t ) is Gaussian with covariance σ 2 z
I and a mean outputted by some function µ G ( z t − 1 , a t − 1 ) .
App.
B.1 gives a detailed descriptions of the function µ G for both types of datasets.
We note that the model is well specified, in the sense that transition model ˆ p ( z t | z t − 1 , a t − 1 ) is also Gaussian with a mean outputted by a MLP.
For both types of datasets, we consider two different graphs, G z (1) and G z (2) for the temporal type, and G a (1) and G a (2) for the action type.
These graphs are specified in App.
B.1.
Performance metrics.
We report four metrics to verify if we can recover the correct graphical structure as well as the representation, up to the proper equivalence class.
To measure complete disentanglement (Def. 3), we report the mean correlation coefficient (MCC), which is obtained by first computing the Pearson correlation matrix K ∈ R d z × d z between the ground-truth representation and the learned representation ( K i,j is the correlation between z i and ˆ z j ).
Then
MCC = max P ∈ permutations 1 d z (cid:80) d z i =1 | ( KP ) i,i | .
We denote by ˆ P the optimal permutation found by MCC.
To evaluate whether the learned representation is linearly equivalent to the ground-truth (Def. 1), we perform linear regression to predict the ground-truth latent factors from the learned ones, and report the mean of the Pearson correlations between the predicted ground-truth latents and the actual ones.
This metric is sometimes called the coefficient of multiple correlation , and happens to be the square root of the better known coefficient of determination denoted by R 2 .
The advantage of using R instead of R 2 is that the former is comparable to MCC, and we always have MCC ≤ R .
Let us denote by ˆ L the matrix of estimated coefficients, which should be thought of as an estimation of L in Def. 1.
To evaluate whether the learned representation is consistent to the ground-truth (Def. 7), as predicted by Thm. 8, we perform linear regression on ˆ P ⊤ ˆ
z while constraining the matrix of coefficient to have the same zeros as C from
Def. 7, and report the mean of the associated coefficients of multiple correlation, denoted by R con .
As a consequence, we have that 0 ≤ MCC ≤ R con ≤ R ≤ 1 .
See App.
B.3 for more details on this novel metric.
Sparsity helps.
Table 1 shows that the sparsity constraint yields significant improvement in MCC and R con .
When the sparsity constraint is used, the gap between R con and R is very small (both are almost 1), indicating that the learned latents that were excluded from the linear regression to compute R con add almost no predictive power.
This indicates that the learned model is consistent to the ground-truth model (Def. 7), as predicted by Thm. 8.
Moreover, the gap between MCC and R con is always significant, indicating that the learned representations are not completely disentangled (Def. 3), but are only partially disentangled (Def. 5), as expected.
The small SHD values indicates the graph is properly learned.
See Fig.
1a,b to visualize typical learned graphs.
In all runs using the sparsity constraint, we set the upper bound to be β :=
||
G || 0 .
In practice, ||
G || 0 is unknown and β must be chosen, e.g. using unsupervised disentanglement ranking (UDR)
[Duan et al., 2020].
CONCLUSION
We introduced a generalization of the theory of disentanglement via mechanism sparsity [Lachapelle et al., 2022] which applies to all ground-truth causal graphs G .
We defined a novel equivalence relation over models, we named consistency (Def. 7), and gave conditions for when the corresponding equivalence class can be identified from observations by enforcing sparsity (Thm. 8).
We showed that the equivalences “ ∼ con ” and “ ∼ perm ” coincide when G satisfies the criterion of Lachapelle et al.
[2022], allowing complete instead of partial disentanglement.
Finally, we proposed to enforce sparsity by solving a constrained optimization problem and validated this approach on synthetic data.
Author Contributions
Sébastien Lachapelle wrote the paper, performed the experiments, came up with the theoretical results and proved them. Simon Lacoste-Julien provided supervision that led to clarifying various aspects of this work.
Acknowledgements
This research was partially supported by the Canada CIFAR
AI Chair Program, by an IVADO excellence PhD scholarship and by a Google Focused Research award. The experiments were in part enabled by computational resources provided by Calcul Quebec and Compute Canada.
Simon Lacoste-Julien is a CIFAR Associate Fellow in the Learning in Machines & Brains program.
B Experiments
B.1 Synthetic datasets . .
B.2 Implementation details of the constrained VAE approach . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Details about the R con metric and its relation to MCC and R . . . . . . . . . . . . . . . . . . . . . . . .
A THEORY A.1 MINIMAL SUFFICIENT STATISTICS FOR EXPONENTIAL FAMILIES
The following defines what a minimal sufficient statistics is for an exponential family.
This property ensures that the parameter of an exponential family is identifiable.
See example Wainwright and Jordan
2008, p
Definition 9 (Minimal sufficient statistic) .
Given a parameterized distribution in the exponential family, as in (2) , we say its sufficient statistic T i is minimal when there is no v ̸ = 0 such that v ⊤ T i ( z ) is constant for all z ∈ Z .
A.2 THEORY FOR DISENTANGLEMENT VIA MECHANISM SPARSITY A.2.1 First insight
Recall that the conditions of Thm. 4 implies that the learned model ˆ θ is linearly equivalent to the ground-truth model θ , i.e.
[Equation]
The following specifies an important consequence of linear identifiability.
Note that this argument is taken from Lachapelle et al.
Lemma 10.
Assume the dimensionality of every sufficient statistics T
i is k = 1 .
If two models θ :=
( f , λ , G ) and ˆ θ = (ˆ f , ˆ λ , ˆ G ) are linearly equivalent, i.e. θ ∼ L ˆ θ (Def. 1), then for all z <t , a <t , τ,⃗ϵ in their respective supports,
[Equation]
where D τz λ and D
τz ˆ λ denote Jacobian matrices with respect to z τ and ∆ τ λ and ∆ τ ˆ λ denote matrices of partial differences with respect to a τ , i.e.
[Equation]
See Equation (95) for the definition of ∆ τℓ λ ( z <t , a <t , ϵ ℓ ) .
Proof.
We can rearrange (5) to obtain
[Equation]
where we defined v :
=
ˆ f − 1 ◦ f .
Taking the derivative of (11) w.r.t.
z , we obtain
[Equation]
We can rewrite (6) 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 (14) in the above equation and rearranging the terms, we get the first desired equation:
[Equation]
[Equation]
To obtain the second equation, we take a partial difference w.r.t.
a τℓ (defined in (95)) on both sides of (15) to obtain
[Equation]
where ϵ is some real number.
We can regroup the partial differences for every ℓ ∈ { 1 , ..., d a } and get
[Equation]
This allows us to rewrite (18) and obtain the second desired equation
[Equation]
Following the exposition of Lachapelle et al.
[2022] to improve readability and present our results in their full generality, consider an arbitrary function of the form
[Equation]
where Γ is some arbitrary set.
Depending on the context, this function Λ( γ ) will correspond either to
D τz λ
( z <t , a <t ) D T ( z τ ) − 1 , where Γ is the support of ( z <t , a <t , τ ) , or ∆ τ λ ( z <t , a <t ,⃗ϵ ) , where Γ is the support of ( z <t , a <t ,⃗ϵ, τ ) .
By doing the following substitutions:
[Equation]
21) we get the equation:
[Equation]
[Equation]
where the argument γ ∈ Γ of the abstract function Λ( γ ) corresponds to ( z <t , a <t , τ ) .
We can do an analogous substitution
[Equation]
<t , a <t
[Equation] [Equation]
which yields
[Equation]
where the argument γ ′ ∈ Γ ′ of the abstract function Λ( γ ′ ) corresponds to ( z <t , a <t ,⃗ϵ, τ ) .
Key observation from Lachapelle et al.
[2022]: Notice how the zeros of Λ( γ ) and ˆΛ( γ ) corresponds to the missing edges in G z and ˆ G z , respectively, and how the zeros of Λ ′ ( γ ′ ) and ˆΛ ′ ( γ ′ ) corresponds to the missing edges in G a and ˆ G a , respectively.
The intuition for why sparsity induce disentanglement is that enforcing sparsity of G results in a sparse ˆΛ( γ ) and ˆΛ ′ ( γ ′ ) , which will result in a sparse L via equations (22) & (24).
Since L relates the ground-truth representation with the learned one, a sparse L means a “more disentangled” representation.
The lemmas and definitions of the following section make this intuition precise.
A.2.2 Central Lemmas and Definitions
In order to formalize the intuition presented in the above section, we need to set up some notation and definitions.
Many notation choices, definitions and results are taken from Lachapelle et al.
Notation.
The j th column of Λ( γ ) and its i th row will be denoted as Λ · ,j ( γ ) and Λ i, · ( γ ) , respectively.
For convenience, we will sometimes treat a binary vector b as a set of indices { i | b i = 1 } and sometimes treat a binary matrix B as a set index couples { ( i, j ) | B i,j = 1 } .
For example, this will allow us to write B · ,j 1 ∩ B · ,j 2 , which should be understood as either the index set { i | B · ,j 1 = 1 } ∩ { i | B · ,j 2 = 1 } or the binary vector B · ,j 1 ⊙ B · ,j 2 (where ⊙ is the element-wise product), depending on the context.
Another example would be the complement of a binary vector b c which should be understood as either { i | b i = 0 } or 1 − b , where 1 denotes a vector filled with ones.
The usefulness of this notation will become apparent later on.
We introduce further notations in the following definitions.
Definition 11 (Aligned subspaces of R m ) .
Given a binary vector b ∈ { 0 , 1 } m , we define
[Equation]
Definition 12 (Aligned subspaces of R m × n ) .
Given a binary matrix B ∈ { 0 , 1 } m × n , we define
[Equation]
Next, we define the sparsity pattern of Λ , which compactly captures which of its entries are always zero.
Definition 13 (Sparsity pattern of Λ [Lachapelle et al., 2022]) .
The sparsity pattern of Λ : Γ → R m × n is a binary matrix
S ∈ { 0 , 1 } m × n such that
[Equation]
An other way to phrase this is to say that the sparsity pattern of Λ is the sparsest binary matrix S such that Λ(Γ) ⊂
R m × n S .
We are now ready to present the lemmas that will be central to the main theorems of this work.
Lemma 14 (Lachapelle et al.
Let S, S ′ ∈ { 0 , 1 } m × m and let ( A ( i,j ) )
( i,j ) ∈ S be a basis of R m × m S .
Let L be a real m × m matrix.
Then
[Equation]
Proof.
We start with direction “ = ⇒ ”.
Choose ( i 0 , j 0 ) ∈ S .
Since e i 0 e ⊤ j 0 ∈ R m × m S (where e i denotes the vector with a 1 at entry i and 0 elsewhere) and the matrices A ( i,j ) form a basis of R m × m S , we can write e i
e ⊤ j 0 = (cid:80) ( i,j ) ∈ S
α i,j A ( i,j ) for some coefficients α i,j .
Thus
[Equation]
where the final “ ∈ " holds because each element of the sum is in R m × m S ′ .
We now show the reverse direction “ ⇐ = ”.
Let A ∈ R m × m S .
We can write
[Equation] [Equation]
where the last “ ∈ ” hold because every term in the sum is in R m × m S ′ .
Lemma 15 (Lachapelle et al.
Let s, s ′ ∈ { 0 , 1 } m and ( a ( i ) )
i ∈ s be a basis of R ms .
Let L be a real m × m matrix.
Then
[Equation]
Proof.
We start with “ = ⇒ ”.
Choose i 0 ∈ s .
We can write the one-hot vector e i 0 as (cid:80) i ∈
s α i a ( i ) for some coefficients
α i (since ( a ( i ) )
i ∈ s forms a basis).
Thus
[Equation]
where the final “ ∈ ” holds because each element of the sum is in R ms ′ .
We now show “ ⇐ = ”.
Let a ∈ R ms .
We can write
[Equation]
where the last “ ∈ ” holds because all terms in the sum are in R ms ′ .
The following simple Lemma will be useful throughout this section.
The argument is taken from Lachapelle et al.
Lemma 16 (Sparsity pattern of an invertible matrix contains a permutation) .
Let L ∈ R m × m be an invertible matrix.
Then, there exists a permutation σ such that L i,σ ( i ) ̸ = 0 for all i .
Proof.
Since the matrix L is invertible, its determinant is non-zero, i.e.
[Equation]
[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]
The exact form of the ground-truth graph G will force some of the entries of the matrix L , which relates the ground-truth and the learned representations, to be zero.
Understanding which entries of L are zero is very important to understand qualitatively how disentangled the learned representation is expected to be.
We now recall the notion of S -consistency (introduced in the main text) which will be crucial to precisely relate the form of the ground-truth graph G to the sparsity pattern of L via the consistency equivalence relation (Def. 7) in Thm. 8.
Note that it is reformulated with the notation introduce in this appendix.
Definition 6 ( S -consistency) .
Given a binary matrix S ∈ { 0 , 1 } m × n , a matrix C ∈ R m × m is S -consistent if [Equation]
[Equation] , where [ · ] + := max { 0 , ·} and 1 is a matrix filled with ones (assuming implicitly its correct size).
The following characterization of S -consistency will be useful later on to prove Lemma 18 & 19, to give an intuitive interpretation of S -consistency (Sec. A.2.8) and to relate S -consistency to the graphical criterion introduced by Lachapelle et al.
[2022] (Sec. A.2.7).
Lemma 17 (Characterizing S -consistency) .
Let C ∈ R m × m and S ∈ { 0 , 1 } m × n .
The following statements are equivalent.
C is S -consistent
(Def. 6); 2.
∀ i, ( C i, · ) ⊤ ∈ R m (cid:84) k ∈ Si, · S · ,k ; 3.
∀ j, C · ,j ∈ R m (cid:84)
k ∈ S
c j, · S c · ,k .
Proof.
We proceed by showing how both the second and third statements are equivalent to the first one.
Choose arbitrary i and j .
[Equation]
One can rephrase (40) as
[Equation]
which proves the first and second statements are equivalent.
One can also rephrase (40) as
[Equation]
which proves the first and third statements are equivalent.
Later in Sec.
A.2.3, we show that the set of invertible and S -consistent matrices form a group under matrix multiplication, i.e. that it is closed under matrix multiplication and inversion.
This will be crucial to show that the relation ∼ con (Def. 7) is an equivalence relation (Sec. A.2.4).
We are now ready to show the central lemmas that can be directly applied to easily prove the main theorem of this work,
Thm. 8.
Note that Lemmas 18 & 19 can be thought of as generalizations of Lemmas 17 & 18 from Lachapelle et al.
[2022], respectively.
The difference is that we do not assume anything about the specific form of S , which yields a different (sometime weaker) conclusion.
Lemma 18 ( L ⊤ Λ( · ) L sparse implies L sparse) .
Let Λ : Γ → R m × m with sparsity pattern S (Def. 13).
Let L ∈ R m × m be an invertible matrix and ˆ S be the sparsity pattern of ˆΛ( . ) :
=
L ⊤ Λ( · ) L .
Let σ be a permutation such that for all i ,
L i,σ ( i ) ̸
=
(Lemma 16) and let P be its associated permutation matrix, i.e. Pe i =
e σ ( i ) for all i .
Assume that
[Sufficient Variability] span (Λ(Γ))
=
R m × m S .
Then S ⊂ P ⊤ ˆ SP .
Further assume that
[Sparsity] ||
ˆ S
|| 0 ≤ || S
|| 0 .
Then S =
P ⊤ ˆ SP and L = CP ⊤ where C is S -consistent and S ⊤ -consistent.
Proof.
We separate the proof in four steps.
The first step leverages the Assumption 1 and Lemma 14 to show that L must contain “many" zeros.
The second step leverages the invertibility of L to show that PSP ⊤ ⊂ ˆ S .
The third step uses
Assumption 2 to establish PSP ⊤ =
ˆ S and the fourth step concludes that L = CP ⊤ where C is both S -consistent
and
S ⊤ -consistent.
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 definition of ˆ S as sparsity pattern of L ⊤ Λ( . )
L (Definition 13), we have for all ( i, j ) ∈ S
[Equation]
Then, by Lemma 14, we must have ∀
[Equation]
i, j
[Equation]
i,
[Equation]
Step 2: Since ∀ i, L i,σ ( i ) ̸ = 0 , (44) implies that for all ( i, j ) ∈ S ,
[Equation]
which, in other words, means that
[Equation] [Equation]
This proves the first claim of the theorem.
Step 3:
By Assumption 2, || ˆ S || 0 ≤ || S
||
= ||
PSP ⊤ || 0 , we must have that
[Equation]
which proves the second statement of the theorem.
Step 4: We notice that, since L j,σ ( j ) ̸ = 0 , (44) implies
[Equation]
We interchange indices i and j in the second equation above (this is purely a change of notation)
[Equation]
which yields
[Equation]
Equations (49) & (51) can be rewritten as
[Equation]
Applying left multiplying both equations above by P ⊤
[Equation]
we obtain
[Equation] [Equation]
j ∈ ( S ⊤ ) i, · ( S ⊤ ) · ,j (55)
Let C :
= LP .
Equations (54) & (55) imply that C is S -consistent and S ⊤ -consistent, respectively (by Lemma 17).
Since
L = CP ⊤ , this completes the proof.
Lemma 19 ( L ⊤ Λ( . ) sparse implies L sparse) .
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 ⊤ Λ .
Let σ be a permutation such that for all i , L i,σ ( i ) ̸
=
(Lemma 16) and let P be its associated permutation matrix, i.e. Pe i =
e σ ( i ) for all i .
Assume that
[Sufficient Variability] For all j ∈ { 1 , ..., n } , span (Λ · ,j (Γ))
=
R mS · ,j .
Then S ⊂ P
⊤ ˆ S .
Further assume that
[Sparsity] ||
ˆ S
|| 0 ≤ || S
|| 0 .
Then S =
P ⊤ ˆ S and L = CP ⊤ where C is S -consistent.
Proof.
We separate the proof in four steps.
The first step leverages the Assumption 1 and Lemma 15 to show that L must contain “many” zeros.
The second step leverages the invertibility of L to show that PS ⊂ ˆ S .
The third step uses
Assumption 2 to show this inclusion is in fact an equality and the fourth step concludes that L can be written as an
S -consistent matrix times P ⊤ .
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 definition of ˆ S as sparsity pattern of L ⊤ Λ( . )
(Definition 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 , which allows us to rewrite as
[Equation]
Step 2:
Since L i,σ ( i ) ̸ = 0 for all i , (58) implies that
[Equation]
i
[Equation]
) ∈ ˆ S ·
[Equation]
) which can be rephrased as
[Equation]
This proves the first statement of the theorem.
Step 3:
By Assumption 2, || ˆ S || 0 ≤ || S ||
= ||
PS || 0 , so the inclusion (60) is actually an equality
[Equation]
which proves the second statement.
Step 4: We notice that (58) can be rewritten as
[Equation]
which is equivalent to
[Equation]
We apply P ⊤ on both sides of the above line and get
[Equation]
) which means, by Lemma 17, that LP is S -consistent.
By defining C :
= LP , we have that L = CP ⊤ , which is what we wanted to prove.
A.2.3 Invertible S -consistent matrices form a group under matrix multiplication
The following theorem shows that, perhaps surprisingly, the set of invertible S -consistent matrices forms a group under matrix multiplication, i.e. that the set is closed under multiplication and inversion.
This will be very useful to show that the consistence relation over models, ∼ con
(Def. 7), is an equivalence relation.
The proof can be safely skipped at first read.
Theorem 20.
Let S ∈ { 0 , 1 } m × n .
The identity matrix I is S -consistent; 2.
For any invertible S -consistent matrices C and C ′ , the matrix product CC ′ is also S -consistent; 3.
For any invertible S -consistent matrix C , C − 1 is also S -consistent.
In other words, the set of invertible matrices that are S -consistent forms a group under matrix multiplication.
Proof.
First, let i ≤ m .
Notice that [ 1 − S ( 1 − S )
⊤ ]
+ i,i =
[1 − S i, · ( 1 i, · − S i, · ) ⊤ ] + = 1 .
Thus, I is S -consistent.
Second, we show closure under matrix multiplication.
Let i, j such that [ 1 − S ( 1 − S ) ⊤ ]
+ i,j = 0 .
Consider ( CC ′ )
i,j = C i, · C ′· ,j .
By Lemma 17, we have that
[Equation]
Notice that if the intersection (cid:16)(cid:84) k ∈ S i, · S · ,k (cid:17) ∩ (cid:16)(cid:84) k ∈ S cj, · S c · ,k (cid:17) is empty, the dot product C i, · C ′· ,j is zero and the second statement of this theorem holds.
By (40) from the proof of Lemma 17, there exists a k such that k ∈ S i, · and k ∈ S cj, · , and, since S · ,k ∩ S c · ,k = ∅ , the initial intersection is itself empty.
Third, we show that the inverse C − 1 is also S -consistent.
Notice that, since C is invertible, there exists a sequence of elementary row operations that will transform C into the identity.
This process is sometimes called Gaussian elimination or
Gauss-Jordan elimination.
The elementary row operations are (i) swapping two rows, (ii) multiplying a row by a nonzero number, and (iii) adding a multiple of one row to another.
These three elementary operation can be performed by left multiplying by an elementary matrix , which have the following forms:
(i) Swapping two rows:
[Equation]
(ii) Multiplying a row by a nonzero number:
[Equation] [Equation]
(iii) Adding a multiple of a row to another:
[Equation]
We will show that it is possible to transform C into the identity by using only elementary matrices that are themselves
S -consistent , i.e. that there exists a sequence of S -consistent elementary matrices E 1 , ..., E p , such that E p ...
E 2 E 1 C
= I .
Since this implies C − 1 = E p ...
E 2 E 1 and all elementary matrices are S -consistent, C − 1 is also S -consistent (using closure under multiplication shown above).
We now construct the sequence of E i using standard Gaussian elimination.
Start by initializing M := C .
Throughout the algorithm, M will be gradually transformed by elementary operations that are S -consistent
(and invertible)
, thus M will remain S -consistent
(and invertible).
We consider every column j = 1 , ..., m from left to right.
If M j,j = 0 , we will show that rows j, ..., m can be permuted to obtain M j,j ̸ = 0 using an S -consistent permutation, but we delay this technical step to the end of the proof to avoid breaking the flow of the exposition.
For now, assume M j,j ̸ = 0 .
Rescale row j so that M j,j =using matrix of the form (68), which is S -consistent.
Then, put zeroes below M j,j by adding a multiple of row j to each row i >
j such that M i,j ̸ = 0 .
Each of these operations corresponds to an elementary matrix of the form (69) where the nonzero entry below the diagonal is at position ( i, j ) .
Since M is S -consistent
and M i,j ̸ = 0 , these elementary matrices must also be
S -consistent.
Once every element below M j,j are zero go to the next column.
Do that for all columns.
At this point, M is upper triangular with a diagonal filled with ones.
We must now remove every nonzero elements above the diagonal by a process similar to what we just did.
Start with column j = n up to j = 1 , from left to right.
To remove every nonzero elements above M j,j , we can add a multiple of row j to the rows i < j that have M i,j ̸ = 0 .
This is equivalent to multiplying M by an elementary matrix of the form (69) with its off diagonal nonzero entry by at position ( i, j ) .
Again, since M i,j ̸ = 0 and M is S -consistent, this elementary matrix must also be S -consistent.
Once all elements above M j,j are zeros, go to the next column and repeat for every columns until column j = 1 is reached.
At this point, M = I , which is what we wanted to show.
We now have to show what to do when M j,j ̸ = 0 .
We know that M has the following form
[Equation]
70) where U ∈ R ( j − 1) × ( j − 1) is an upper triangular matrix with only ones on its diagonal and B is a square matrix with B 1 , 1 = 0 .
Since M is invertible, B is invertible too (otherwise, det( M )
= det( U ) det( B ) = 0) .
Thus, by Lemma 16, there exists a permutation σ such that for all i , B σ ( i ) ,i ̸
= 0 .
Consider its corresponding permutation matrix P :=
[ e σ (1) · · · e σ ( m − j +1) ] .
Notice that the matrix
[Equation]
is S -consistent, since otherwise M is not.
We know that the cyclic group { P k | k ∈ Z } forms a subgroup of the group of permutations, and thus has finite order.
Thus, there exists ℓ ∈ N such that P ℓ = I , and thus P − 1 =
P ℓ − 1 [Artin, 2013,
Section 2.4].
Recall P − 1 =
P ⊤ since P is a permutation.
This means
[Equation]
is S -consistent, since it is a product of S -consistent matrices.
Notice how ( P ⊤ B ) i,i =
e ⊤ σ ( i ) B · ,i =
B σ ( i ) ,i ̸
= 0 .
In particular ( P ⊤ B ) 1 , 1 ̸ = 0 .
We can thus update M by applying matrix (72) to it to get a nonzero entry at ( j, j ) :
[Equation]
which completes the proof.
A.2.4 The consistency relation (Def. 7) is an equivalence relation
We start by showing a fact that will be useful to show that ∼ con is an equivalence relation.
Lemma 21.
Let S ∈ { 0 , 1 } m × n .
A matrix C ∈ R m × m is S
-consistent if and only if C is SP -consistent, where P is an n × n permutation matrix.
A matrix C ∈ R m × m is S
-consistent if and only if PCP ⊤ is PS -consistent, where P is a m × m permutation matrix.
When m = n, a matrix C ∈ R m × m is S
-consistent if and only if PCP ⊤ is PSP ⊤ -consistent, where P is a m × m permutation matrix.
Proof.
To show the first statement, we simply have to notice that
[Equation]
To show the second statement, we start with
[Equation]
The third statement, is a combination of the first two.
Proposition 22.
The consistency relation, ∼ con (Def. 7), is an equivalence relation.
Proof.
First, recall the fact that an intersection of subgroups is a subgroup.
This means that, the set of invertible matrices that are G z
-consistent
( G z ) ⊤ -consistent and G a -consistent is a group, and thus is closed under matrix multiplication and inversion.
Reflexivity.
It is easy to see that θ ∼ con θ , by simply setting L := I .
Symmetry.
Assume θ ∼ con ˜ θ .
Hence, we have G z = P ⊤ ˜
G z P and G a = P ⊤ ˜ G a as well as
[Equation]
where the matrix C is G z -consistent, ( G z )
⊤ -consistent and G a -consistent.
In order to show symmetry, we just need to show that the inverse of CP ⊤ can be written as ˜ C ˜ P ⊤ where ˜ P is some permutation and ˜ C is ˜
G z -consistent, ( ˜ G z )
⊤ -consistent and ˜ G a -consistent.
Notice that ( CP ⊤ ) − 1 = PC − 1 and that C −is consistent to G z , ( G z ) ⊤ and G a by closure under inversion.
Thus, by Lemma 21, we have that ˜ C :=
PC − 1 P ⊤ is ˜
G z -consistent, ( ˜ G z ) ⊤ -consistent, ˜ G a -consistent.
Hence
[Equation]
Transitivity.
Suppose θ ∼ con ˜ θ and ˜ θ ∼ con ˆ θ .
This means
[Equation]
P
[Equation]
G
[Equation] [Equation]
P
[Equation]
85) where C 1 is consistent to G z , ( G z ) ⊤ and G a ; and
[Equation]
and ˜ G
[Equation]
P
[Equation]
[Equation] [Equation]
[Equation]
x
[Equation]
, a
[Equation]
<
[Equation]
(88) where C 2 is consistent to ˜
G z , ( ˜ G z ) ⊤ and ˜ G a .
To show that θ ∼ con ˆ θ , we first combine (83) with (86) to get
G z = P ⊤ 1 P ⊤(cid:124) (cid:123)(cid:122) (cid:125) P ⊤ ˆ G z P 2 P 1 (cid:124) (cid:123)(cid:122) (cid:125)
P and G a
[Equation]
Moreover, we can combine (84) with (87) to get
[Equation]
and the same can be done for (85) and (88).
We must now show that C 1 P ⊤ 1 C 2 P ⊤ 2 = CP ⊤ where C is some matrix consistent to G z , ( G z ) ⊤ and G a (Def. 6).
Notice that
[Equation]
where ˆ C :
= P 2 C 2 P ⊤ 2 is consistent to ˆ G z , ( ˆ G z )
⊤
and ˆ G a , by Lemma 21.
We can further write
[Equation]
where C ′ :
= P ⊤ 1 P ⊤ 2 ˜ CP 2 P 1 is consistent to G z , ( G z ) ⊤ and G a , by Lemma 21 and (89).
Since C 1 is also consistent to G z , ( G z ) ⊤ and G a , the product C :
= C 1 C ′ also is, because of closure under multiplication (Thm. 20).
This concludes the proof that θ ∼ con ˆ θ .
A.2.5 Proof of Theorem
Finally, we can prove Thm. 8.
Note that its proof reuses many arguments initially introduced by Lachapelle et al.
In fact, the statement of Thm. 8 is identical to Thm. 5 of Lachapelle et al.
[2022] except for (i) the absence of the graphical criterion (Def. 23) and (ii) the conclusion, which is θ ∼ con
ˆ θ instead of θ ∼ perm ˆ θ .
App.
A.2.7 shows how Thm.
can be seen as a generalization of Thm.
5 from Lachapelle et al.
Theorem 8 (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 sufficient statistic T is d z -dimensional ( k = 1 ) and is a diffeomorphism from Z to T ( Z ) .
[Sufficient 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 G z ⊂
P ⊤ ˆ G z P .
Further assume that
[Sufficient action-variability] For all ℓ ∈ { 1 , ..., d a } , there exist { ( z ( p ) , a ( p ) , ϵ ( p ) , τ ( p ) ) } | Ch aℓ | p = belonging to their respective support such that
[Equation]
where Ch aℓ is the set of children of a ℓ and ∆ τℓ λ ( z <t , a <t , ϵ ) is a partial difference defined by
[Equation]
where ϵ ∈ R and E ℓ,τ ∈ R d a × t is the one-hot matrix with the entry ( ℓ, τ ) set to one.
Thus, (95) is the discrete analog of a partial derivative w.r.t.
a τℓ .
Then G a ⊂ P ⊤ ˆ G a .
Further assume that
[Equation]
Then, ˆ θ is consistent with θ , i.e. θ ∼ con
ˆ θ (Def. 7).
Proof.
First of all, since the assumptions of Thm. 4 hold, we have that θ
and ˆ θ are linearly equivalent.
Since k =(assumption 1), we can apply Lemma 10 to obtain the following equations:
[Equation]
, and
[Equation]
where we use the labelling of Sec.
A.2.1 with Λ functions.
Let us introduce S (1) , ˆ S (1) , S (2) and ˆ S (2) , the sparsity patterns of Λ (1) , ˆΛ (1) , Λ (2) and ˆΛ (2) , respectively.
As was hinted at in Sec.
A.2.1, the relationship between the sparsity patterns and the graphs is
[Equation]
Because of assumptions 2 & 3, we must have that
[Equation]
Notice how assumption 2 corresponds to assumption 1 of Lemma 18 and how assumption 3 corresponds to assumption 1 of
Lemma 19.
This means we can obtain the first conclusion of both Lemmas 18 & 19, i.e. that
S
[Equation]
⊂ P ⊤
[Equation] [Equation]
P , and S ( [Equation]
⊂ P ⊤ ˆ S
[Equation]
(101) which implies
[Equation]
All the above together with the sparsity assumption ( || ˆ G || 0 ≤
|| G || 0 ) allows to write
[Equation]
[By (98)] (103)
[Equation] [Equation]
[Equation] [Equation]
[By assumption 4 (Sparsity)] (106)
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
Since the l.h.s.
of (103) equals the r.h.s.
of (110), all the above inequalities are actually equalities.
Hence we have || ˆ S (1) ||
= || S (1) ||
and || ˆ S (2) ||
= || S (2) || 0 , (111) as well as
[Equation]
| 0 and ||
ˆ S
[Equation] [Equation]
[Equation]
The latter, combined with the r.h.s.
of (98) and (99), implies that
[Equation]
The equalities of (111) respectively implies the inequalities of the sparsity assumption of Lemmas 18 & 19, which allows us to obtain their second and most important conclusion i.e. that S (1) =
P ⊤ ˆ S (1) P , S (2) =
P ⊤ ˆ S (2) and that L = CP ⊤ where
C is S (1) -consistent, ( S (1) )
⊤ -consistent (Lemma 18) and S (2) -consistent
(Lemma 19).
Notice that because S (1) =
G z ,
S (2) =
G a , ˆ S (1) =
ˆ G z and ˆ S (2)
=
ˆ G a , these are equivalent to what we wanted to show, i.e. that θ ∼ con ˆ θ .
To gain a better understanding of sufficient time-variability and sufficient action-variability assumptions of Thm. 8, we provide examples of transition functions λ that do not satisfy them. The synthetic datasets used in our experiments are examples of processes satisfying the sufficient variability assumption, their exact form can be found in App. B.1.
For the sake of simplicity, assume the latent variables are Gaussian with a variance fixed to one, which implies that T is the identity.
Further assume that the system is Markovian, meaning λ ( z <t , a <t )
=
λ ( z t − 1 , a t − 1 ) .
The sufficient time-variability thus reduces to: There exist { ( z ( p ) , a ( p ) ) }
|| G z
|| 0 p
= belonging to their respective support such that
[Equation]
Now, assume λ ( z t − 1 , a t − 1 ) :
=
Wz t − 1 , with W ∈ R d z × d z G z .
This implies that D z λ ( z t − 1 , a t − 1 ) =
W , which clearly means that the sufficient time-variability assumption is not satisfied.
In this context, this assumption requires that λ is sufficiently nonlinear, in the sense that its Jacobian matrix varies sufficiently.
We postulate that this assumption is a reasonable one, given how complex real world dynamics can be.
Similarly, assume that λ ( z <t − 1 , a t − 1 ) =
Wa t − 1 , with W ∈ R d z × d a G a .
We thus have that
[Equation]
Unless every a ℓ has exactly one child, the sufficient action-variability assumption is violated, which, again, shows how linearity can cause problem.
A.2.7 Connecting to the graphical criterion of Lachapelle et al.
We now clarify how the graphical criterion of Lachapelle et al.
[2022], which guarantees complete disentanglement, is related to Thm. 8.
Let us first recall what this criterion is about.
Definition 23 (Graphical criterion of Lachapelle et al.
A graph G =
[ G z G a ] ∈ { 0 , 1 } d z × ( d z + d a ) satisfies the criterion of Lachapelle et al.
[2022] if, for all i ∈ { 1 , ..., d z } ,
[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ℓ is the set of children of a ℓ in G a .
We note that the above definition is slightly different from the original one, since the intersections run over Ch zi , Pa zi and Pa ai instead of over some sets of indexes
I , J ⊂ { 1 , ..., d z } and L ⊂ { 1 , ..., d a } .
This slightly simplified criterion is equivalent to the original one, which we now demonstrate for the interested reader.
Proposition 24.
Let G =
[ G z G a ] ∈ { 0 , 1 } d z × ( d z + d a ) .
The criterion of Def. 23 holds for G if and only if the following holds for G :
For all i ∈ { 1 , ..., d z } , there exist sets I , J ⊂ { 1 , ..., d z } and L ⊂ { 1 , ..., d a } such that
[Equation]
(cid:92)
[Equation]
Proof.
The direction “ = ⇒ ” is trivial, since we can simply choose I :=
Ch zi , J :=
Pa zi and L :
= Pa ai .
To show the other direction, we notice that we must have I ⊂ Ch zi , J ⊂ Pa zi and L ⊂ Pa ai , otherwise one of the sets in the intersection would not contain i , contradicting the criterion.
Thus, the criterion of Def.
23 intersects the same sets or more sets.
Moreover these potential additional sets must contain i because of the obvious facts that j ∈ Ch zi ⇐⇒ i ∈
Pa zj and ℓ ∈ Pa a i ⇐⇒ i ∈
Ch a ℓ , thus they do not change the result of the intersection.
We can now derive the fact that, if all assumptions of Thm. 8 and the graphical criterion of Def. 23 hold, then the learned representation will be completely disentangled:
Proposition 25 (Complete disentanglement as a special case) .
Suppose all assumptions of Thm. 8 and the graphical criterion of Def. 23.
Then, ˆ θ is completely disentangled, i.e. ˆ θ and θ are permutation-equivalent.
Proof.
Since assumptions of Thm. 8 holds, we have that θ
and ˆ θ are equivalent up to CP ⊤ where C is G z -consistent, ( G z )
⊤ -consistent and G a -consistent.
Using Lemma 17, we have that, for all i ,
[Equation]
Thus C is in fact a diagonal matrix, and hence ˆ θ is completely disentangled.
A.2.8 Interpreting the meaning of Theorem 8 and the ∼ con -equivalence
(Def. 7)
To interpret the conclusion of Thm. 8, which is that the learned model ˆ θ is consistent to the ground-truth model θ , i.e. θ ∼ con
ˆ θ (Def. 7), we recall the example introduced in Sec.
3.2.1: Consider the case where the ground-truth graphs G z =(no temporal dependencies) and G a is
[Equation]
=            , which does not satisfy the graphical criterion of Def. 23.
Then, θ ∼ con ˆ θ implies that: (i) ˆ G is the same as G , up to a permutation, and (ii) both representations f − 1 and ˆ f − 1 are linked by a linear transformation L = CP ⊤ (assuming T ( z ) :
= z for simplicity) where the matrix C is G z -consistent, ( G z )
⊤ -consistent and G a -consistent.
The conditions of G z -consistency and ( G z ) ⊤ -consistency are vacuous, since [ 1 − 0 ( 1 − 0 )
⊤ ] + = 1 , i.e. they do not enforce anything on C .
However,
G a -consistence forces C to have the same zeros as
[Equation]

[Equation]
11111111111   .
Lemma 17 gives a different perspective by telling us that C being G a -consistent is equivalent to having C i,j = 0 whenever j ̸∈ (cid:84) ℓ ∈ G ai, · G a · ,ℓ , which is equivalent to having j ∈ (cid:83)
ℓ ∈
Pa ai ( Ch aℓ ) c .
This allows us to see that the ground-truth factor z i is not a function of the learned factor ˆ z j ( C i,j = 0 ) whenever there exists an action a ℓ that targets z i , but not z j .
B EXPERIMENTS B.1 SYNTHETIC DATASETS
We now provide a detailed description of the synthetic datasets used in experiments of Sec. 4, which exactly match those of Lachapelle et al.
[2022], except for the graphs used.
We nevertheless provide a full description of the datasets used here for completeness.
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)
[Equation]
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, 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. 8.
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 = 10 and d a = 5 .
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.
Transition function of the time-sparsity datasets (left of Table 1).
The mean function in this case is given by
[Equation]
where G zi is the i th row of the ground-truth causal graph G z , the sin function is applied element-wise, the · is the dot product between two vectors and the summation in the sin function is broadcasted.
The various frequencies and phases in the sin functions ensures the sufficient time-variability assumption of Thm. 8 is satisfied.
Graphs of the datasets with temporal dependence (left of Table 1).
[Equation]

[Equation]  [Equation] 
[Equation]
 [Equation] 111111111111111    [Equation]
Transition function of the action-sparsity datasets (right of Table 1).
The mean function is given by
[Equation]
which is analogous to (122).
Graphs of the datasets with actions (right of Table 1).
G a (1) :
=         G a (2) : =
           (125)
B.2 IMPLEMENTATION DETAILS OF THE CONSTRAINED VAE APPROACH
All details of our implementation matches those of Lachapelle et al.
[2022] (except for the constrained optimization which is novel to our work).
We nevertheless repeat all details here for completeness.
Learned mechanisms.
Every coordinate z i of the latent vector has its own mechanism ˆ p ( z t i | z <t , a <t ) that is Gaussian with mean outputted by ˆ µ i ( z t − 1 , a t − 1 ) (a multilayer perceptron with 5 layers of 512 units) and a learned variance which does not depend on the previous time steps.
For learning, we use the typical parameterization of the Gaussian distribution with µ and σ 2 and not its exponential family parameterization.
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 first 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 .
Learned graphs ˆ G z
and ˆ G a .
As explained in Sec. 3.3, 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 fixed to 0 , i.e. it is not learned.
Analogously, in action-sparsity experiments, there is no temporal dependence so ˆ G z is fixed to 0 .
Encoder/Decoder.
In all experiments, including baselines, both the encoder and the decoder is modelled 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 .
Constrained optimization.
Let ELBO (ˆ f , ˆ λ , ˆ G, q ) be the ELBO objective evaluated on the whole dataset.
The constrained optimization we want to solve is
[Equation]
where ˆ G ∼ σ ( γ ) means that ˆ G i,j are independent and distributed according to σ ( γ i,j ) .
Because E ˆ G ∼ σ ( γ ) ||
ˆ G
||
= ||
σ ( γ ) ||where σ ( γ ) is matrix, the constraint becomes || σ ( γ ) || 1 ≤ β .
To solve this problem, we perform gradient descent-ascent on the Lagrangian function given by
[Equation]
(127) where the ascent step is performed w.r.t.
ˆ f , ˆ λ , ˆ G and q ; and the descent step is performed w.r.t.
Lagrangian multiplier α , which is forced to remain greater or equal to zero via a simple projection step.
As suggested by Gallego-Posada et al.
[2021], we perform dual restarts which simply means that, as soon as the constraint is satisfied, the Lagrangian multiplier is reset to 0 .
We used the library Cooper [Gallego-Posada and Ramirez, 2022], which implement many constrained optimization procedure in Python, including the one described above.
Note that we use Adam [Kingma and Ba, 2015] for the ascent steps and standard gradient descent for the descent step on the Lagrangian multiplier α .
We also found empircally that the following schedule for β is helpful:
We start training with β = max
ˆ G
|| G
||
and linearly decreasing its value until the desired number of edges is reached.
This avoid getting a sparse graph too quickly while training, thus letting enough time to the model parameters to learn.
In each experiment, we trained for 300K iterations, and the β takes 150K to reach to go from its initial value to its desired value.
B.3 DETAILS ABOUT THE R con METRIC AND ITS RELATION TO MCC AND R
To evaluate whether the learned representation is consistent to the ground-truth (Def. 7), as predicted by Thm. 8, we came up with a novel metric, denoted by R con .
Computing R con goes as follows:
First, we permute the learned representations ˆ z using the permutation
ˆ P found by MCC (Sec. 4), i.e. ˆ z perm :
= ˆ P ⊤ ˆ z .
Then, we compute the sparsity pattern imposed by the consistency equivalence (Def. 7), denoted by S C .
Then, for every i , we predict the ground-truth z i given only the factors allowed, i.e. ( S C ) i, · ⊙ ˆ z perm , and compute the associated coefficient of multiple correlations R con ,i and report the mean, i.e.
R con :
= 1 d z (cid:80) d z i =1 R con ,i .
It is easy to see that we must have R con ≤ R , since R con was computed with less features than R .
Moreover, MCC ≤ R con , because MCC can be thought of as computing exactly the same thing as for R con , but by predicting z i only from ˆ z perm ,i , i.e. with less features than R con .
This means we always have 0 ≤ MCC ≤ R con ≤ R ≤ 1 .
This is a nice property which allows to compare all three metrics together.