Efficiently Disentangle Causal Representations

TL;DR

This work tackles the problem of learning disentangled causal representations by avoiding costly adaptation during transfer. It introduces a generalization-loss-based criterion, encapsulated in the score $\\mathcal{S_G} = \\mathcal{G}_{B \rightarrow A} - \\mathcal{G}_{A \rightarrow B}$, to identify the correct causal direction and learn causal representations under distribution shift. The authors provide theoretical justification linking conditional divergence and generalization gaps, and demonstrate substantial empirical gains in both causality direction prediction (up to $11.0\\x$ sample efficiency and $32.4\\x$ speed) and representation learning (up to $1.9\\x$ sample efficiency and $9.4\\x$ speed) across tasks, including real-data scenarios. The approach integrates smoothly with standard ML workflows, reduces hyperparameter burdens, and presents a robust, scalable alternative to adaptation-speed baselines.

Abstract

This paper proposes an efficient approach to learning disentangled representations with causal mechanisms based on the difference of conditional probabilities in original and new distributions. We approximate the difference with models' generalization abilities so that it fits in the standard machine learning framework and can be efficiently computed. In contrast to the state-of-the-art approach, which relies on the learner's adaptation speed to new distribution, the proposed approach only requires evaluating the model's generalization ability. We provide a theoretical explanation for the advantage of the proposed method, and our experiments show that the proposed technique is 1.9--11.0$\times$ more sample efficient and 9.4--32.4 times quicker than the previous method on various tasks. The source code is available at \url{https://github.com/yuanpeng16/EDCR}.

Figures (8)

• Figure 1: Experiments to evaluate efficiency in causality direction prediction. The model is discrete, and both $A$ and $B$ are ten-dimensional variables ($N=10$). The proposed approach is green. The baseline approach bengio2020a is orange.
• Figure 2: Experiments to evaluate efficiency in representation learning with the same setting. The curve of the proposed approach (green) is above that of the baseline approach (orange), indicating the proposed method has both better sample and time complexity.
• Figure 3: Experiments to evaluate efficiency in causality direction prediction. The model is discrete, and both $A$ and $B$ are 100-dimensional variables ($N=100$).
• Figure 4: Experiments for other metrics. The model is discrete and $N=10$. Curve (blue) is median over 100 runs, with 25-75% quantiles intervals, and it is significantly above zero (red), indicating these metrics are good indicators for causality learning.
• Figure 5: Experiments to evaluate robustness. $A \rightarrow B$ is the correct causality direction. Here, $N=10, M=10, K=200$. Curves are median over 100 runs, with 25-75% quantiles intervals. The result shows that the models with correct causalities are slow to adapt in the baseline approach (a), which means the baseline approach does not work, but the proposed approach (b) works.

Theorems & Definitions (4)

• Proposition 1
• proof
• Proposition 2
• proof