Measure-Theoretic Anti-Causal Representation Learning

Causal representation learning in the anti-causal setting—labels cause features rather than the reverse—presents unique challenges requiring specialized approaches. We propose Anti-Causal Invariant Abstractions (ACIA), a novel measure-theoretic framework for anti-causal representation learning. ACIA employs a two-level design: low-level representations capture how labels generate observations, while high-level representations learn stable causal patterns across environment-specific variations. ACIA addresses key limitations of existing approaches by: (1) accommodating perfect and imperfect interventions through interventional kernels, (2) eliminating dependency on explicit causal structures, (3) handling high-dimensional data effectively, and (4) providing theoretical guarantees for out-of-distribution generalization. Experiments on synthetic and real-world medical datasets demonstrate that ACIA consistently outperforms state-of-the-art methods in both accuracy and invariance metrics. Furthermore, our theoretical results establish tight bounds on performance gaps between training and unseen environments, confirming the efficacy of our approach for robust anti-causal learning.

Our approach introduces a novel measure-theoretic framework for learning in anti-causal domains ($Y \rightarrow X \leftarrow E$). The method consists of three key algorithms: Causal Dynamics Algorithm, Causal Abstraction Algorithm, and OOD Optimization Algorithm. The framework decomposes the problem into two key representation spaces:

• Low-Level Representation ($\mathcal{Z}_L$): Captures the causal latent dynamics of how the label $Y$ generates the features $X$ (as per Theorem 3 in the paper).
• High-Level Representation ($\mathcal{Z}_H$): Learns the causal invariant abstraction (as per Theorem 4), ensuring the final prediction is independent of environment factors.

This two-level approach enables OOD generalization by learning from both perfect and imperfect interventions.

# Clone the repository
                git clone https://github.com/ArmanBehnam/ACIA.git
                cd ACIA

                # Install the package
                pip install -e .

                # Or install with development dependencies
                pip install -e ".[dev]"

from acia import ColoredMNIST, CausalRepresentationNetwork, CausalOptimizer
                from torch.utils.data import DataLoader, ConcatDataset

                # Load datasets
                train_e1 = ColoredMNIST(env='e1', train=True)
                train_e2 = ColoredMNIST(env='e2', train=True)
                train_loader = DataLoader(ConcatDataset([train_e1, train_e2]), batch_size=128)

                # Train model
                model = CausalRepresentationNetwork()
                optimizer = CausalOptimizer(model, batch_size=128)

                for x, y, e in train_loader:
                    metrics = optimizer.train_step(x, y, e)

python examples/example_colored_mnist.py

We thank the anonymous reviewers for their valuable and constructive feedback. This work was supported in part by the Cisco Research Award and by the National Science Foundation under Grant Nos. ECCS-2216926, CCF-2331302, CNS-2241713, and CNS-2339686.