CausalMACE: Mixture-of-Experts for CACE

• The paper introduces a mixture-of-experts architecture that estimates complier average causal effects (CACE) by modeling latent compliance types without needing monotonicity or exclusion restrictions.
• It employs a two-stage EM algorithm for estimating the compliance gating network and expert outcome regressions, ensuring identifiability and √n-consistency under mild conditions.
• Simulations and the PROBIT trial demonstrate that CausalMACE yields lower bias and variance compared to classical IV methods, especially when non-compliance challenges standard assumptions.

CausalMACE is the causal inference method introduced in “Estimating Complier Average Causal Effects with Mixtures of Experts,” where the name combines Causal with MACE, short for Mixture-of-Experts for the Complier Average Causal Effect (CACE). It is a mixture-of-experts architecture for randomized experiments with treatment non-compliance, in which principal strata are treated as latent classes, a gating network estimates stratum membership probabilities from covariates, and expert networks model stratum-specific potential outcomes. Its central purpose is to estimate the complier effect without relying on the classical instrumental-variable assumptions of monotonicity and exclusion restriction, replacing them with principal ignorability and structural mixture modeling (Grolleau et al., 2024).

1. Causal formulation and target estimand

CausalMACE is formulated for experiments with baseline covariates $X$, randomized assignment $Z \in \{0,1\}$, actual treatment taken $T \in \{0,1\}$, and observed outcome $Y$. Randomization is modeled as $Z \mid X \sim \text{Bernoulli}(\eta(X))$, with allocation ratio $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$. The methodological difficulty is that assignment and treatment can differ, so the intention-to-treat contrast in $Z$ is generally not the scientific quantity of interest when non-compliance is substantial (Grolleau et al., 2024).

The method adopts the principal-stratification view. Four latent compliance types are defined: compliers, always-takers, never-takers, and defiers. These are represented by a one-hot latent vector

$\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$

with $S_k \in \{0,1\}$ and $\sum_k S_k = 1$. Their covariate-dependent probabilities,

$Z \in \{0,1\}$0

are the core outputs of the gating network. Treatment actually taken is written as

$Z \in \{0,1\}$1

so compliers satisfy $Z \in \{0,1\}$2, always-takers satisfy $Z \in \{0,1\}$3, never-takers satisfy $Z \in \{0,1\}$4, and defiers satisfy $Z \in \{0,1\}$5 (Grolleau et al., 2024).

The target estimand is the complier average causal effect,

$Z \in \{0,1\}$6

To accommodate violations of exclusion restriction, the framework uses elementary potential outcomes $Z \in \{0,1\}$7, which allow outcomes to depend separately on assignment and treatment. Consistency is stated as

$Z \in \{0,1\}$8

Under principal ignorability and positivity of compliers, the paper derives the key decomposition

$Z \in \{0,1\}$9

where $T \in \{0,1\}$0 and $T \in \{0,1\}$1. Estimating CACE is therefore reduced to estimating the complier probability $T \in \{0,1\}$2 and the two complier outcome regressions $T \in \{0,1\}$3 and $T \in \{0,1\}$4 (Grolleau et al., 2024).

2. Identification strategy relative to classical IV

The standard LATE or IV formulation for CACE uses four assumptions: instrument relevance, instrument exogeneity, monotonicity, and exclusion restriction. Under these assumptions, CACE is identified by the Wald estimator. The paper’s starting point is that, in many applications, monotonicity can fail because defiers exist, and exclusion restriction can fail because assignment can have direct or placebo effects. In those settings, classical IV can be biased or only partially identifying (Grolleau et al., 2024).

CausalMACE departs from this framework in two explicit ways. First, it does not require monotonicity: defiers are retained as a distinct latent stratum and enter the treatment model through $T \in \{0,1\}$5. Second, it does not require exclusion restriction: elementary potential outcomes retain separate $T \in \{0,1\}$6 and $T \in \{0,1\}$7 arguments, so assignment may affect outcomes beyond treatment receipt. The method instead relies on principal ignorability, consisting of $T \in \{0,1\}$8 and $T \in \{0,1\}$9, together with positivity of compliers, $Y$0 (Grolleau et al., 2024).

This reallocation of assumptions is central to the method’s logic. Rather than imposing behavioral restrictions such as “no defiers” or “no direct assignment effects,” CausalMACE imposes a covariate-driven latent-mixture structure. The paper explicitly characterizes this as a trade of classical IV assumptions for modeling assumptions in a machine-learning spirit. A plausible implication is that the method is most attractive when rich covariates are available and monotonicity or exclusion restriction is substantively doubtful (Grolleau et al., 2024).

3. Mixture-of-experts parameterization

The first mixture-of-experts component models treatment receipt. Define

$Y$1

Lemma 4 shows that

$Y$2

where the experts are fixed and known: $Y$3 Thus,

$Y$4

and the unknown part of the first mixture is entirely the gating function $Y$5. In the parametric version,

$Y$6

so the gating network is a softmax over the four strata (Grolleau et al., 2024).

The second mixture-of-experts component models observed outcomes inside the two informative cells $Y$7 and $Y$8. Let

$Y$9

These decompose as

$Z \mid X \sim \text{Bernoulli}(\eta(X))$0

$Z \mid X \sim \text{Bernoulli}(\eta(X))$1

where the gating weights inside each cell are known functions of the first-step estimates: $Z \mid X \sim \text{Bernoulli}(\eta(X))$2

$Z \mid X \sim \text{Bernoulli}(\eta(X))$3

Accordingly, the outcome model is a known-gating two-expert mixture in each observed cell (Grolleau et al., 2024).

The expert functions $Z \mid X \sim \text{Bernoulli}(\eta(X))$4 are parametric or nonparametric outcome regressions. For binary outcomes, the paper uses logistic GLMs,

$Z \mid X \sim \text{Bernoulli}(\eta(X))$5

and for continuous outcomes it uses linear-Gaussian models,

$Z \mid X \sim \text{Bernoulli}(\eta(X))$6

with residual variance. Theoretical identifiability results are established for these mixtures when the expert functions are linear or expit functions. This architecture is the sense in which the method is a genuine mixture-of-experts estimator rather than a relabeled IV estimator (Grolleau et al., 2024).

4. Estimation, identifiability, and asymptotics

CausalMACE is estimated in three components. Step 1 estimates the compliance-type gating network $Z \mid X \sim \text{Bernoulli}(\eta(X))$7. Step 2 estimates the expert outcome functions $Z \mid X \sim \text{Bernoulli}(\eta(X))$8 and $Z \mid X \sim \text{Bernoulli}(\eta(X))$9 together with their companion experts $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$0 and $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$1. The final step substitutes these nuisance estimates into the principal-ignorability plug-in estimator

$\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$2

The paper describes this estimator as a partial M-estimator under correct parametric specification and regularity conditions (Grolleau et al., 2024).

Both major fitting stages are carried out by dedicated EM algorithms. In Step 1, stratum labels are treated as latent; the E-step computes posterior responsibilities $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$3, and the M-step updates $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$4 through multinomial logistic regression with soft labels. In Step 2, the latent labels are “complier versus always-taker” in cell $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$5 and “complier versus never-taker” in cell $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$6; the E-step computes posterior responsibilities, and the M-step fits weighted GLMs or weighted regressions for the experts. Nonparametric variants replace multinomial logistic regression and GLMs with more flexible learners while preserving the EM-like weighting scheme (Grolleau et al., 2024).

The theoretical contribution is unusually explicit for a flexible latent-mixture causal method. Theorem 4.1 gives identifiability of the treatment mixture $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$7 under mild conditions. Theorems 4.2 and 4.3 establish identifiability of the outcome mixtures for linear and logistic expert functions, respectively. In the appendix, the paper shows that $\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$8 and the nuisance parameters jointly solve a stacked system of estimating equations, leading to

$\eta(X) \in (\epsilon_\eta,1-\epsilon_\eta)$9

Thus the estimator is $Z$0-consistent and asymptotically normal, with variance estimated in practice by the nonparametric bootstrap (Grolleau et al., 2024).

5. Specialized regimes, simulations, and the PROBIT trial

Although the general formulation allows both defiers and direct assignment effects, the paper also defines specialized estimators for restricted regimes. Setting $Z$1 yields a monotonicity-only version, $Z$2. Adding exclusion restriction yields further simplifications and the estimator $Z$3. Separate algorithms are also given for ER-only regimes. These restricted versions are not the primary proposal, but they clarify how the general model nests more classical cases and can gain efficiency when stronger assumptions are credible (Grolleau et al., 2024).

The simulations consider four scenarios: neither monotonicity nor ER holds; ER holds but monotonicity fails; monotonicity holds but ER fails; and both monotonicity and ER hold. The data-generating process uses 14 correlated covariates, specifically 7 Bernoulli and 7 log-normal variables, with sample sizes $Z$4. The comparisons include $Z$5, $Z$6, $Z$7, $Z$8, and two classical IV benchmarks, the Wald estimator and IV matching, under both correct specification and misspecification induced by dropping important covariates (Grolleau et al., 2024).

The main empirical pattern is consistent across the scenarios. When IV assumptions hold, IV estimators perform very well, and $Z$9 is close in RMSE. When IV assumptions fail, Wald and matching show very large biases, on the order of roughly 40–60 percentage points, whereas CausalMACE estimators have much smaller bias and substantially lower RMSE, including under misspecification. In Scenario 1, where neither monotonicity nor ER holds, $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$0 has the lowest RMSE among the CausalMACE estimators; in Scenario 3, where monotonicity holds but ER fails, $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$1 is best. The variance is reported to scale roughly like $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$2, consistent with the $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$3-rate, and $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$4 often has the smallest variance (Grolleau et al., 2024).

The real-data illustration uses the Promotion of Breastfeeding Intervention Trial (PROBIT), a large randomized trial in Belarus. Assignment $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$5 is breastfeeding promotion versus standard care, treatment $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$6 is actual attendance at the breastfeeding program, outcome $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$7 is infant weight at 3 months dichotomized at 6000g, and $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$8 includes baseline covariates such as location, education, maternal allergy, smoking, C-section, sex of child, mother’s age, and birth weight. The intention-to-treat ATE of assignment on $\boldsymbol{S} = (S_c,S_a,S_n,S_d)^T,$9 is about 0.07, but 36% of those assigned to the program did not attend, making compliance-sensitive estimands relevant. The IV-based CACE, under monotonicity and ER, is about 0.11 with tight confidence intervals; $S_k \in \{0,1\}$0, which assumes neither monotonicity nor ER, is about 0.20 with wider confidence intervals. The paper interprets this gap as plausible because assignment to a breastfeeding program may have direct behavioral or psychological effects and because monotonicity may be violated if some mothers respond adversely to institutional encouragement (Grolleau et al., 2024).

6. Practical considerations and nomenclature

The implementation guidance in the paper is modular. It recommends starting with a multinomial logistic model for $S_k \in \{0,1\}$1, then fitting GLM-based or more flexible expert regressions for $S_k \in \{0,1\}$2, estimating the model in two stages via EM, and finally computing $S_k \in \{0,1\}$3 with bootstrap confidence intervals. It also emphasizes that covariates should include variables suspected to explain both compliance behavior and outcomes, that the number of strata is fixed by design at four, and that convergence monitoring for the EM algorithm should focus on log-likelihood and parameter stability. Joint estimation of all components is explicitly not recommended because of identifiability issues (Grolleau et al., 2024).

The principal limitation is the strength of principal ignorability: conditional on $S_k \in \{0,1\}$4, all variables jointly affecting compliance type and potential outcomes must be measured. The paper also notes that misspecification of $S_k \in \{0,1\}$5 or $S_k \in \{0,1\}$6 can bias CACE, that scalability to very high-dimensional $S_k \in \{0,1\}$7 may require regularization or modern machine-learning methods, that nonparametric identifiability is not yet fully characterized, and that extensions to multi-arm trials, longitudinal compliance, or censoring require further development. This suggests that CausalMACE is best viewed as a theoretically grounded middle ground between classical IV estimators and fully nonparametric principal-stratification models (Grolleau et al., 2024).

The term “CausalMACE” is not unique in the broader arXiv literature. It has also been used informally for the Model-Averaged Causal Estimation Transformer Neural Process (MACE-TNP), a meta-learned estimator of Bayesian model-averaged interventional distributions under uncertain causal graphs (Dhir et al., 7 Jul 2025), and it appears as the title of “CausalMACE: Causality Empowered Multi-Agents in Minecraft Cooperative Tasks”, a causality-aware multi-agent planning framework for Minecraft (Chai et al., 26 Aug 2025). In the sense defined here, however, CausalMACE denotes the mixture-of-experts method for estimating CACE under treatment non-compliance introduced in (Grolleau et al., 2024).