Automatic, Debiased, and Invariant Counterfactual Generation under General Interventions Raphael C. Kim1,2 , Jingsen Zhu1 , Ramin Zabih1 , and Michele Santacatterina2 1
arXiv:2606.07399v1 [stat.ML] 5 Jun 2026
2
Cornell Tech, Cornell University, New York, NY Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY
Abstract Generative models for counterfactual outcomes have great potential to support decision-making under complex interventions, but existing approaches are limited by unstable estimation, poor generalization across environments, and bias from nuisance model misspecification. We introduce ADIGen, a framework for automatic, debiased, and invariant counterfactual generation under general interventions, including highdimensional interventions and outcomes. ADIGen combines Riesz regression to avoid unstable density-ratio estimation, causal invariance to improve generalization under distribution shift, and orthogonal statistical learning to obtain doubly robust guarantees against nuisance model misspecification. We provide excess-risk bounds showing that ADIGen controls counterfactual risk under general interventions, with a product-bias nuisance remainder and an invariant risk bound across environments.
1
1
Introduction
Decision-making in complex systems often requires understanding counterfactuals of general, potentially highdimensional, interventions with limited data. Collecting sufficient data for every counterfactual in complex systems may be near impossible due to cost or ethical reasons. With the recent growth in expressivity and power in generative modeling, generative models that can synthesize counterfactual outcomes under generalized interventions stand as a viable solution for supporting robust decision-making in real-world systems. In an ideal world, we may simply train a generative model with the data we have, and sample from the generator under the intervention of interest. Counterfactual generative modeling may fail with such an approach due to confounding bias. Correlations observed in the sampled data may be mistaken for true causal effects, yielding incorrect downstream decisions. For example, generating medical images under changes in intervention dose can help track disease progression and identify optimal dosing strategies. However, if the training data primarily consisted of those who were responsive to intervention (e.g., younger populations), then the generator would identify the ranges in the data as effective even if this does not hold for different populations (e.g. older populations). This problem arises in other settings too. If we wish to generate images of faces under different lighting conditions, it may be biased towards faces seen in darker or brighter conditions. If we wish to understand sequential decision-making policies in robotic environments, similar confounding bias may arise. Counterfactual generative modeling provides a principled solution by explicitly accounting for confounding bias and modeling how outcomes would change under alternative interventions rather than what was simply observed. Despite the remarkable progress, three main challenges hinder the ability of counterfactual generative models for complex real-world, decision-making tasks. First, general interventions may include binary, multivalued, continuous, or high-dimensional interventions. The standard tool for approaching this requires estimating density ratio-type objects, but this can yield high instability for non-binary interventions. Second, counterfactual generators may fail to generalize across environments due to distribution shifts from training to test time. Third, misspecification bias can arise when methods adjust for confounding using imperfect models. Contributions. In this paper, we introduce ADIGen, a framework for Automatic, Debiased, and Invariant counterfactual Generation under general interventions, including high-dimensional interventions and outcomes. In light of the concerns raised above, we provide the following contributions: (Co1) Flexible, automated generative modeling of generalized interventions without requiring (i) training a separate model for each intervention, (ii) no explicit causal structural model, and (iii) circumventing unstable estimation of importance sampling-type ratios. This supports ‘on-the-fly’ generation of counterfactuals of interest for potentially high-dimensional intervention and outcome relationships. (Co2) Causally invariant counterfactual generative modeling, supporting high-quality and generalizable generative modeling that remains stable across distribution shifts. (Co3) Doubly-robust, invariant counterfactual modeling of generalized interventions. We provide theoretical results demonstrating that ADIGen controls excess invariant counterfactual risk under general interventions, with a doubly robust (DR) nuisance remainder and robustness guarantee across distributions. In contrast to non-DR approaches, we relax convergence requirements for the generative model of interest.
1.1
Related Works
Doubly robust generation of counterfactuals. Most closely related, Luedtke and Fukumizu (2025) propose a doubly robust algorithm for counterfactual generative modeling of Y (a∗ ) under binary interventions, A = {0, 1}. Their work shows that orthogonalization can be used to control nuisance-model bias in counterfactual generation. Our work extends this framework beyond binary interventions to general intervention spaces, including multivalued, continuous, and high-dimensional interventions. We further incorporate causal invariance to target mechanisms that remain stable across environments. Counterfactual generation under general interventions. Several lines of work are related to counterfactual generation under multivalued, continuous, or structured interventions. Agarwal et al. (2023) study combinatorial interventions using latent factors and a matrix-completion perspective. Nie et al. (2021) propose
neural-network methods for continuous interventions based on varying coefficient models. Kladny et al. (2023) study causal effect estimation from observational and interventional data using matrix weighted linear estimators. Israel et al. (2023) use variational backdoor adjustment for high-dimensional confounders, interventions, and outcomes. These approaches address increasingly general intervention spaces, but they are not designed as invariant, doubly robust generative models for high-dimensional counterfactual outcomes. Causal generative models for image counterfactuals. A second related line of work uses causal generative models for image counterfactuals. De Sousa Ribeiro et al. (2023) generate high-fidelity image counterfactuals using probabilistic causal models and hierarchical variational autoencoders. Yang et al. (2020) propose CausalVAE, which imposes causal structure on latent variables in a variational autoencoder and supports interventions on latent causal factors. Shen et al. (2020) develop weakly supervised disentangled causal representation learning, allowing interventions on learned causal latent variables. Sauer and Geiger (2021) introduce Counterfactual Generative Networks, combining structural causal models with generative adversarial networks to generate images under interventions on causal factors. Pawlowski et al. (2020) use neural conditional distributions to define deep structural causal models for tractable counterfactual inference in high-dimensional settings. More recently, Komanduri et al. (2024) combine structural causal modeling with diffusion autoencoders for counterfactual image generation. These methods provide powerful causal generative models, but they typically rely on explicit structural causal models, latent causal assumptions, and model correctness. Invariance, transfer, and generative modeling. Our work is also related to invariance and transfer learning. Janetzky et al. (2026) study continual learning of domain-invariant representations. Gardner et al. (2024) use language modeling for large-scale transfer learning with tabular data. Cai et al. (2024) study active exploration through autoregressive generation of missing data. These works are related in their use of representation learning, transfer, or generative modeling across distributions, but they do not directly target invariant and doubly robust counterfactual generation under general interventions. Table 1 summarizes how these methods relate to our three target contributions. Table 1: Comparison of related approaches for counterfactual generation. Co1: supports general interventions; Co2: targets invariant mechanisms across environments; Co3: provides doubly robust protection against nuisance misspecification. Method
Co1
Co2
Co3
DoubleGen (Luedtke and Fukumizu, 2025) Synthetic Combinations (Agarwal et al., 2023) VCNet (Nie et al., 2021) Matrix Weighted Linear Estimators (Kladny et al., 2023) Variational Backdoor Adjustment (Israel et al., 2023) High-Fidelity Image Counterfactuals (De Sousa Ribeiro et al., 2023) CausalVAE (Yang et al., 2020) DEAR (Shen et al., 2020) Counterfactual Generative Networks (Sauer and Geiger, 2021) Deep Structural Causal Models (Pawlowski et al., 2020) Causal Diffusion Autoencoders (Komanduri et al., 2024)
× ✓ ✓ ✓ ∼ ∼ ∼ ∼ ∼ ∼ ∼
× × × × × × × × × × ×
✓ × × ∼ × × × × × × ×
ADIGen
✓
✓
✓
Legend. ✓: addressed; ×: not addressed; ∼: partially addressed or related, but not the primary target.
2
ADIGen
In this section, we introduce our approach ADIGen. Let Z = (X, A, Y, E) denote our observed data tuple i.i.d.
with covariates X, intervention A, outcome Y , and environment E. We assume Zi ∼ P. Our aim is to sample from a causal invariant of P[Y (a∗ )] for a∗ ∈ A, where A ⊆ RdA is potentially high-dimensional, sampled from distribution πb . The learned mechanism will be invariant across environments e ∈ E (formalized in
(Inv1) below), promoting out-of-distribution performance. At a high level, our approach involves a first step of nuisance model training of the propensity score and conditional density model, and a second step of generative model training using a doubly-robust loss derived from these nuisances. Together, this yields a counterfactual generative model of interest. However, our approach has key features that address the fundamental challenges to counterfactual modeling raised above. (Co1) The primary challenge for moving from binary interventions to learning flexible, generalized interventions, is stable learning of density ratio-type objects. Our proposal is to learn a ‘universal’ Riesz representer (URR) rather than an inverse propensity score during nuisance training. This enables us to stabilize learning for complex interventions while generalizing across several interventions of interest without having to retrain a separate model for each intervention which is unscalable. More formally, the Riesz Representation theorem states that for any continuous linear functional E[Y (a∗ )], there exists some α called the Riesz representation of our functional that can be expressed as an inner product between the representer α and regression function ψ (typically E[Y | A, X]). E[Y (a∗ )] = E[α(X; a∗ ) · ψ(X, A)] We extend this to multiple interventions using a URR, a function of intervention, covariate, and environment, inducing the following risk minimization problem, over function class Fα . This function class can be a neural network, random forest, or any model class of interest: (1) α0 = arg min EA,X,Y,E α(A, X, E)2 − 2 ψ(X, A) α(A, X, E) α∈Fα
(Co2) Generalization across distributions is challenging. Cost or ethical reasons yielding limited intervention data exacerbates the concerns for decision-making in high-stakes settings. To overcome this, we propose to learn a causally invariant mechanism. Intuitively, this means that across environments e ∈ E, we seek to learn the underlying causal mechanism that is stable across environments. More formally, we assume that there exists some feature representations on covariate and intervention S0 (X), T0 (A) s.t. the following holds: Y (a∗ ) ⊥ E | S0 (X), T0 (a∗ ). This means that the causal mechanism is invariant to the environment, or (Inv1) Causal Invariance across Environments: For all environments E, E ′ ∈ E, PE (Y (a∗ ) | S0 (X), T0 (a∗ )) = PE ′ (Y (a∗ ) | S0 (X), T0 (a∗ )).
(2)
(Co3) Correct model specification in real-world settings is nearly impossible to guarantee. We propose a doublyrobust approach for invariant counterfactual generation under general interventions, reducing the reliance on nuisance estimation by having error rely on a product bias between the density model and the propensity score model (Chernozhukov et al., 2018). More concretely, the error G(θ) (defined more closely below) takes the following form. With high probability, Singly-Robust:
ADIGen:
2 G(θ̂n ) ≲ inf G(θ) + δn,inv + θ∈Θ
1 + d2Ψ,inv (ψ̂n , ΨP ) + λinv rn,I , n
1 2 G(θ̂n ) ≲ inf G(θ) + δn,inv + + ∥α̂n − α∥22 d2Ψ,inv (ψ̂n , ΨP ) + λinv rn,I θ∈Θ n | | {z } {z } | {z } Invariance Doubly-Robust Oracle
where drawing from empirical process theory (Pollard, 1990), δ is our localization radius, d is a suitable distance measure over our generative model class from the true density, and λinv rn,I are invariance specific penalties. What should be noted from this result is that ADIGen yields oracle excess risk with product bias on the URR convergence rate and the generator convergence rate, ∥α̂n − α∥22 d2Ψ,inv (ψ̂n , ΨP ), while the non-DR approach is reliant on the generator convergence rate alone d2Ψ,inv (ψ̂n , ΨP ). Consequently, the convergence rate requirements for the generator would be weaker in ADIGen in comparison to non-DR approaches under sufficient regularity of the URR. This is formalized in Section 3.
At a high level, ADIGen has two steps. In the nuisance training step, we estimate three objects on a separate fold: the universal Riesz representer α̂ of the target functional, the invariant outcome model ψ̂ (a conditional density, and the invariant maps (Ŝn , T̂n ). In the estimation phase, we minimize a cross-fitted doubly robust risk over the generative class Θ with an invariance penalty. The full procedure is in Algorithm 1. The doubly robust Step 2 has three loss terms. The first is the URR-weighted loss in the estimation phase, residual α̂(A, X, E) ℓ(θ; Y, A) − ζψ̂ (θ, X, A, E) . The second is the standard is the plug-in loss. The superscript j indicates the data partition used during cross-fitting: nuisance models are trained on fold Z j so that the models trained are ‘fixed’ and independent from the evaluation fold Z 3−j . As shown in Theorem 1, these terms yield a doubly-robust error that weakens the generator convergence rate requirements under sufficient regularity of URR. The final term is an invariance penalty λinv În (θ, η̂), added to enforce stability across environments. We remark that we let u ∼ Π denote input noise to the generative model ψ̂(u | ·). This generalizes across diffusion, flow-matching, and autoregressive generators. Algorithm 1 ADIGen. Require: iid
• Data Z1 , . . . , Zn ∼ P , where Zi = (Xi , Ai , Yi , Ei ) • Partition indices into two folds Zn1 , Zn2 of size ⌊n/2⌋ • Choice of generative modeling framework • Representations S and T for the invariance classes on X and A respectively • Invariance penalty I and tuning parameter λinv ≥ 0 1. Nuisance and representation estimation. For j ∈ {1, 2}, use observations in Znj to estimate: 1. Invariant outcome map: learn ψnj such that ψnj (u, Snj (X), Tnj (A)) = P (Y (A) | π(A), X) 2. Universal Riesz representer: estimate αnj (a, x, e) solving α̂ = arg min En α(A, X, E)2 − 2 α(A, X, E) α∈Fα
(3)
3. Covariate and intervention representations: Snj (X),
Tnj (A).
2. Invariant risk minimization. Define θn as the minimizer of Rn,inv (θ, S, T ) =
Z 2 h i 1X X αnj (a, x, e) ℓ(θ(x, a), y) − ℓ θ(x, a), ψnj (u | ϕjn (x), ρjn (a)) n j=1 3−j z∈Zn + ℓ θ(x, a), ψnj (u | Snj (x), Tnj (a)) Π(du) + λinv Ibn3−j (θ, ψnj , Snj , Tnj )
where z = (x, a, y, e). 3. Return ϕn := τ (θn ), Sn , Tn Unbiased stochastic gradients can be obtained by sampling (j, z, u) and environment pairs (e, e′ ).
3
Theoretical Results
In this section, we wish to provide an excess risk bound for ADIGen. Let π be the intervention distribution of interest, and πb be the observed intervention distribution, defined as the generalized propensity score over the
invariant class T . Define the risk of an estimator θ by Z R(θ) = E[ℓ(θ; Y (a), a)] =
E[ℓ(θ; Y (a), a)]dν(a)
for loss ℓ and measure ν. Note the expectation here is taken over environments too. The aim is to bound the excess risk given by G(θ) = R(θ) − inf R(θ′ ). ′ θ ∈Θ
′ It will be helpful R to define the intervention-specific excess risk and risk given by Ga (θ) = Ra (θ) − inf θ′ ∈Θ Ra (θ ), meaning R = Ra dν(a). In order to train our model, we utilize a doubly-robust loss that is a function of nuisance models. The nuisances will be denoted by η = (α, ψ, T, S), is a tuple consisting of the URR, outcome density, treatment invariant mapping, and covariate invariant mapping respectively. Now define the DR-loss by h i Lη̂ (θ)(Z) = α̂(a, X, E) ℓ(θ; Y, a) − ζψ̂ (θ, X, a) + ζψ̂ (θ, X, a),
where
Z Z ℓ(θ; ψ̂(u | x, a), a)Π(du)dν(a)
ζψ̂ (θ, x, a) =
√ Notation. We let E be the population expectation, En be the empirical average, and Gn = n(En − E) be the empirical process. η̂ will denote estimated quantities in a cross-fit manner, and η0 will denote the true functions. All constants introduced below are required to hold uniformly over E ∈ E. We will reason about entropy conditions on the loss class given by fθ = Lη0 (θ) − Lη0 (θ0 ),
fˆθ = Lη̂ (θ) − Lη̂ (θ0 )
The uniform entropy integral for function class F at scale δ will be denoted by Z δp J(δ, F ) = sup 1 + log(N (ε∥G∥2 , F, L2 (Q)))dε Q
0
for N (s, F, Q) the covering number at scale s on function class F and measures Q (Pollard, 1990). We make the following assumptions (A1) Existence of oracle minimizer. There exists θ0 ∈ arg min R(θ). θ∈Θ
(A2) Bounded loss. There exists a constant C2 < ∞ such that for all (a, E) ∈ A × E, sup ∥ℓ(θ; ·, a) − ℓ(θ0 ; ·, a)∥L∞ (P) ≤ C2
θ∈Θ
(A3) Curvature. There exists C3,a < ∞ such that ∥ℓ(θ; ·, a) − ℓ(θ0 ; ·, a)∥2L2 (P(Y (a)) ≤ C3,a Ga (θ) where P(Y (a)) =
R
P (Y | a, x, E)dPX|E for all θ ∈ Θ and a ∈ A.
(A4) Entropy/localization condition. The conditional loss class F = {z 7→ ℓ(θ; z) − ℓ(θ0 ; z) : θ ∈ Θinv } has J(δ, Fδ ) < ∞ for sufficiently small δ where Fδ = {f ∈ F :√∥f ∥2 ≤ δ∥F ∥2 } for envelope F of class. The critical radius δn,inv is the smallest δ solving J(δ, Fδ ) ≤ nδ 2 governing the fundamental difficulty accordin to the function class.
(A5) Conditional mixed-Lipschitz condition. There exists C5 < ∞ such that Z 2 [ℓ(θ; ψ(u | x, a, E), a) − ℓ(θ; ψ0 (u | x, a, E), a)] Π(du) ≤ C5 Ga (θ)dΨ (ψ, Ψ0,a,E )2 for all (x, a, E). (A6) Nuisance convergence rates. There exist positive rα , rψ > 0 such that uniformly over E ∈ E, ∥α̂ − α0 ∥L2 (πb ×PX,E ) = Op (n−rα ) dΨ,inv (ψ̂, Ψ0 ) = Op (n−rψ ) (A7) Invariance Regularity. We assume realizability, and our penalty is learnable: I.e. there exists a sequence rn,I → 0 such that with high probability, sup ∥În (·) − I(·)∥ ≤ rn,I θ,ψ,T,S
Assumptions (A1)-(A4) are standard for empirical risk minimzation results and analogous to C3-C6 of DoubleGen (Luedtke and Fukumizu, 2025). Together, these permit empirical process control (Term II in Section A.1). Assumptions (A5)-(A6) help control our DR-remainder term (Term I in Section A.1). Finally, assumption (A7) ensures we are learning a causal invariant problem with this ERM and we have proper approximation of our invariance penalty (Term I in Section A.1). If we wish to practically utilize this approach for general interventions, we require flexible modeling and suitable existence of our possible interventions of interest. Let π be our intervention distribution of interest, and πb be our observed distributions. We will overload the notation to be the generalized propensity score over the invariant class. In other words, π(A) = P[T0 (A) = T ]. Now, assume, (Int1) Overlap There exists ν ≥ 0 such that h π(A | X) 2+ν i Eπb
πb (A | X)
<∞
(Int1) states that we have sufficient probability to observe interventions of interest in the data and guarantees learnability over our distributions. We are now set to state our result. Theorem 1 (Invariant AutoDoubleGen excess risk). Assume ERM regularity assumptions (A1)-(A4), nuisance assumptions (A5)-(A6), causal invariance (Inv1), and overlap (Int1) across environments E ∈ E. Let Θinv ⊆ Θ denote the invariant generator class. Then the estimator θ̂ derived from Algorithm 1 satisfies, with probability 1 − exp(−s), s 2 G(θ̂) − inf G(θ) ≲ δn,inv + + ∥α̂ − α0 ∥2L2 (πb ×PX,E ) dΨ,inv (ψ̂, Ψinv )2 + λinv rn,I θ∈Θ n where δn,inv is the critical radius. The proof is found in Section A.1. Theorem 1 shows that the learned generator approaches the best invariant counterfactual generator in the chosen class, with error controlled by oracle risk, nuisance estimation error, and the cost of enforcing stability across environments. Thus, ADIGen controls counterfactual risk under distribution shift while preserving a doubly robust product-bias structure having the potential to relax convergence rates significantly. This implies a bound on the divergence measure of our generative model: Corollary 1.1 (Divergence upper bound on Generative Model Performance). Assume the conditions of Theorem 1. Let the divergence measure of interest be denoted by D. Then, there exists s, b, ϵ > 0 such that with probability 1 − exp(−s), the following holds ib h 1/2 D P, P̂ (θ̂) ≤ C δn,inv + ns + ∥α̂ − α0 ∥2L2 (πb ×PX,E ) dΨ,inv (ψ̂, Ψinv )2 + λinv rn,I + ϵ This proof follows by combining Theorem 1 with C1 of DoubleGen Luedtke and Fukumizu (2025), and can be specialized to generative models of interest in a similar manner.
4
Conclusion
We introduce ADIGen, a framework for counterfactual generation under general interventions that combines Riesz regression, causal invariance, and orthogonal statistical learning. Our results show that ADIGen controls excess counterfactual risk while preserving a doubly robust product-bias structure and accounting for distribution shifts across environments.
References Agarwal, A., Agarwal, A., and Vijaykumar, S. (2023). Synthetic combinations: A causal inference framework for combinatorial interventions. arXiv preprint arXiv:2303.14226. Ajith, A. and Gopakumar, G. (2023). Domain adaptation: A survey. In Tistarelli, M., Dubey, S. R., Singh, S. K., and Jiang, X., editors, Computer Vision and Machine Intelligence, pages 591–602, Singapore. Springer Nature Singapore. Cai, T. T., Namkoong, H., Russo, D., and Zhang, K. W. (2024). Active exploration via autoregressive generation of missing data. arXiv preprint arXiv:2405.19466. Chernozhukov, V., Chetverikov, D., Demirer, M., Duflo, E., Hansen, C., Newey, W., and Robins, J. (2018). Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1):C1–C68. De Sousa Ribeiro, F., Xia, T., Monteiro, M., Pawlowski, N., and Glocker, B. (2023). High fidelity image counterfactuals with probabilistic causal models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 7390–7425. PMLR. Gardner, J., Perdomo, J. C., and Schmidt, L. (2024). Large scale transfer learning for tabular data via language modeling. arXiv preprint arXiv:2406.12031. Israel, D., Grover, A., and Van den Broeck, G. (2023). High dimensional causal inference with variational backdoor adjustment. arXiv preprint arXiv:2310.06100. Janetzky, P., Schlagenhauf, T., and Feuerriegel, S. (2026). Continual learning of domain-invariant representations. arXiv preprint arXiv:2605.15775. Kladny, K.-R., von Kügelgen, J., Schölkopf, B., and Muehlebach, M. (2023). Causal effect estimation from observational and interventional data through matrix weighted linear estimators. arXiv preprint arXiv:2306.06002. Komanduri, A., Zhao, C., Chen, F., and Wu, X. (2024). Causal diffusion autoencoders: Toward counterfactual generation via diffusion probabilistic models. arXiv preprint arXiv:2404.17735. Luedtke, A. and Fukumizu, K. (2025). Doublegen: Debiased generative modeling of counterfactuals. arXiv preprint arXiv:2509.16842. Nie, L., Ye, M., Liu, Q., and Nicolae, D. (2021). VCNet and functional targeted regularization for learning causal effects of continuous treatments. arXiv preprint arXiv:2103.07861. Pawlowski, N., Castro, D. C., and Glocker, B. (2020). Deep structural causal models for tractable counterfactual inference. arXiv preprint arXiv:2006.06485. Pollard, D. (1990). Empirical Processes: Theory and Applications, volume 2. NSF-CBMS Regional Conference Series in Probability and Statistics. Sauer, A. and Geiger, A. (2021). Counterfactual generative networks. arXiv preprint arXiv:2101.06046. Shen, X., Liu, F., Dong, H., Lian, Q., Chen, Z., and Zhang, T. (2020). Weakly supervised disentangled generative causal representation learning. arXiv preprint arXiv:2010.02637. van der Vaart, A. and Wellner, J. (1996). Weak Convergence and Empirical Processes. With Applications to Statistics. New York: Springer.
van der Vaart, A. and Wellner, J. A. (2011). A local maximal inequality under uniform entropy. Electronic Journal of Statistics, 5(none):192 – 203. Yang, M., Liu, F., Chen, Z., Shen, X., Hao, J., and Wang, J. (2020). CausalVAE: Structured causal disentanglement in variational autoencoder. arXiv preprint arXiv:2004.08697.
SUPPLEMENTARY MATERIAL TO “Automatic, Debiased, and Invariant Counterfactual Generation under General Interventions” Raphael C. Kim1,2 , Jingsen Zhu1 , Ramin Zabih1 , and Michele Santacatterina2 1
2
Cornell Tech, Cornell University, New York, NY Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY
A
Proofs
Lemma 2. Let F be a function class with measurable envelope function F s.t. ∥F ∥2 < ∞ and a finite,√supremum entropy integral J(δ) (envelope normalized). Define δn to be the critical radius, or δ s.t. J(δ) ≤ δ 2 n. Then, for any s > 0, with probability 1 − exp(−s), the following holds: sup ∥(E − En )f ∥ ≲ δn ∥f ∥2 + δn2 +
f ∈F
s n
Proof. Define the localized supremum process at radius r by vn (r) = supf ∈F :∥f ∥2 ≤r ∥(E − En )f ∥. Then, by Theorem 2.1 of van der Vaart and Wellner (2011), we have J(r) J(r) vn (r) ≲ √ (1 + √ 2 )∥F ∥2 n nr ∥F ∥2 Next, we claim that vn (r) ≲ δn2 + δn r for r ≥ 0. Note that by Lemma S3 of Luedtke and Fukumizu (2025), J(r)/r is non-increasing in r. Consider when √ r ≥ δn . It suffices to show that (1 + √J(r) ) is bounded. Since δn is the critical radius, J(δn )/δn = nδ so nr 2 √ J(r)/r ≤ J(δn )/δn ) =⇒ J(r)/ n ≤ δn r . Hence, dividing by r2 the quantity is bounded, and the upper bound on δn2 follows. When r < δn , vn (r) ≤ vn (δn ) ≲ δn2 . This gives the claim. By Lemma 3.5.9 of van der Vaart and Wellner (1996) (Talagrand’s), with probability 1 − exp(−s), vn (r) ≲ Evn (r) + r
p
s/n +
s n
p Now, we can further bound using Young’s to show r s/n ≲ rδn + rs/(nδn ) and the result from the first step showing vn (r) ≲ δn2 + δn r. Combining with the union bound on levels of r, we conclude the result. Lemma 3 (Invariance Penalty Estimation). Suppose we have an invariant measure of the form I(θ, η) = Φ({Eh(θ, η)}E∈E ) for Φ an aggregation operator h a function of sample points. Assume H = {h(θ, η; ·) : θ ∈ Θinv } is uniformly bounded with finite uniform-entropy integral J(H) < ∞. Then, sup ∥În (θ̂) − În (θ0 )∥ = op (1) Proof. This setup encompasses various measure such as MMD, CORAL, moment differences, and others (Ajith and Gopakumar, 2023). For example, for MMD, we may choose h to be the kernel of choice and Φ to be a pairwise sum over environment discrepancies. To prove our result, we will bound the differences using standard empirical process theorems such as the Dudley integral (Pollard, 1990). More concretely, ∥În (θ, η̂) − I(θ, η)∥ ≲ ∥En h(θ, η̂) − Eh(θ, η)∥ ≲ ∥(En − E)h(θ, η̂)∥ + ∥E(h(θ, η̂) − h(θ, η))∥ The first term is an empirical process term, and the samples are cross-fit so η̂ is fixed. Under our controlled entropy assumption on h, this is op (1). For the second term, assume h is neyman-orthogonal with respect to η (Chernozhukov et al., 2018). Then, we have second order error permitting similar convergence requirements as the DR-approach above yielding an op (1) rate. Note, if the function h is not neyman-orthogonal, we may orthogonalize.
A.1
Proof of Theorem 1
Proof. The proof proceeds in four primary steps. First, we employ the basic inequality to decompose the risk into a doubly-robust remainder, empirical process term, and penalty term. Then, we control each term in three separate steps, with the doubly-robust remainder using orthogonal statistical learning, the empirical process term controlled using chaining, and the penalty using standard maximal inequalities. Decomposition
Define the centered losses by fθ = Lη0 (θ) − Lη0 (θ0 ),
fˆθ = Lη̂ (θ) − Lη̂ (θ0 )
Now, we decompose the loss as follows, adding and subtracting moments
G(θ) = Efθ =
+ (E − En )fˆθ̂ + En fˆθ̂ E[f − fˆ ] | | {z } {z } | θ̂{z θ̂} I: DR-Remainder II: Empirical Process III: Invariance Penalty
Term I: DR-Remainder E[fθ̂ − fˆθ̂ ] = [E[Lη0 (θ̂)] − E[Lη̂ (θ̂)] − [E[Lη0 (θ0 )] − E[Lη̂ (θ0 )] = [R(θ̂) − E[Lη̂ (θ̂)]] + [R(θ0 ) − E[Lη̂ (θ0 )]] For the first term, R(θ) − E[Lη̂ ] = E[α0 ζ0 ] − E[α̂(ζ0 − ζψ̂ ] − E[α0 ζψ̂ ] = E[(α0 − α̂)(ζ0 − ζψ̂ ] ≤ α̂ − α0 · ζ0 − ζψ̂ ≲ ∥α̂ − α0 ∥ · dΨ,inv (ψ̂, Ψinv ) G(θ)1/2 where the first inequality follows from Cauchy-Schwarz, the second inequality (final step) follows by (A5). Note that expectations here are once again taken over measures on πb × PX,E . Term II: Empirical Process The proof proceeds in four steps. First, we demonstrate that there exists a finite envelope function for our function class. Second, we invoke an empirical process bound formalized in Theorem 2. Third, we refine our bound using (A3). Step 1. Finite Envelope. Define the function class F1 = {fˆθ : θ ∈ Θinv }. This is fixed conditional on the training folds. Under bounded loss (A2), ∥fˆθ ∥ ≤ F = C2 (1 + α̂(a, X, E)). By our overlap condition (Int1), the claim follows: ∥F ∥2 ≲ (1 + E[α02 ]) < ∞ Step 2. Empirical Process Bound. With a finite envelope and (A4), we can invoke Lemma 2 to conclude that with probability 1 − exp(−s) for s > 0, 2 sup ∥(E − En )fˆθ ∥ ≲ δn,inv + δn,inv ∥fˆθ ∥2 +
θ∈Θinv
s n
Step 3. Refinement. We can bound the norm of ∥fˆθ ∥ more closely. Using Holder’s, we have ∥fˆθ ∥ ≲ ∥α0 ∥ · ∥(ℓ(θ, ·) − ℓ(θ0 , ·))2 ∥ By (Int1) and (A2)-(A3), ∥fˆθ ∥ ≲ G(θ)1/2 yielding s 2 sup ∥(E − En )fˆθ ∥ ≲ δn,inv + δn,inv G(θ)1/2 + n θ∈Θinv
Term III: Invariance Penalty θ̂ minimizes the empirical risk, so it follows that En Lη̂ (θ̂) + λinv În (θ̂) ≤ En Lη̂ (θ0 ) + λinv În (θ0 ) We then have, En fˆθ̂ ≤ λinv (În (θ0 ) − În (θ̂) = λinv rn,I = oP (1) where the last line follows by Lemma 3. Final Bound
Combining the inequalities above, we have
2 G(θ̂) ≲ δn,inv G(θ)1/2 + ∥α̂ − α0 ∥ · dΨ,inv (ψ̂, Ψinv )G(θ̂)1/2 + δn,inv +
sinv + λinv rn,I n
Note, by Young’s, 1/2
δn,inv G(θ)1/2 ≲ G(θ̂) + δn,inv
∥α̂ − α0 ∥dΨ,inv (ψ̂, Ψinv )G(θ̂)1/2 ≲ G(θ̂) + ∥α̂ − α0 ∥2 dΨ,inv (ψ̂, Ψinv )2
Hence, we can plug-in these bounds to conclude our final result.