Zhenyu (Zach) Wang
Home Blog Teaching
← All posts

Making invariance-based causal discovery actually scalable

Dec 2024, revised Jan 2026 · with Yifan Hu (equal contribution), Peter Bühlmann, Zijian Guo · arXiv:2412.11850

The short version

Setting
Data from several environments. We want the direct causes of \(Y\) and their effect sizes \(\beta^\star\) — not merely a good predictor — and we allow hidden confounding.
Principle
The causal model is the one whose prediction risk is equal across environments. Invariance identifies causality; prediction accuracy alone cannot.
Obstacle
Turning that principle into a method has meant testing all \(2^p\) subsets of covariates (ICP), or solving a mixed-integer program (EILLS). Both die well before \(p\) reaches practical sizes.
Idea
Recast the search as one continuous minimax problem, NegDRO, in which the environment weights are allowed to go negative. Negative weights enforce risk parity rather than mere robustness.
The catch, resolved
Negative weights destroy convexity. But we prove a benign landscape: every stationary point is close to \(\beta^\star\), so plain gradient descent–ascent suffices. No global optimization needed.
Payoff
Identification with as few as two environments (versus \(|\mathcal E| \ge p\) for IRM), polynomial runtime, and estimation error flat in \(p\) out to \(p = 100\) where ICP and EILLS have long since timed out.

The problem: prediction cannot see causation

A company wants to know which ad channel causally drives sales \(Y\): clicks from search ads \(X_1\), or engagement from social ads \(X_2\)? A model fit purely to predict \(Y\) will happily use both, because both correlate with sales. But correlation cannot separate the mechanisms below, and spending against a non-causal metric wastes money.

Y X₁ X₂ (a) X₁ → Y → X₂ Y X₁ X₂ (b) X₂ → Y, X₂ → X₁ Y X₁ X₂ (c) X₁ → Y ← X₂
direct cause of Y
Figure 1. Three causal structures that can produce the same correlations between \(X_1\), \(X_2\) and \(Y\). Least squares fits all three equally well. Only in (c) are both variables causes; in (a) \(X_2\) is a child of \(Y\), and a predictor will use it eagerly — it is highly informative and completely useless for intervention.

Invariance, and which flavour of it

A powerful principle rescues us: if a set of variables truly causes \(Y\), the way \(Y\) depends on them stays invariant across environments — time slots, regions, regimes — even as everything else shifts. Environments are modelled as additive interventions on the covariates: a linear structural equation model whose noise gains an environment-specific shift \(\delta^e\) on the \(X\) coordinates, never on \(Y\)'s own noise.

This paper uses the weakest useful form, risk invariance: only the residual second moment must match across environments. Writing \(R_e(b) = \mathbb E[(Y^e - b^\top X^e)^2]\), the invariant set is

\[ \mathcal B_{\mathrm{inv}} \;=\; \big\{\, b \in \mathbb R^p \ :\ R_e(b) = R_f(b) \quad \text{for all } e, f \in \mathcal E \,\big\}. \]

This matters more than it looks. The conditional-mean version of invariance used by EILLS, and the stronger independence-based version underlying ICP — which asks that \(\varepsilon_Y\) be independent of \(X_{S^\star}\) with an identical distribution across environments — both break under hidden confounding. Risk invariance does not: the systematic noise \(\eta\) may correlate \(Y\) and \(X\) freely, and \(\beta^\star\) remains the unique risk-invariant model. Theorem 1 shows \(\mathcal B_{\mathrm{inv}} = \{\beta^\star\}\) under an identification condition, and Theorem 2 shows that condition is nearly necessary.

The identification condition (Condition 1b) is refreshingly concrete — it asks only that the environments differ enough in a spectral sense. With \(\boldsymbol A(w) = \sum_e \big(w_e - \tfrac{1}{|\mathcal E|}\big)\,\mathbb E[X^e X^{e\top}]\), it requires some simplex weight \(w^0\) with

\[ \lambda \;:=\; \lambda_{\min}\big(\boldsymbol A(w^0)\big) \;>\; 0 . \]

The scalar \(\lambda\) is a measure of environmental heterogeneity, and it governs every rate below. Crucially, this can hold with as few as two environments — where IRM's guarantee wants \(|\mathcal E| \ge p\), and earlier invariance results wanted infinitely many.

Why the obvious algorithm is exponential

Invariant causal prediction turns the principle into a procedure: for every subset \(S \subseteq \{1,\dots,p\}\), test whether the relationship is invariant, and intersect the sets that pass. It is correct, and it is a search over \(2^p\) subsets. EILLS replaces the tests with a regularized least-squares objective, but that objective is discontinuous and nonconvex, and solving it still amounts to enumeration — essentially a mixed-integer program. Its continuous Gumbel-softmax relaxation buys speed by giving up the identification guarantee. So the field had a choice between a method with theory and no runtime, and a method with runtime and no theory.

NegDRO: one continuous problem, with negative weights

Start from the constrained program "minimize risk subject to risk invariance", and penalize the (nonconvex) equality constraint. What comes out is a minimax problem that looks like distributionally robust optimization, except that the weights on environments may be negative:

\[ b^\gamma_{\mathrm{Neg}} \;\in\; \arg\min_{b\in\mathbb R^p}\ \max_{w \in \mathcal U(\gamma)}\ \sum_{e\in\mathcal E} w_e\, R_e(b), \qquad \mathcal U(\gamma) = \Big\{ w : \textstyle\sum_e w_e = 1,\ \min_e w_e \ge -\gamma \Big\}. \]

Set \(\gamma = 0\) and \(\mathcal U(0)\) is the simplex: this is exactly Group DRO, which merely protects the worst environment. Let \(\gamma > 0\) and something different happens. The inner maximum is attained at a vertex that loads \(1 + \gamma(|\mathcal E|-1)\) on the worst environment and \(-\gamma\) on every other, so the objective has the closed form

\[ \Phi_\gamma(b) \;=\; \big(1 + \gamma|\mathcal E|\big)\,\max_{e}\, R_e(b) \;-\; \gamma \sum_{e\in\mathcal E} R_e(b). \]

Read it as: minimize the worst risk, while being rewarded for total risk. The only way to win on both counts is to pull the maximum risk down toward the average — that is, to equalize risks. As \(\gamma \to \infty\) the penalty enforces exact risk invariance and \(b^\gamma_{\mathrm{Neg}} \to \beta^\star\); at finite \(\gamma\), Proposition 2 gives \(\|b^\gamma_{\mathrm{Neg}} - \beta^\star\|_2 \lesssim \big[\lambda(1+\gamma|\mathcal E|)\big]^{-1}\). Robustness has become identification.

Nothing is free: negative weights make \(\Phi_\gamma\) a difference of convex functions, hence nonconvex. But unlike ICP and EILLS, it is a continuous problem in \(b\) — no subsets, no integers. In practice one adds a ridge term \(-\mu\|w\|_2^2\) inside the maximum, which makes the inner problem strongly concave and \(\Phi_\mu\) differentiable by Danskin's theorem, then alternates a closed-form simplex update for \(w\) with a gradient step on \(b\).

risk in each environment0.250.50.7511.25env 1env 2env 3env 4NegDRO equalizes risk (spread 0.002)ERM spread 0.647 · Group DRO 0.126ERMGroup DRONegDROestimated coefficients vs β⋆-0.5-0.250.250.5X₁X₂X₃X₄X₅dashes = truth. X₄, X₅ are children of Y:ERM and Group DRO use them, NegDRO does not.‖b − β⋆‖: ERM 0.38 · GDRO 0.17 · NegDRO 0.02
Figure 2. The paper's simulation model, solved at the population level. Here \(X_1 \to X_2 \to X_3\), \(Y = \tfrac12 X_1 - \tfrac12 X_3 + \varepsilon_Y\), and \(X_4, X_5\) are children of \(Y\) — so \(\beta^\star = (\tfrac12, 0, -\tfrac12, 0, 0)\). Left: ERM's risk swings across environments (spread 0.647); Group DRO flattens it somewhat (0.126); NegDRO equalizes it almost exactly (0.002). Right: that risk parity is what kills the spurious coefficients. ERM and Group DRO both lean on the children \(X_4, X_5\); NegDRO sets them to zero and recovers \(\beta^\star\) to within 0.02. All values computed from the model, not drawn by hand.

Why nonconvexity isn't fatal here

Solving a nonconvex problem to global optimality is generally hopeless, so a guarantee about the global minimizer would be of little use. The paper proves something stronger and stranger — a benign landscape. Under the same heterogeneity condition, for any point \(b\) whatsoever,

\[ \|b - \beta^\star\|_2 \;\lesssim\; \frac{1}{\lambda}\Big( \frac{1}{1+\gamma|\mathcal E|} \;+\; \|\nabla \Phi_\mu(b)\|_2 \Big) \;+\; \sqrt{\mu/\lambda}. \]

Read the right-hand side: it is small whenever the gradient is small. Every stationary point — not just the global minimum, not just local minima — sits near \(\beta^\star\). Saddle points and spurious local minima, if they exist, are all in the right neighbourhood. This is what licenses a plain first-order method, which is only ever guaranteed to find stationary points. Note the bound is on the distance to \(\beta^\star\), not to the optimum: it is not a classical error bound, and it is why global optimization is unnecessary rather than merely unattractive.

Feeding in sampling error and a finite iteration budget, Theorem 5 gives the non-asymptotic rate. With \(T\) iterations, \(n\) samples per environment, and the optimal \(\mu \asymp T^{-1/2}\),

\[ \|\widehat b^\gamma - \beta^\star\|_2 \;\lesssim\; \underbrace{\frac{1}{1+\gamma|\mathcal E|}}_{\text{penalty}} \;+\; \underbrace{T^{-1/4}}_{\text{optimization}} \;+\; \underbrace{n^{-1/4}}_{\text{statistics}} . \]

So \(\epsilon\)-accuracy needs \(\gamma = \Omega(\epsilon^{-1})\), \(T = \Omega(\epsilon^{-4})\) and \(n = \Omega(\epsilon^{-4})\) — all polynomial, with each iteration costing one gradient evaluation over \(p\) coordinates and \(|\mathcal E|\) risks. Whether the \(n^{-1/4}\) and \(T^{-1/4}\) can be improved to \(n^{-1/2}\), \(T^{-1/2}\) is stated as open.

Does it work? Accuracy that ignores \(p\), runtime that doesn't explode

The evidence is simulation — the paper has no real-data application. The headline experiment sweeps \(p\) from 5 to 100 with \(n = 20{,}000\), four environments, and a 30-minute cap per run. NegDRO and EILLS both use their recommended \(\gamma = 20\); the remaining causal-invariance baselines are tuned in an oracle manner, choosing whichever hyperparameter minimizes \(\|\widehat b - \beta^\star\|_2\) — a handicap in NegDRO's disfavour.

estimation error ‖b̂ − β⋆‖₂00.20.40.65255075100NegDROERM / AnchorICPtimes out at p = 15EILLStimes out at p = 25number of covariates pruntime (seconds, log scale)0.011100180030-minute cap5255075100NegDROAnchorERMICPEILLSnumber of covariates pExhaustive-search methods hit the wall; NegDRO’s cost grows polynomially.
Figure 3. Left: estimation error. NegDRO stays near 0.02–0.04 for every \(p\). ERM and Anchor regression plateau around 0.18 — they never remove the children of \(Y\). ICP is conservative (error ≈ 0.57) and stops at \(p = 10\). EILLS is the most accurate of all where it runs, then stops. Right: why they stop. ICP and EILLS grow exponentially and hit the 30-minute wall at \(p = 15\) and \(p = 25\); NegDRO's runtime is roughly flat at 10–20 seconds. Values read from the paper's Figure 8; the timing curves are approximate between measured points.
MethodConfounding?EnvironmentsSearchLargest \(p\) in 30 min
NegDRO (ours)yesas few as 2continuous100+
ICPno≥ 2all \(2^p\) subsets10
EILLSno≥ 2mixed-integer20
IRM\(|\mathcal E| \ge p\)continuous
ERM / Anchorconvex100+

"Confounding?" asks whether \(\beta^\star\) is still identified when the systematic noise correlates \(Y\) with \(X\); ICP's independence-based invariance and EILLS's conditional-mean invariance both fail there. In the paper's confounded simulation, NegDRO's error stays at 0.03–0.10 across \(p\), while ICP and EILLS both sit near 0.7. ERM and Anchor regression scale to any \(p\) but never identify \(\beta^\star\) at all — they are fast at solving the wrong problem.

A second experiment probes the regime the theory cares about: how strong must the interventions be? With two environments, four covariates, and interventions of varying strength, the competitors that need a dominated reference environment (DRIG) or an invertible difference of second moments (CausalDantzig) degrade badly, while NegDRO is nearly unaffected.

Estimation error, 2 environmentsLimited interventionsWeak interventionsStrong interventions
NegDRO (ours)≈ 0.05≈ 0.05≈ 0.05
CausalDantzig≈ 0.75≈ 0.75≈ 0.05
DRIGgrows to ≈ 2.5grows to ≈ 2.7≈ 0.05

Values read from the paper's Figure 5 at large \(\gamma\), \(n_e = 10{,}000\). "Limited" means some covariates are never intervened on; "weak" means the intervention variance is 0.01. DRIG's error increases with \(\gamma\) in these regimes — its guarantee needs a strictly dominated environment that does not exist here. Section 5 extends NegDRO's identification to the limited regime under a weaker condition: it suffices to intervene on the children of \(Y\).

2Environments that can suffice for identification, versus \(|\mathcal E| \ge p\) for IRM
10–20 sNegDRO's runtime, essentially flat from \(p = 5\) to \(p = 100\)
p = 15, 25Where ICP and EILLS hit the 30-minute wall

Why it matters

This closes a gap between an elegant statistical principle and usable methodology. Invariance-based discovery had the right theory but an exponential price tag; recasting it as a continuous nonconvex program — and then proving the landscape is benign — makes it feasible at realistic dimensions, with the multi-environment data modern applications already generate: multiple hospitals, time periods, regional markets, experimental regimes. The conceptual takeaway is that a negative weight is not a bug in the DRO formalism but the precise device that converts worst-case robustness into causal identification. To our knowledge this is the first causal-invariance method that provably reaches approximate global optimality for its nonconvex objective in polynomial time.

Honest caveats. The theory is for linear structural equation models, and the authors position it as a stepping stone toward nonlinear invariance learning. The \(n^{-1/4}\) and \(T^{-1/4}\) rates may well be suboptimal; closing them is open. Risk invariance hinges on the outcome's own noise never being intervened on, so all heterogeneity must come from additive interventions on the covariates. Identification conditions are sufficient and only nearly necessary — necessity is proven for single-coordinate interventions. And all evidence is simulated: there is no real-data application in the paper.

© 2026 Zhenyu (Zach) Wang 王振宇