Compositional Models for Estimating Causal Effects

TL;DR

This work addresses the challenge of estimating individual-level causal effects in heterogeneous, modular systems by introducing a compositional framework that represents each unit as an instance-specific composition of multiple components. Using modular neural network architectures, the approach learns component-wise potential outcomes and aggregates them through a structured interaction graph to obtain unit-level effects, enabling compositional generalization to unseen configurations. The paper formalizes the compositional data-generating process, defines unit- and component-level causal estimands, and establishes identifiability under ignorability, overlap, and consistency, including a tractable additive-parallel special case. Through an experimental infrastructure with real-world benchmarks (query execution, manufacturing, matrix operations) and synthetic data, the authors demonstrate improved CATE estimation, sample efficiency, and robustness to observational bias, while also analyzing how component-level data access and composition structure influence performance. Overall, compositional causal modeling offers scalable, instance-specific reasoning for structured systems and highlights both its potential and conditions under which it may not yield gains, guiding future work on interventions at the component level and broader applications.

Abstract

Many real-world systems can be usefully represented as sets of interacting components. Examples include computational systems, such as query processors and compilers, natural systems, such as cells and ecosystems, and social systems, such as families and organizations. However, current approaches to estimating potential outcomes and causal effects typically treat such systems as single units, represent them with a fixed set of variables, and assume a homogeneous data-generating process. In this work, we study a compositional approach for estimating individual-level potential outcomes and causal effects in structured systems, where each unit is represented by an instance-specific composition of multiple heterogeneous components. The compositional approach decomposes unit-level causal queries into more fine-grained queries, explicitly modeling how unit-level interventions affect component-level outcomes to generate a unit's outcome. We demonstrate this approach using modular neural network architectures and show that it provides benefits for causal effect estimation from observational data, such as accurate causal effect estimation for structured units, increased sample efficiency, improved overlap between treatment and control groups, and compositional generalization to units with unseen combinations of components. Remarkably, our results show that compositional modeling can improve the accuracy of causal estimation even when component-level outcomes are unobserved. We also create and use a set of real-world evaluation environments for the empirical evaluation of compositional approaches for causal effect estimation and demonstrate the role of composition structure, varying amounts of component-level data access, and component heterogeneity in the performance of compositional models as compared to the non-compositional approaches.

Figures (17)

• Figure 1: Overview of key ideas: (a) Structured units: Units are composed of multiple heterogeneous components. Each color represents a distinct component. Treatment $T$ is applied to the unit, and the compositional system processes the inputs under intervention, returning potential outcomes. (b) Unitary approach: Standard approaches to effect estimation flattens the underlying structure. They use a fixed-size representation for each unit, aggregating component-level information to estimate unit-level potential outcomes. (c) Compositional approach: The compositional approach models each unit with an instance-specific structure. Component-level covariates $X_j$ and outcomes $Y_j$ are used to train each component model, and component-level outcomes are hierarchically aggregated to estimate unit-level potential outcomes. Each color represents a distinct component model with different parameters.
• Figure 2: Results on manufacturing domain ($10,000$ samples). We report $R^2$ between CATE estimates and ground-truth effects (higher is better). (a) Sample-size efficiency (WID): Compositional models are more accurate and sample-efficient in CATE estimation for within-distribution settings. (b) Compositional generalization (OOD): Models are trained on units with tree depth $\leq K$ and evaluated on a test-set with depth $8$. Compositional models generalize to the unseen combinations due to compositional structure whereas non-compositional baselines perform comparably only after training on units similar to test data. (c) Effect of increasing observational bias (OOD): Models are trained and tested on data with increasing observational bias strength between assigned treatment and tree depth. Compositional models and X-learner are less affected by increased observational bias.
• Figure 3: Role of component-level data access and composition structure in the performance of compositional models:$R^2$ score for models evaluated on compositional generalization task with varying degrees of component-level data access. (Higher is better; PEHE errors are reported in Figure \ref{['fig:prior_knowledge_pehe_app']} of the supplementary material). We observe that end-to-end trained models incorporating just modular structure compositionally generalize as trained on more module combinations. Unitary models show compositional generalization for additive parallel composition but perform comparably only for in-distribution combinations ($K$=$10$) for sequential composition, except X-learner. Note that the number of training samples increases as training depth increases.
• Figure 4: Results for real-world data sets: (a) Query execution data set: Compositional model estimates the effect more accurately as observational bias increases. (b) Matrix operations data set: All baselines perform similarly for this data set due to a single shared covariate, homogenous component outcome functions, and dominant contribution of matrix multiplication.
• Figure 5: Graphical representation of compositional causal models: Each plate model represents the data-generating process of unit-level and component-level variables for a given instance-specific composition structure ($G_i$), shown for three different structures here. $X_j$ denotes the component-specific covariates, $Y_j$ denotes the component-specific outcomes, and $T$ denotes the unit-level shared treatment. Each distinct color represents the fixed data generating process for a specific component, that might appear in multiple units.

Theorems & Definitions (6)

• Definition 1
• Definition 2
• Definition 3
• Lemma 4
• Definition 5
• Definition 6