E FFICIENTLY D ISENTANGLE C AUSAL R EPRESENTATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we propose a novel approach to efﬁciently learning disentangled representations with causal mechanisms, based on the difference of conditional probabilities in original and new distributions.
We approximate the difference with model’s generalization abilities so that it ﬁts in standard machine learning framework and can be efﬁciently computed.
In contrast to the state-of-the-art approach, which relies on learner’s adaptation speed to new distribution, the proposed approach only requires evaluating the generalization ability of the model.
We provide theoretical explanation for the advantage of the proposed method, and our experiments show that the proposed technique is 1.9–11.0 × more sample efﬁcient and 9.4–32.4 × quicker than the previous method on various tasks.
The source code is in supplementary material.
I NTRODUCTION
Causal reasoning is a fundamental tool that has shown great impact in different disciplines (Rubin &
Waterman, 2006; Ramsey et al., 2010; Rotmensch et al., 2017), and it has roots in work by David
Hume in the eighteenth century (Hume, 2003) and in classical AI (Pearl, 2003).
Causality has been mainly studied from a statistical perspective (Greenland et al., 1999; Pearl, 2009; 2018) with Judea
Pearl’s work on the causal calculus leading its statistical development.
More recently, there has been a growing interest to integrate statistical techniques into machine learning to leverage their beneﬁts.
Welling raises a particular question about how to disentangle correlation from causation in machine learning settings to take advantage of the sample efﬁciency and generalization abilities of causal reasoning (Welling, 2015).
Although machine learning has achieved important results on a variety of tasks like computer vision and games over the past decade (e.g.,
Mnih et al.
; Silver et al.
(2017); Szegedy et al.
(2017); Hudson & Manning (2018)), current approaches can struggle to generalize when the test data distribution is much different from the training distribution (common in real applications).
Further, these successful methods are typically “data hungry”, requiring an abundance of labeled examples to perform well across data distributions.
In statistical settings, encoding the causal structure in models has been shown to have signiﬁcant efﬁciency advantages.
For example, in the context of speech recognition, Chorowski et al.
(2014) and
Bengio (2017) show that a model trained to learn P ( acoustics ) performs worse than a model trained to also learn a conditional distribution containing phonemes, P ( acoustics | phonemes ) · P ( phonemes ) .
The model performs better when the causality—phonemes cause acoustics—is encoded in the model.
Further, Schölkopf et al.
(2012) show that encoding the causal structure leads to more robust estimation (and thus, the performance improves).
Independent of any change in P ( phonemes ) , the causal mechanism P ( acoustics | phonemes ) is not expected to change across tasks or datasets.
These robustness and generalization characteristics lead to quicker training, because once the causal mechanism is learned, it requires minimal updates to transfer to new data.
In support of the advantages of encoding causal mechanisms, Bengio et al.
(2020) recently introduced an approach to disentangling causal relationships in end-to-end machine learning by comparing the adaptation speeds of separate models that encode different causal structures.
The core idea is explained with an atmospheric physics example of learning P ( A, B ) , where A ( Altitude ) causes
B ( Temperature ) (Peters et al., 2016).
The marginal probability of A can be changed depending on, for instance, the place (Switzerland to less mountainous country like Netherlands), but P ( B | A ) as the underlining causality mechanism does not change.
Therefore, the causality structure can be inferred from the robustness of predictive models on out-of-distribution (Peters et al., 2016; 2017).
The invariance of P ( B | A ) across train and transfer distributions is the key assumption in their work, and we also follow it in this work.
P ( A, B ) can be either modeled as P ( B | A ) · P ( A ) or P ( A | B ) · P ( B ) , reﬂecting the possible causal directions.
Since P ( B | A ) · P ( A ) reﬂects the right causal structure, it is more efﬁcient to adapt to a change in P ( A ) since parameters for P ( B | A ) need not to be updated.
On the other hand, the model that assumes the opposite causal direction, P ( A | B ) · P ( B ) , must adapt both its marginal and conditional distribution parameters, leading to slower adaptation speed.
The previous method proposed the state-of-the-art meta-learning approach to disentangling causal mechanisms using the adaptation speed as a meta-objective to optimize.
In particular, the task is to determine whether the causal direction is A → B or B → A .
In this work , we propose a more efﬁcient approach to disentangling causal mechanisms than the state-of-the art one based on the same assumption.
The baseline approach essentially captures the changes in P ( B | A ) and P ( A | B ) , by examining how fast the joint probability models can adapt or learn after intervention.
In contrast, our approach only evaluates the generalization ability of the model without the need of the adaptation process, therefore it is much quicker in disentangling the causal mechanisms.
Our experiments show the proposed technique is 1.9–11.0 × more sample efﬁcient and 9.4–32.4 × quicker than measuring adaptation speed on various tasks.
We also argue that the proposed approach has less hyper parameters and it is straight-forward to implement the approach within the standard machine learning workﬂows.
Our contributions can be summarized as follows.
• We propose an efﬁcient approach to disentangling representations for causal mechanisms by measuring generalization.
• We theoretically prove that the proposed estimators can identify causal direction and disentangles causal mechanisms.
• We empirically show that the proposed approach is signiﬁcantly quicker and more sample efﬁcient for various tasks.
Sample efﬁciency is important when data size in transfer distribution is small.
A PPROACH
To begin, we reﬂect on the tasks and the disentangling approach (as baseline) described by previous work (Bengio et al., 2020).
The invariance of conditional distribution for the correct causal direction
P ( B | A ) is the key assumption in their work, and we also follow it in this work.
We notice that their baseline approach compares the adaptation speed of models on a transfer data distribution, and hence requires signiﬁcant time for adaptation.
We propose an approach to learn causality mechanisms by directly measuring the changes in conditional probabilities before and after intervention for both
P ( B | A ) and P ( A | B ) .
Further, we optimize the proposed approach to use generalization loss rather than a divergence metric—because loss can be directly measured in standard machine learning workﬂows—and we show that it is likely to correctly predict causal direction and disentangle causal mechanisms.
2. C AUSALITY D IRECTION P REDICTION
As the ﬁrst step towards learning disentangled representations for causal mechanisms, we start with the binary classiﬁcation task.
Given two discrete variables A and B , we want to determine whether A causes B , or vice-versa.
The training (transfer) data contains samples ( a, b ) from training (transfer) distribution, P 1 ( P 2 ).
The baseline approach deﬁnes models that factor the joint distribution P ( A, B ) into two causality directions
P A → B ( A, B ) =
P A → B ( B | A ) P
A → B ( A ) and
P B → A = P B
→ A ( A | B ) P B
→ A ( B ) .
It then compares their speed of adaptation to transfer distribution.
Intuitively, the factorization with correct causality direction should adapt more quickly to the transfer distribution.
Suppose A → B is the ground-truth causal direction.
For the correct factorization, they assume that conditional distribution P A → B ( B | A ) does not change between the train and transfer distributions, so that only the marginal P A → B ( A ) needs adaptation.
In contrast, for the factorization with incorrect causality direction, both the conditional P B → A ( A | B ) and marginal distributions
P B → A ( B ) need adaptation.
They analyze that updating a marginal distribution P ( A )
A → B is likely to have lower sample complexity than the conditional distribution P B → A ( A | B ) (only a part of
P B → A ( A, B ) )
because the later has more parameters.
Therefore, the model with correct factorization will adapt more quickly, and causality direction can be predicted from adaptation speed.
To leverage this observation, the baseline method deﬁnes a meta-learning objective.
Let L A → B and L B → A be the log-likelihood of P A → B ( A, B ) and
P B
→ A ( A, B ) , respectively.
Their baseline approach optimizes a regret, R , to acquire an indicator variable γ ∈ R .
If γ > 0 , the prediction is
A → B , otherwise B → A .
[Equation]
2. P ROPOSED M ECHANISM : O BSERVE THE M ODEL ’ S C ONDITIONAL D ISTRIBUTION D IVERGENCE
Rather than relying on gradients of end-to-end predictions to identify causal direction, we propose to directly observe the divergence of each model’s conditional distribution predictions under the intervention of the transfer dataset.
We consider the distribution of a single model’s predictions on the training and transfer distributions.
These distributions depend on both the marginal and conditional distributions learned by the models.
It is expected that a model encoding the correct causal direction will have learned a correct conditional distribution, but the marginal distribution might shift from the training to transfer datasets.
Hence, we would like to ignore changes in the model’s predictions that are due to marginal distribution differences and focus directly on the conditional distribution.
Given access to the model’s conditional distribution predictions, we could use the KL-divergence to directly measure the conditional distribution differences between the training and transfer datasets.
A model that encodes the correct causal structure should show no change in its conditional distribution predictions, so the divergence of these predictions should be small.
On the other hand, a model that encodes the incorrect causal structure is likely to show signiﬁcantly large divergence.
We tie these observations together in the following proposition:
Proposition 1.
Given two data distributions with the same directed causality between two random variables, A and B , the difference of the KL-divergences of their conditional distributions is an unbiased estimator of the correct causality direction between A and B .
This approach shares the same intuition as the baseline approach.
The model with correct causal direction should only witness minimal changes in the conditional distribution of its predictions between train and transfer distributions.
Thus, the KL-divergence of the predictions should not change for models with correct causal structure.
The proof for Proposition 1 is included in Appendix A.
2. P ROPOSED S IMPLIFICATION : G ENERALIZATION L OSS A PPROXIMATES THE D IVERGENCE
Although we can use the KL-divergence of the conditional probabilities of two data distributions as an unbiased estimator of causality direction, it requires additional computation to model conditional distribution in transfer distribution.
However, in many practical end-to-end learning settings, it is more efﬁcient and straight-forward to acquire generalization loss, such as the cross-entropy, that can be used to approximate the conditional KL-divergence.
More speciﬁcally, we compare the generalization gaps between two causal models to approximate the causality direction.
Let the generalization gap, G , be the difference of a model’s losses on the training and transfer datasets:
G A → B = L transfer A → B − L train A → B .
Here, L · A → B is the loss on the speciﬁed set.
Further, we deﬁne a directionality score as the difference of generalization gaps:
[Equation]
We show that, for an appropriately chosen loss, such as cross-entropy, S G is a biased but reasonable estimator of the correct causal direction.
When S G > 0 , it is likely that A → B is the correct causal direction–the generalization gap for the incorrect causal model dominates the score.
We formalize this notion in Proposition 2 in Appendix B:
Proposition 2.
Given two data distributions with the same directed causality between two random variables, A and B , the difference of their generalization gaps is a biased estimator of the correct causality direction between A and B .
In practice for the tasks we test in next section, we ﬁnd that although S G is biased, it always indicates the correct causal direction.
An intuitive understanding is that this approach measures how well models trained on train distribution can predict in transfer distribution–their generalization ability.
Algorithm 1 summarizes the process for identifying a correct causal model.
In Appendix C, we describe that the generalization-based approach should converge more quickly than a gradient-based approach, especially in practical settings. 1: Train f
A → B on training data, and get train loss L train A → B . 2: Train f B → A on training data, and get train loss L train B → A . 3: Get transfer loss L transfer
A → B with f A → B on transfer data.
4: Get transfer loss L transfer B
→ A with f B → A on transfer data. 5: Get generalization loss G A → B = L transfer A → B − L train A → B . 6:
Get generalization loss G B → A = L transfer B → A − L train B → A . 7: Compute S G = G B → A − G A → B . 8: If S G > 0 return A → B , else return return B → A . The baseline approach and our proposed approach are based on the same intuition of unchanging conditional distributions for correct causality direction.
However, the baseline approach must observe gradients during transfer distribution training, while the proposed approach looks directly at the changes in model outputs on the transfer data distribution.
Informally, we note that these approaches are roughly equivalent by observing the following: For the correct causal direction,
[Equation]
Since ∇L train = 0 , ∇G ·→· = ∇L transfer ·→· .
The previous work shows that ∇ θ i L transfer ·→· = 0 for model parameters, θ i , representing correct causal structures (i.e., joint distribution), so ∇L transfer ·→· should also be small.
Similar to that, we expect G ·→· to be small relative to the incorrect causal model.
2. R EPRESENTATION L EARNING
So far, we assume that the causal variables are already disentangled in the raw inputs.
We are interested in extending it to situations where the variables are entangled in input, such as in pixels or sounds.
Here, we want to learn representations that disentangle causal variables.
Following the previous work, we suppose the true causal variables ( A, B ) generate input observation ( X, Y ) with ground truth decoder, D .
Then, we want to train an encoder, E , that converts input ( X, Y ) to a hidden representation, ( U, V ) .
The decoder and encoder are rotation matrices.
[Equation]
Then we treat ( U, V ) as observed input, and build causal modules on them in a way similar to causality direction prediction problem (Section 2.1).
If the encoder is learned correctly, we expect to obtain ( U, V ) =
( A, B ) with θ E = − θ D , or ( U, V ) =
( B, A ) with θ E = θ D .
The baseline approach addresses this problem with loss function in Equation 1, but extending meta-parameters to encoder parameters in the meta-learning framework.
We apply our method to learn representation without using adaptation process.
Suppose the conditional distributions P
( V | U ) is modeled by f U → V , and P ( U | V ) by f V → U .
Also, the original losses are L U → V for f U → V and L V → U for f V → U on train data.
We optimize the following objective.
[Equation]
λ is a hyper parameter for interpolation.
Note that we compute L U → V , L V → U with train data, and
G U → V , G V → U with transfer data.
When we use mini-batch, this means that there are two types of iteration do 5: end for mini-batches for data from train and transfer datasets, respectively.
Train data is used to learn f U → V and f V → U , and transfer data is used to learn encoder θ E .
Please see Algorithm 2 for details.
An intuitive understanding is that we want to learn a representation that works well in both causal directions in train distribution, and at least one causal direction in transfer data distribution.
Therefore, the models are trained well in train distribution, and the representation recovers causal variables, because it works on both distributions in at least one causality direction.
E XPERIMENTS
We run experiments to show that predicting causal direction and disentangling causal representations can be more efﬁcient with the proposed approach. For fair comparison, we focus on the experimental setups in the previous work.
3. C AUSALITY D IRECTION P REDICTION
[Figure]
[Figure] [Figure] [Figure]
[Figure]
We ﬁrst compare the approaches in the same meta-learning setting.
For the proposed approach, we replace log-likelihood of joint probabilities with the generalization loss.
The overall loss for the proposed approach R (cid:48) is similar to Equation (1) with L A → B = G A → B and L B → A = G B → A .
[Equation]
The ﬁrst case we consider is the discrete model with both A and B being 10 dimensional variables ( N = 10 ).
For other cases, please refer to Appendix E.
The model architecture, same as in the baseline, has four modules: P A → B ( A ) , P A → B ( B | A ) , P B → A ( B ) ,
P B
→ A ( A | B ) .
The baseline uses all of them, and the proposed approach only uses the conditionals.
We parameterize these probabilities via Softmax of unnormalized tabular quantities, and train the modules independently with maximum likelihood estimator.
We then predict causality direction on transfer data with each approach.
For each approach, we ﬁx ground-truth causal direction as A → B , and run 100 experiments with different random seeds.
We plot the mean value of σ ( γ ) in Figure 1a.
The proposed method (green) is above the baseline method (orange).
This comparison shows that using generalization loss is more efﬁcient than the log-likelihood of joint probabilities.
Especially in the ﬁrst several steps, the proposed approach moves toward the correct γ more quickly than the baseline approach.
This quicker convergence is important when the transfer data is expensive.
We also observe that for the baseline method, σ ( γ ) is less than 0.5 for the ﬁrst several steps, which might be because the factorization with incorrect causality direction has more parameters to be updated, leading to faster adaptation.
Since the proposed approach does not need an adaptation process, it actually does not need to update σ ( γ ) with gradients.
Instead, we can calculate a score, S G = G B
→ A − G A → B , to detect causality direction.
This score is not directly comparable with σ ( γ ) in baseline approach, because they have different scales and semantics.
To compare the approaches, we reformulate the causality direction prediction problem as a binary classiﬁcation problem, and evaluate accuracy of the classiﬁcation.
For each approach, we run 100 experiments with the same setting as above.
We then compute the accuracy as the number of experiments with successful prediction over the number of all experiments.
For the baseline approach, we count an experiment as successful when γ > 0 .
For the proposed approach, we count when S G > 0 .
We plot accuracy along with the number of iterations to compare sample efﬁciency as shown in
Figure 1b.
The experimental result demonstrates that the proposed approach (green) achieves 100% accuracy after just 3 iterations with 5.1 ms, while the baseline approach (orange) takes 33 iterations to achieve 100% accuracy with 165.1 ms.
Since our method only requires estimating a scale value of loss from transfer data, so that it is more sample efﬁcient than the original method that requires learning / updating all parameters in a model during adaptation.
The experiments show that our method is 11.0 times sample efﬁcient and 32.4 times quicker compared with the baseline approach.
Also, the baseline accuracy for the ﬁrst several steps is less than 50%.
This is consistent with the previous observation that the σ ( γ ) < 0 .
5 for the ﬁrst several steps.
3. R EPRESENTATION L EARNING
[Figure]
[Figure]
The curve of the proposed approach green is above that of baseline approach (orange), the proposed have both better sample
We evaluate representation learning in the same setting as in previous work.
We set the correct causal direction as A → B .
With a sample ( A, B ) , a ﬁxed decoder D =
R
( θ D ) converts it to an observation ( X, Y ) , where R is a rotation matrix.
We then use both approaches to learn encoders E
= R ( θ E ) that map ( X, Y ) to representation ( U, V ) .
We set decoder parameter θ D = − π/ 4 , so that the expected encoder parameter should be π/ 4 or − π/ 4 .
For both approaches, we use the same experiment setting and hyper parameters.
In each iteration, the baseline approach uses transfer samples to adapt the model by learning gradient descent for several steps, and obtain regret to update encoder.
On the other hand, we directly use loss instead of analyzing with statistical test, because this can be neurally integrated with neural network training for disentangled representation learning.
We plot the relation of θ values along with the number of iterations and time to compare sample and time efﬁciency.
The result for sample efﬁciency is in Figure 2a.
The proposed approach (green) reaches near π/ 4 after around 220 iterations.
Baseline approach (orange) reaches near π/ 4 after around 410 iterations.
The proposed approach is around 1.9 times sample efﬁcient compared with the baseline approach.
The result for time efﬁciency is in Figure 2b.
The proposed approach (green) reaches near π/ 4 in around 34 seconds.
Baseline approach (orange) reaches near π/ 4 in aroundseconds.
The proposed approach is 9.4 times quicker than the baseline approach.
The results show that proposed approach is both sample and time efﬁcient for the representation learning task.
Especially, the time efﬁciency shows the proposed method’s advantage that it does not need multiple steps of adaptation in each iteration.
It even does not need to compute gradient with back propagation or update model parameters.
D ISCUSSIONS
To further understand the proposed approach, we discuss alternative metrics for distribution differences, and the working conditions of the proposed approach and the baseline.
We also extend experiments on real data and neural network models.
Alternative Metrics
The proposed approach may even work when using other metrics that are easily accessed in standard machine learning workﬂows.
Many loss functions share properties of distance metrics and smoothness conditions that make them useful for comparing the generalization ability of different models.
Such metrics may be biased but can also be used to assess causality direction.
We describe and experiment with metrics such as the KL-divergence loss and the gradient norm in Appendix F. Empirical tests show that both alternative metrics can identify causality direction in the prediction task.
One may choose these alternative metrics depending on the workﬂow.
Robustness We compare the proposed and the baseline approach on the working conditions.
One potential problem of the baseline approach is the dependency on marginal distributions.
This not only requires more samples and computation, but when the marginal distribution for the correct factorization P ( A ) is complicated, it may produce incorrect causality predictions (suppose A → B is the correct causality direction).
We show an example for such a case in Appendix G. We design marginal distribution of cause variable P ( A ) with a complicated hidden structure, which makes it harder to adapt after intervention than the marginal distribution of effect variable P ( B ) , and conditional distribution P ( A | B ) .
We use a mixture model for P ( A ) with K dimensional hidden variables C , and do not change models for other distributions.
The result shows that the proposed approach is more robust than the baseline one in this case.
Real Data
In addition to comparing the proposed approach with the previous work in the same setting to understand fundamental mechanism for causality learning, we also evaluate whether the propose method extends naturally to broader situations.
We use real data of temperature and altitude (Mooij et al., 2016).
We apply the proposed approach with both linear models and neural networks.
We also add noise to the data and observe the changes on performance.
The details can be found in Appendix H.
The results show that the proposed approach works for the real data, and also applies for both linear models and neural networks.
It resists some level of noise (standard deviation ≤ 1 ), but not when the noise is too large.
R ELATED W ORK
Our work is closely related to causal reasoning, disentangled representation learning and its application to domain adaptation and transfer learning.
Causal Reasoning There is a rich literature on causal reasoning based on pre-conceived formal theory since Pearl’s seminal work on do-calculus (Pearl, 1995; 2009).
The major approach is structure learning in Bayesian networks.
Discrete search and simulated annealing are reviewed in (Heckerman et al., 1995).
Minimum Description Length (MDL) principles are commonly used to score and search (Lam & Bacchus, 1993; Friedman & Goldszmidt, 1998).
Bayesian Information Criterion
(BIC) is also used to search when models have high relative posterior probability (Heckerman et al., 1995).
Prior work uses observation, but does not consider interventions, and focus on likelihood learning or hypothesis equivalence classes (Heckerman et al., 1995).
Later, approaches have also been proposed to infer the causal direction only from observation (Peters et al., 2017), based on not generally robust assumptions on the underlying causal graph.
The impact of interventions on probabilistic graphical models is also introduced in (Bareinboim & Pearl, 2016).
Another early work (Peters et al., 2016) uses intervention to detect causal direction, but it does not address representation learning, which is an important problem in modern neural network, and the main topic of this paper.
Further, our method should be more efﬁcient than the fast approximation method (Peters et al., 2016), which needs to ﬁt a model and takes many iterations.
In contrast, we only run one forward pass on a ﬁxed model.
Recently, meta-learning approach is adopted to draw causal inferences from purely observational data.
Dasgupta et al.
(2019) focuses on training an agent to learn causal reasoning by meta reinforcement learning; Bengio et al.
(2020) focuses on predicting causal structure based on how fast a learner adapts to new distributions.
In contrast, we propose an efﬁcient approach for disentangling causal mechanisms, based on the difference of conditional probabilities in original and new distributions, assuming causal mechanisms hold in both distributions.
More recently, generative models have been augmented with causal capabilities to learn to generate images and also to plan (Kocaoglu et al., 2018; Kurutach et al., 2018).
Disentangled Representation Learning This work discovers the underlying causal variables and their dependencies.
This is related to disentangling variables (Bengio et al., 2013) and disentangled representation learning (Higgins et al., 2017).
Some work point out that assumptions, including priors or biases should be necessary to discover the underlying explanatory variables (Bengio et al., 2013;
Locatello et al., 2018).
Another work (Locatello et al., 2018) reviews and evaluates disentangling and discusses different metrics.
A strong assumption of disentangling is that the underlying explanatory variables are marginally independent from each other.
Many deep generative models (Goodfellow et al., 2016) and independent component analysis models (Hyvärinen et al., 2004; Higgins et al., 2017; Hyvarinen et al., 2018) are built on this assumption.
However, in realistic cases, this assumption is often not likely to hold.
Some recent work addresses compositionality learning (Li et al., 2019; 2020) without such assumption, but they do not address how to use causal mechanisms for representation learning.
Domain Adaptation and Transfer Learning Previous work ﬁnds a subset of features to best predict a variable of interest (Magliacane et al., 2018).
However, the work is about feature selection, and our work is about causality direction and representation learning.
Other work also examine incorrect inference and turn the problem into a domain adaptation problem (Johansson et al., 2016).
Counterfactual regression is proposed to estimate each effect from observation (Shalit et al., 2017).
Another approach is to ﬁnd a subset which makes the target independent from the variable selection (Rojas-Carulla et al., 2018).
To do that, they assume the invariance of conditional distribution in train and target domains.
Also, competition may be introduced to recover a set of independent causal mechanisms, which leads to specialization (Parascandolo et al., 2017).
C ONCLUSION
In this paper, we propose an efﬁcient simpliﬁcation to a recent technique to disentangle causal representations.
Unlike the baseline approach which requires signiﬁcant adaptation time and sophisticated comparisons of the models’ adaptation speed, the proposed approach directly measures generalization ability based on the divergence of conditional probabilities from the original to transfer data distributions.
We provide theoretical explanation for the advantage of the proposed method, and our experiments demonstrate that the technique is signiﬁcantly more efﬁcient than adaptation speed based approach in causality direction prediction and representation learning tasks.
We hope this work in causal representation learning will be helpful for more advanced artiﬁcial intelligence.
A P ROOF THAT KL- DIVERGENCE C AN B E U SED TO I DENTIFY C AUSALITY D IRECTION
Here, we prove that we could use conditional KL-divergence to detect causality direction.
Proposition 1.
Given two data distributions with the same directed causality (e.g., A → B ) between two random variables, A and B , the difference of the KL-divergences of their conditional distributions (i.e., P ( B | A ) ) is an unbiased estimator of the correct causality direction.
Proof.
First, let P 1 and P 2 be the distributions such that P 1 (cid:54) =
P 2 , but their conditional distributions are the same.
Let D KL ( ·(cid:107)· ) be conditional KL-divergence.
CrossEntropy A → B ( P i , P j ) is the conditional cross-entropy, − (cid:80) ∀ (
a,b ) P i ( a, b ) log P j ( b | a ) =
E A,B ∼ P i [ − log P j ( B | A )] .
Let H A
→ B ( · ) be conditional entropy.
By the Kraft–McMillan theorem, we have
[Equation]
Without loss of generality, suppose A → B is the correct causality direction.
Since P 1 ( B | A ) =
P 2 ( B | A ) , we have that CrossEntropy
A → B ( P 2 , P 1 ) =
H A → B ( P 2 ) and
[Equation]
CrossEntropy
[Equation]
→
[Equation]
P
[Equation]
P
[Equation]
Finally, if we deﬁne a directionality score, S D KL , as the difference of KL-divergence of the two distributions, we can identify the incorrect causal direction by its greater divergence.
Namely,
[Equation]
Here we use the property that entropy is less than cross entropy when the two distributions are not equal.
B P ROOF THAT G ENERALIZATION L OSS C LOSELY A PPROXIMATES D IVERGENCE
Here, we prove that we can replace conditional KL-divergence with a standard generalization loss calculation to approximate the causal directionality score, S D KL , in Proposition 1.
This approximation introduces bias into the score, but we expect this bias to be small relative to the difference of conditional KL-divergences of the incorrect causal model.
Thus, it can be used as a good estimator of causality direction.
Proposition 2.
Given two different data distributions, P 1 and P 2 , with the same directed causality between random variables, A and B , consider the generalization loss deﬁned as the conditional cross-entropy of a given distribution:
[Equation]
And deﬁne the generalization gap as G ·→· = L P 2 ·→· − L P 1 ·→· .
Finally, deﬁne the modiﬁed causality score, S G = G B
→ A − G A → B .
If S D KL = D B
→ A − D A → B as deﬁned in Proposition 1, then S
G = S D KL − [∆ H ( B ) − ∆ H ( A )] .
S G is a biased estimator of which model has correct causality direction.
Proof.
Let H ( · , · ) be the joint entropy.
The generalization gap is
[Equation]
Here ∆ H ( A, B ) , ∆ H ( A ) and ∆ H ( B ) are the change in entropy from P 1 to P 2 for ( A, B ) , A and
B , respectively.
Further
[Equation]
the modiﬁed causality score is the difference of these generalization gaps:
[Equation]
Without loss of generality, suppose A → B is the correct causal direction, then D A → B = 0 .
[Equation]
When S G > 0 (and [∆ H ( B ) − ∆ H ( A )] is small), the prediction of causality direction is correct.
Informally, we expect the bias term, [∆ H ( B ) − ∆ H ( A )] , to be small relative to the divergence,
D B → A , in real application settings.
If A → B , then A and B will have some (positive or negative) correlation–they will vary in a correlated way under intervention on A .
Unless the conditional distribution, P 1 ( B | A ) =
P 2 ( B | A ) , contains large portions with high entropy, it is likely that a change in A will cause a commensurate change in B .
Thus, it is likely that | ∆ H ( A ) | ≈ | ∆ H ( B ) | , so it is unlikely that ∆ H ( A ) and ∆ H ( B ) are both large and have opposite sign.
Further, in cases where ∆ H ( A ) is large, but ∆ H ( B ) is small, it is frequently an indicator that the selected P i are “lucky” (uncommon) or the causality between A and B is weak.
Consider the example in which H ( P 1 ( A )) is small, but H ( P 1 ( B )) is large.
This case can occur if H ( P 1 ( B | A )) is small, but it is unlikely, because the (low-entropy) conditional distribution must impose the disorder on
P 1 ( B ) .
It is easy to see that the reverse situation ( H ( P 1 ( B | A )) is large) implies that the causality between A and B is weak, such that causality direction prediction may be difﬁcult anyway.
Finally, if we want to relax the assumption that either A → B or B → A (e.g., a null-hypothesis that
A and B are not causally related), we could introduce a signiﬁcance test to decide whether we can accept the selection of causal direction from either method.
Such a signiﬁcance test would need to be mindful of the magnitude of any biases introduced in the approach.
C G ENERALIZATION L OSS IS L IKELY TO BE M ORE P RACTICAL
Here, we explain how generalization loss is likely to be more practical than gradient-based approaches for identifying a model with correct causal direction.
Both the baseline and proposed approaches rely on convergence of estimators toward a ground truth model, but gradient-based approaches require much stronger assumptions to guarantee convergence, and they are more susceptible to practical challenges like tuning optimizer hyperparameters.
First, both gradient- and generalization-based approaches have the same asymptotic convergence characteristics, with the prediction error proportional to 1 √ n , where n is the number of steps or samples.
Many theorems state that gradient-based approaches can have asymptotic convergence to within an ε error based on number of gradient descent steps, k : ε =
O ( 1 √ k ) .
These proofs rely on many assumptions that cannot be guaranteed in real-world datasets, such as convexity, regularity, or smoothness of error surface, using optimal settings of optimization parameters for best-case step sizes, etc.
(Kuhn & Tucker, 1951; Nesterov, 2004; Yin et al., 2018).
These bounds can also tend to be loose for real data distributions.
On the other hand, non-parametric estimators, such as the sample average generalization loss proposed in this work, do not have a model, so their convergence rates can depend only on the structure of the dataset (e.g., ﬁrst or second moments).
For instance, early estimator convergence results show that after summarizing n samples, the error in an estimator can be within ε of the true value:
ε =
O ( 1 √ n ) (Bernstein, 1946; Hoeffding, 1963).
The constants in these bounds are just the moments that characterize the dataset (e.g., mean, variance).
Like the gradient-based convergence rates, these non-parametric convergence rates can also be loose for real data distributions.
Thus, in practical settings, it is likely that a generalization-based approach (non-parametric estimators) will converge more quickly on causality identiﬁcation in real-world datasets.
The proposed approach does not require further model optimization (ﬁne-tuning), so it is not susceptible to poorly chosen hyperparameters.
Further, it is unlikely to encounter complex error surface topology.
Thus, in practice, it is likely to converge more quickly than gradient-based causality identiﬁcation techniques.
D R EPRESENTATION L EARNING
We run representation learning experiments for both approaches on a GeForce GTX TITAN Z with
PyTorch.
We use the original implementation for the baseline approach, and extend it for the proposed approach.
The hyperparameters are the same as the original setting (Bengio et al., 2020).
E M ORE E XPERIMENTS ON C AUSALITY D IRECTION P REDICTIONS
[Figure]
[Figure] [Figure]
We would like to extend causality direction prediction experiments to high dimensional variable and continuous variables.
We still follow the experiment settings in previous work (Bengio et al., 2020), and compare the proposed and the baseline approaches.
We run all the experiments of causality direction prediction with Intel Core i9 9920X 12-Core Processors.
For high dimensional variable, we run experiments with N = 100 .
We use the same setting as the experiments with N = 10 , but only changed the N .
The result for sample efﬁciency is in Figure 3a.
Similar to the case of N = 10 , the proposed approach (green) is signiﬁcantly more efﬁcient than the baseline approach (orange) in both sample.
For continuous variable, we used the same setting as the previous work (Bengio et al., 2020).
For the proposed method, we only use one sample for each episode.
The plot shows that the proposed approach (green) is more efﬁcient than the baseline approach (orange) in sample number (Figure 3b).
F A LTERNATIVE M ETRICS
The proposed approach may even work when using other metrics that are easily accessed in standard machine learning workﬂows.
Many loss functions share properties of distance metrics and smoothness conditions that make them useful for comparing the generalization ability of different models.
Such metrics may be biased but can also be used to assess causality direction.
We describe and experiment with metrics such as the KL-divergence loss and the gradient norm.
Empirical tests show that both alternative metrics can identify causality direction in the prediction task.
One may choose these alternative metrics depending on one’s workﬂow.
For KL-divergence, we deﬁne the metric as the follows, and show an example result in Figure 4a.
[Equation]
For gradient L 2 norm, we deﬁne the metric as the follows, and we show an example result in
Figure 4b.
[Equation]
Note that the model or optimization algorithm does not have special requirement on gradient, such as gradient clipping.
These results indicate that other loss functions can behave similarly to the conditional cross-entropy when identifying causality direction.
The initial samples already show a large gap between the correct and incorrect causal models, and after just a few tens of examples, the difference is close to convergence.
We expect that the same intuition applies for these other loss metrics.
In particular, these loss metrics both satisfy distance metric properties of non-negativity and the triangle inequality.
Further, they have useful continuity and smoothness characteristics that suggests any biases they may have should behave like conditional cross-entropy.
Our empirical results conﬁrm these intuitions.
[Figure]
G R OBUSTNESS
To compare the robustness of the proposed approach with the baseline, we design P ( A ) with a complicated hidden structure, so that it is harder to adapt after intervention than P ( B ) and
P ( A | B ) .
For P ( A ) , we use a mixture model with K dimensional hidden variables C , and keep other distributions unchanged.
[Equation]
where θ ∈ R K , φ ∈ R K × N .
We set N = 10 , M = 10 , K = 200 .
For the baseline approach, we plot the improvement of log likelihood during adaptation along with the number of samples ( × 100 ) in Figure 5a.
The model with correct causality direction (blue) adapts slower than that of incorrect direction (red), which means the fundamental intuition of baseline approach does not apply in this situation.
For the proposed approach, we plot the difference of generalization losses
S G = G B
→ A − G A → B along with the number of samples ( × 100 ) in Figure 5b.
The difference (blue) is signiﬁcantly larger than zero (red).
This is natural because the proposed approach does not use P ( A ) , so that it is not inﬂuenced by whether P ( A ) has complicated hidden structure.
This experiment shows that the proposed approach is more robust than the baseline one in this case.
[Figure]
[Figure]
approach (a), which means the baseline approach does not work, but the proposed approach (b) works.
H R EAL D ATA E XPERIMENTS
We use altitude and temperature dataset (Mooij et al., 2016), where altitude ( A ) causes temperature ( B ).
We generate train and transfer datasets in the following way.
We ﬁrst divide the data to two sets D 1 and D 2 , according to their A values.
samples in D 1 has larger or equal A values than those in D 2 .
We respectively shufﬂe and separate D 1 and D 2 to two subsets with ratio of 9:1, generating
D 11 , D 21 , D 12 , D 22 .
We than merge sets to generate train data ( D 11 ∪ D 22 ) and transfer data ( D 21 ∪ D 12 ).
We use both linear regression model with L 2 loss and neural network model with two layers.
The hidden layer has 100 nodes with ReLU activation.
We use scikit-learn (Pedregosa et al., 2011) for implementation.
We run experiments 1,000 times with different random seeds, and the predicted causal directions are always correct.
For different noise levels, we use Gaussian noise with different standard deviations.
In the altitudetemperature experiment with the linear model, the success rate start to drop to 99 .
8 ± 0 .
1% at 0.5, and becomes 96 . 0 ± 0 .
4% at 1.0, 63 .
9 ± 1 . 8% at 5.0, 55 .
3 ± 1 .
2% at 10.0, and 50 .
1 ± 1 .
3% at 50.
(Figure 6).
This means the method is robust to some noise (around 1 or less), but does not work when the noise is too large.