Overlap-Adaptive Regularization for Conditional Average Treatment Effect Estimation
Authors: Valentyn Melnychuk
Dennis Frauen
Jonas Schweisthal
Stefan Feuerriegel
Venue: ICLR 2026 Poster (2026)
Keywords: causal machine learning
treatment effect estimation
Primary area: causal reasoning
TL;DR: We introduce a novel framework of Overlap-Adaptive Regularization that tackles the issue of low overlap with an adaptive regularization.

Abstract: The conditional average treatment effect (CATE) is widely used in personalized medicine to inform therapeutic decisions. However, state-of-the-art methods for CATE estimation (so-called meta-learners) often perform poorly in the presence of low overlap. In this work, we introduce a new approach to tackle this issue and improve the performance of existing meta-learners in the low-overlap regions. Specifically, we introduce Overlap-Adaptive Regularization (OAR) that regularizes target models proportionally to overlap weights so that, informally, the regularization is higher in regions with low overlap. To the best of our knowledge, our OAR is the first approach to leverage overlap weights in the regularization terms of the meta-learners. Our OAR approach is flexible and works with any existing CATE meta-learner: we demonstrate how OAR can be applied to both parametric and non-parametric second-stage models. Furthermore, we propose debiased versions of our OAR that preserve the Neyman-orthogonality of existing meta-learners and thus ensure more robust inference. Through a series of (semi-)synthetic experiments, we demonstrate that our OAR significantly improves CATE estimation in low-overlap settings in comparison to constant regularization.

Published as a conference paper at ICLR 2026
OVERLAP-ADAPTIVE REGULARIZATION FOR CONDI- TIONAL AVERAGE TREATMENT EFFECT ESTIMATION
Valentyn Melnychuk, Dennis Frauen, Jonas Schweisthal & Stefan Feuerriegel LMU Munich & Munich Center for Machine Learning
Munich, Germany
me meInychuke lmu.
ABSTRACT
The conditional average treatment effect (CATE) is widely used in personalized medicine to inform therapeutic decisions. However, state-of-the-arr methods for CATE estimation (so-called meta-learners) often perform poorly in the presence of low overlap. In this work, we introduce and improve the performance of existing meta-leamers in the low-overlap regions. Specifically, we introduce Overlap-Adaptive Regularization (OAR) that regularizes target models proportionally to overlap weights so that, informally, the regulariza- tion is higher in regions with low overlap. To the best of our knowledge, our OAR is the first approach to leverage overlap weights in the regularization terms of the meta-learners Our OAR approach is flexible and works with any existing CATE meta-leamnr: we demonstrate how OAR can be applied to both parametric and non-parametric second-stage models. Furthermore, we propose debiased versions of our OAR that preserve the Neyman-orthooonallity of existing meta-learners and thus ensure more robust inference. Through a series of (semi-)synthete experi- ments, demonstrate that our OAR significantly improves CATE estimation in low-overlap settings in comparison to constant regularization.
approach tackle this issue
INTRODUCTION
Estimating the conditional average treatment effect (CATE) from observational data a core challenge in causal machine learning (ML). Especially in medical applications, the CATE estimates help to guide personalized therapeutic decisions by predicting how different patients might respond to a given treatment (Feuerriegel et al. 2024).
State-of-the-ar art methods for CATE estimation are based on two-stage Neyman-orthogond meta- learners (Curth & van der Schaar 202 Morzywolek et al., 2023). As such, meta-learners have several practical benefits. Specifically, they are model-agnostic (Künzel et al. 2019) (i.e., they can be instantiated with arbitrary predictive models such neural networks). Furthermore, by using Neyman-orthogonal risks (Chernozhukov et al. 2017: Foster & Syrgkanis 2023), meta-learners can achieve favorable theoretical properties. In particular, the second-stage model becomes less sensitive to eтrоrs in the nuisance function estimates, which improves robustness.
However, the performance of meta-learners is constrained by the degree of overlap in the data (''Amour et al. 2021 Matsouaka et al. 2024)- that is, the extent to which patients with similar covariates receive different treatments.
our work, overlap is represented as the product of the conditional probabilities of receiving each treatment, namely, overlap weights. Overlap is often violated in medicine when patients with certain covariate profiles almost exclusively receive one treatment (e.g., due to adherence to medical guidelines). Hence, the low-overlap regions of covariate space are sparse in counterfactual outcomes, and, thus, learning CATE gets increasingly challenging.
To address issues from low overlap, existing meta-learners suggested two main approaches: (1) re- targeting and (2) constant regularization. (1) Retargeting incorporates the overlap weights into error terms of the target risks (Morzywolek et al. 2023 Nie & Wager 2021 Fisher 2024), so that the error term truncated or weighte in the low overlap regions. In contrast, (2) constant regularization aims to reduce CATE heterogeneity towards more averaged causal quantities (c. g.. ATE). While effective to some extent, these strategies have key limitations (as we show later). In the case of (1) retargeting, the fitted target models struggle in low-overlap regions: they either (i) have

---

Published as a conference paper at ICLR 2026
Stage Enttraaiing reficass frmeikore
Stage 2 Fitting model
Siage
Eikane seaple
Entroitt lnceess
Dopad
TANN
TABN
Hil
Mlage
Tiskaa sesple
CaT Vock
Ovad CATE
Figure I: Motivational example showing how OAR (in red) performs better in low-overlap regions (in yellow). Here, we used our OAR together with a DR-learner. We adapted the synthetic data generator from (Melnychuk et al. 2023) (Terain 250; see Appendix F and used kernel ridge regression (KRR) as a target model. We see that target model fitted w/ our OAR()m) (shown in red) has a better performance in the low-overlap regions, compared to a target model w/ constant regularization (CR, shown blue).
unpredictable behavior (ii) larget at different causal quantity (e. g., R-IIVW-leamers (Nie & Wager 2021 Fisher| (024) with the constant regularization yield a weighted average treatment effect [WATE] in the low-overlap regions). Further, (2) constant regularization does not take into account the degree of overlap and "blindly" regularizes all the regions of the covariate space.
In our work, introduce Overlap-Adaptive Regularization (OAR), novel approach that builds on top of existing two-stage meta-learners and tackles the low-overlap issue through adaptive regularization. Our OAR helps prevent over- and underfitting of the target models by varying the amount of regularization depending on the degree of overlap (see the illustrative example in Fig. [T] As a result, our OAR applies stronger regularization in the regions with low overlap and weaker regularization where overlap is high (i. e., when the propensity scores are close to 0.5).
Our OAR thus addresses the above limitations of existing meta-learners. First, (1) unlike targeting. our OAR makes the predictions of the OAR-fitted target models smoother in low-overlap regions ei. enforces simpler models in those regions). Second, (2) unlike the constant regularization, allows for more CATE modeling flexibility in the overlapping regions (e. g.. for DR-learner). Also, unlike -eaanese, it can yield the average treatment effect (ATE) the low-overlap regions (which is arguably more meaningful causal quantity than the WATE). To the best of our knowledge, ours is the first approach to address the low-overlap problem by directly adapting the regularization term in the target risk.
Our OAR is flexible and can be applied together with any two-stage meta-learner. We provide several versions of our OAR approach: (a) for parametrie target models (e.g., neural networks) and (b) for non non-paramttric target models (c.g., kernel ridge regression). For (a), we introduce two practical implementations via: (i) OAR noise regularization and (ii) OAR dropout. In addition, propose a one-step bias-corrected (debiased) estimator of our OAR. This correction is important because it makes our OAR first-order insensitive errors in the estimated overlap weights (which is especially relevant in observational studies where the ground-truth overlap weights are unknown). As a result, when combined with debiased OAR preserves their syyman-othoogonality We further provide an extension of our OAR to (b) non- -araameeetiicc target models (e.g., kernel ridge regressionn) Appendix C
learners (c. g., DR-, R-, and IVW-learners). our
In sum, our contributions are as follows (1) We introduce a novel approach, which we call Overlap- Adaptive Regularization (OAR), to address the performance of the existing CATE meta-learners in low- overlleraap regions. (2) We propose several versions of our OAR for both parametric and non- parametric target models, as well as debiased version that preserves Neyman-orthogonall show empirically that our OAR improves the performance in CATE estimation over other alternatives.
(3) We
2 RELATED WORK
In the following, briefly review the existing methods for CATE estimation and the ways they tackle the low-overlap issue. For more extended overview of related work, we refer to Appendix A
Two-stage meta-learners. State-of-the-art methods for CATE estimation can be broadly divided into two general categories: (a) plug-in learners (also known as model-based methods) and (b) (two-stage)
meta-learners KKünzel et al. 2019 Curth & van der Schaar 2021b Morzywolek et al. 2023
Frauen
Code available at wwt.bbbeehtttge////h.hb.../////a 199.a/nntnn199991ln
2

---

Published as a conference paper at ICLR 2026
al. 2025). Here, we refer to the overview of (a) plug-in learners to Appendix|A and rather focus on (b) meta-leamers/ Meta-learnert aim to find the best projection of CATE on a second-stage (target) model class and require the estimation of the nuisance functions at the first stage (Künzel et al.) 2019). Importantly, they are fully model-agen model-agnostic meaning that any MI. model can be employed at either of the stages. Also, they can possess many favorable theoretical properties if they Neyman-orthogonall CChennozuukov et al., | meta-learners include DR-learner (van der Laan| 2006; Curth et al. 2020 Kennedy 2023), R-learner
Foster & Syrgkanis 2023): notable examples of Neyman-orthogonal
(Nie & Wager, 2021), and IVW-learner (Fisher (024). Jeyman-orthogome
is property of the target risks that makes it first-order insensitive to the errors of nuisance functions estimation, and,
therefore, in this work, we only focus on those.
How meta-learners deal with low overlap. Low overlap poses serious problem for any causal effect estimation (D'A Amoum et al. 2021 Matsouaka et al. 2024), including CATE. In the low-overlap regions, meta-learners mainly suffer from high variance of the pseudo-outcomes (Morzywolek et al., 2023): it stems either from bad extrapolation of the first-stage models or from large inverse propensity scores. There are two general ways to tackle low overlap: (1) retargeting and (2) constant regularization.
(1) Retargeting. The retargeting approach estimates CATE only for a sub-population by modifying the target risk/loss (Morzywolek et al. 2023 Matsouaka et al. 2024). For example, trimming and truncation (Crump et al.
2009) discard too low propensity scores of the DR-learner loss; and overlap- weighting (Crump al. 2006 Morzywolek et al. (202) of the target risk yields either R-learne (Nie) & Wager [021] Hess et al. (2026) or IVW-Iearner (Fisher 2024). However, (I) retargeting on itself
does not regulate how target models would generalize beyond the target sub-population. Therefore, the above-mentioned works suggested combining it with a (2) constant regularization.
(2) Constant regularization. Constant regularization improves low-overlap predictions by forcing lower CATE heterogeneity (Morzywolck et al. 2023) in the whole covariate space. Yet. this approach also has drawbacks. For example, when combined with the DR-learner, it does not distinguish the variability of pseudo-outcomes in high- and low-overlap regions. DR- DR-learner, thus, can overfit and underfit at the same time due to the constant regularization. On the other hand, when combined with R-IIVW-learners, this approach leads to a different causal quantity (i. c., WATE) when too much regularization is applied.
Adaptive regularization in traditional ML. As discovered by Wager et al. (2013), dropout (Hiton cet al.] 2012; Srivastava et al., 2014) and noise regularization (Matsuoka) 1992] Bishop 1995) can be seen instances of the adaptive regularization. The authors have shown that for generalized linear models, dropout and noise regularization are first-order equivalent to the 1/2 regularization applied to the features scaled with inverse diagonal Fisher information matrix. This result
later extended
for dropout regularization in NN-based models (Mou et al., 2018 Mianjy et al.) 2018 Mianjy & Arora 2019 Wei et al. 2020 Arora al., 2021); for noise regularization in NN-based models
(Rothfuss et al. 2019 C Camuto et al. 2020); and for other types of regularization (Dieng et al. 2018 Mou al. 2018 LeJeune et al. 2020 Zhang et al. [202]] Nguyen et al. 2021]. In our paper, we also draw connections to the seminal results of (Wager et al. 2013). However, we provide -for the first time - the connection of adaptive regularization to CATE estimation. To the best of our knowledge, overlap weights have not been used to explicitly define regularization for CATE estimation.
PRELIMINARIES
Notation. Random variables are denoted by uppercase letters such Z, their realizations by lowercase letters such as 2, and their sample spaces by calligraphic symbols such as 2. We write P(Z), IP(Z 2). and E[Z)) to refer, respectively, density at 2, and its expectation. We denote an l2 norm as || = +oG+ reproducing kernel Hilbert space (RKHS) norm as I| (H, for f HK. where Hx is an RKHS induced by a kernel K (.,,)). We employ two nuisance functions: the propensity score for treatment A is 7 (z)
distribution of Z. its probability mass or
x² for x E IR:: a
P(A = 1 X z). and a conditional expected outcome for the response Y is
BB naively estimating f(x) = (x) (ar), plug-in learners suffer from so so-called plug-in bias Kennedy 2023) (c. g.. Ao (2) is badly estimated for treated population). This is addressed in (two-stage) leamerr Unlike the plug-in in learners, meta-learnes allow to solve the variance trade-off for the nuisance functions and the target CATE separately (Morzywolek et al. 2023). Hence, we focus (two-stage) meta-leamers throughout our paper.
3

---

Published as a conference paper at ICLR 2026
Ma E(Y = a). We also consider a marginalized conditional expected outeome
u(z) = EY = X = z) and overlap weights v(z) = Var(A X 2) as alternative nuisance
function (yet, they can be expressed through the former two: μız) = (1 -1-110 (z)=(()))o(()) + π ( )r (æ)
and (xx) = T (1 π (.)))). Throughout, we work within the Neyman-Rubin potential outcomes
framework Rubin 1974 Specifically, Y[a] denotes the potential outcome under the intervention.
Problem setup. To estimate the CATE, we rely on an observational sample DD = (a(9), )/ where X E X C R are high-dimensional covariates, A E {0, 13 is a binary treatment, and Y EYEC R is a continuous outcome. For example, in cancer care, Y measures tumor growth, A indicates whether chemotherapy is administered, and X records patient attributes such age and sex. Furthermore, bution P(X,AA Y). We also denote a correlation matrix of the covariates X as Σ = E[XX]] and a X()-weighhed correlation matrix as Σi = [I X " for function À [0, x X R.
assume the n triplets in D are drawn i.i.d. from the joint distri-
Causal estimand and assumptions. We are interested in estimating the conditional avveaage treatment effect (CATE): T(x) [1Y Y(] X = -r.. To consistently estimate it from the observational data D, we make the standard causal assumptions of the Neyman- Soyyman-Rubin framework (Rubin) 1974): (i) consistency: Y [[A]= Y; (ii) strong overlap: P( < (XX) < 1- = for some E (0. 1/2);
and (iii) unconfoundedness: (Y (1) 11 A X. Then, under the assumptions (i)(iii), the CATE identified from P(X, A, Y) as T(x) ) = p4 (x) No (r). In this work, we estimate the CATE from
observational data 2> with meta-learners (Künzel et al. 2019: Morzywolek et al., 2023).
Meta-learners for CATE, Formally,, two-stage meta-learners aim to find the best projection 9 (x) of the ground-truth CATE 7(x) on pre-specified model class G = {g X + IR by minimizing a target risk Lig. n) wrt. g. Here, = (Ho. (4, T) are nuisance functions: they are fitted at the first stage and then used at the second stage to learn the optimal 9" arg min Lgg, 17). The majority
of existing CATE meta-learners can be described by the following target risks, which have the same minimizers given the ground-truth nuisance functions:
Original risk:
C(g. n) -E ((x( (4) (x) mo( - 9(X)) ((x) (1)
Neyman-orthogonall risk: C(g. n) = E (X) (6(2.n) - -(x))] + A(o:
(2)
regularization (A)
errer (E)
where (xT(X)) 2 is a weighting function, p(A, (X)) = (A- (A (T(() + 2 0
a debiased weighting function, o6Z, 17) is a pseudo-outcome with a property R[Ø 7) x
(ar). While the two risks (original and Neyman-orthogone the ground-truu truth nuisance functions 17. they yield significantly different results 9 when the nuisance functions are estimated ij.
have the same minimizers g" given
Neyman--rtthognall meta-learners. the following, we will focus on three Neyman-orthogonal
meta-learners (Eg. 2 ): DR-leareer (Kennedy (023). R-learner (Nie & Wager 2021). and IVW-
learner (Fisher 2024). The DR- -R-learner (Kennedy) 2023) is given by w (X)) = p(A, π (X)) = 1
and dx Z, n) 2021) by we T (x)) = v(X), pl A, (X)) = (4 - (X))² and (Z. = (Y =(x)))(A (())::
n (Y (x(X))) + nn (X) N0 (X); the R-learner (Ne & Wager
and the IV IVW-leamne (Fisher 2024) by a combination of the former: the weighting functions of the
R-learner and the pseudo-outcome of the DR-leamer (see details on meta-learners in Appendix B
Low overlap. We speak of low overlap, whenever either (xx) or (1 (x)) (and thus v(x)) are close to O. Conversely, perfect overlap regions have Tzx) = 0.5 and u(x) = 1/4.. Importantly, low overlap negatively affects the convergence of any meta-learner. For example, for the De-learner, inflates the inverse propensity scores and, for R-IVW-lleamers, it retargets the target risk at a different quantity than CATE (Morzywolek et al. 2023).
Constant regularization The regularization term in Eq. ( A = ((X),, should be spec- ified depending on the target model class g. For example, if the second-stage model (a) para- metric. namely G = {g(:; B.e) X' R 3 E R2, c E R). (z-regularization is popular choice: A(g: P(X)) X 1| Here, is an intercept, and X 0 is a regularization constant. Similarly, for a (b) non-parametric second-stage model (e. g., kernel ridge regression), G is the
3 Here, the regularization term A might also depend on P(X) (c. g., standard dropout implicitly depends on the correlation matrix E). In our context, we call it constant regularization does not depend the overlap.
4

---

Published as a conference paper at ICLR 2026
RKHS Hk+c and ((g:((X) >لااااا Here (with slight abuse of notation), is an added constant kemnel. It is easy to see that, in both cases (a) and (b), increasing À - oo leads to 9" - E (())) /)E/BEp(AA xX))] c* (see Remark as the intercept/constant kernel is not regularized. Yet, the constant regularization À (i) does not directly address the low-overlap issue for DR-learner or (ii) leads WATE in low-overlap regions for R- and IVW-learners This motivates our core idea of an adaptive regularization that depends on the distribution of the covariates and the treatment P(X, A), and the level of overlap vXX).
[ in Appendix [.2).. This happens
OVERLAP-ADAPTIVE REGULARIZZTTON
In the following, we introduce our approach of Overlap-Adaptive Regularization (OAR) (Sec. 4.1p and several specific versions for (a) parametric target models (Sec. (). Further, we provide an extension for non- non-parametric versions of our OAR in Appendix A Proofs are provided in Appendix D
4.1 GENERAL FRAMEWORK
Here, we define our Everlap-Adaptinn Regularization (OAR), a novel general approach that (1) ad- dresses the low-overlap issue of existing meta-leamers, and (2) is model-agnostic.
Definition 1 stage model g(-) E G and a target risk Lig. 1i) = E+A (Eq. (2) overlap-adaptive regularization
regularization (explicit form)). For a meta-learner with second-
(OAR) in explicit form is given by
NOAR A(y A);
where M1)) > 0 is a regularization function that defines the amount of the regularization and is
proportional the inverse overlap: X(v) x 1/v. We further distinguish three general classes of
regularization conttttoonfnniiono functions: multipiccative (m), logarithmic (log), and squared multiplicative (m²),
Am Am(((x))= = 1/4u(x)-11
Atog (v(x)) = - (x)(4v())); Amz (v=()()))= = 1/166(z)²-1. (4)
Our OAR explicitly depends on the overlap through the regularization function which is the main difference from the constant regularization.
Interpretation.
Informally, our OAR increases regularization in the regions of the covariate space X when v(x) + 0). Analogously, the regularization becomes
with low overlap (namely, A(r) smaller when perfect overlap is achieved e., A(v) - 0 when v(z) + 1/4). This introduces a desired behavior in a practical application. For example, in medical context, low-overlap regions imply higher certainty about the treatment decisions (as optimal treatment might already be known there). Our OAR then allows to focus the model flexibility the overlapping sub-population (namely, the individuals for whom the CATE/optimal treatment is unknown).
Implicit form. Notably, our OAR can also be defined in the implicit form: it can enter through the error term of the target risk &OAR (c. g., noise regularization and dropout). Still, as we will demonstrate later, the two formulations are equivalent for linear models. That is, is possible to find an equivalent form of the target risk with the original еттоr term & and the regularization term in the explicit form AOAR- This equivalence can also be partially extended to some deep neural networks: noise regularization (Camuto et al.] (020) and dropout (Wei et al.) 2020) were shown to have an explicit, first-order equivalent form.
Flexibility. Our OAR can be combined with any (Neyman-orthogonal) meta-learner. Also, we intentionally did not specify the dependency on g and the observational distribution of P(X, A). This was done as we want our OAR to be model-agnostic and fit into a wide range of the second-stage model classes G. For examplc, many standard regularization techniques like dropout and noise regularization can have very different dependencies g and PXX, A) in their explicit forms (Wager
et al. 2013 Camuto et al. 2020 Mianjy & Arora 2019
Wei et al. 2020).
4.2 DIFFERENCE FROM THE LITERATURE
Difference to retargeting. A natural question arises on whether our OAR is related to retargeting, a standard approach in meta-learners to handle low overlap (see Sec. 2 Specifically, retargeting is
'Our regularization function can be seen as an example of a selection (or tilting) function Li et al 2018 Assaad et al. 202T Matsouaka et al. 2024). So far, these were only used in the error terms of the target risks (e. g., R-Iearnr uses bias-corrected overlap weights).
5

---

Published as a conference paper at ICLR 2026
Table 1: Summary of OAR versions. Here, & is the error term defined Eq. and A(v) is the regularization function of OAR (see Definition D.
Instantiation OAR
Implicit form (EnaR) explicit form (Anax)
Equivalent explicit form a linear model class g
Noise regularization
SOAR = (Eq.
5
AoaR = 0
EOAR & AOAR = ||31 Blp(A,x(X)) A(x(X))
Dropout
EOKR
(Eq.
AcoR 0
EOAR =
AOAR = 33 diag (loc.)) Manl
RKHS
EOAR &
Aoar
A(p)
Undefined
implemented in both R- and IVW-learners, as they down-weight their error term of the target risk proportionally to overlap: (1 xx))) X x] B[(A π (X))² X x] =(x). On the other hand, our OAR up-weights the regularization term wrt. overlap. However, both approaches, in general, lead to different risk minimizers: while the re- re-weighted error term incorporates overlap with all the aspects of the observed distribution P(X,, A, Y), our regularization term in OAR combines overlap only with P(X, A). Interestingly, the two approaches match only in very simple case (as we will show later) when the propensity score and, thus, the overlap are constant.
Difference to balancing. Another way to tackle low overlap suggested by balancing repre- sentations with empirical probability metrics (Johansson et al. 2016; Shalit et al. 2017 Johansson et al. 2022 Assaad al., (202T) in neural network (NN)-based plug-ins (see Appendix Ap. Here, the average amount of regularization is proportional to a distributional distance between untreated and treated covariates, ( A = 0); A 1)). In our case, the average amount of regularization can also be represented through distributional distances, yet different from the one used in balancing regularization.
Proposition 1 (Average regularization function as distributional distance). The average amount of overlap-adaptive regularization E. [[((v())) is equal to or apper-bounded by I-divergences between P(X) and (PX A a) for E {0, 1}.
We immediately see that our OAR is implicitly based on the distributional distances between P(X) and P(X A = a), which different from those used in balancing. Furthermore, OAR is thus simpler in implementation than balancing because we only need to estimate the propensity score but not the distributional distance for high-dimensional X.
5 INSTANTIATION OF OUR OAR
In the following, we provide several versions of our OAR for parametric target models, and we carefully tailor existing regularization techniques so that they become "overlap-adaptive" (see the overview in Table D. For each version, we also (1) show an equivalent explicit form AOAR when the target model is linear and (2) derive a debiased (one-step bias-corrected) version of the regularization. The latter beneficial to remove the first-order dependency on the estimated propensity score. We also provide version for non-parametric target model in Appendix All proofs are in Appendix [.I]
We consider target models in the following parametric form G (g(:3,cc) : X
RR, c E RJ. where B are parameters to be regularized and c is an intercept. For this very general parametric class, we tailor two general regularization techniques based noise injection: (i) OAR noise regularization and (ii) OAR dropout. We also consider w.l.o.g. that noise is injected into the inputs of (if g is an NN, then noise can be injected into any layer, Fig. 3 in Appendix E
5.1
OAR NOISE REGULARIZATION
Our OAR with Gaussian noise regularization is given by
C (g. 1)) = EOAR B / I -(x
(5)
where NN0 o²) is normal distribution with variance o².
Thus, by construction of our OAR noise regularization, the variance of additive noise is proportional to the inverse overlap o² x 1/v(z), and the model g(x² is regularized more in low-overlap regions. We further show an explicit form of OAR noise regularization AOAR for linear models g.

---

Published as a conference paper at ICLR 2026
Proposition 2 (Explicit form of OAR noise regularization in linear g). For a linear model g(z) B x + C, OAR noise regularization has the following explicit form AOAR:
Loiggg g. 11) = + Nour = + | ((AA ((X)) M((X))):
(6)
where & is given by the original ervor term from Eq.
Interpretation. We observe that, for linear models, OAR noise regularization coincides with a ridge regression with the constant regularization À Eip(A, (X)) X((())] However, for other, more complex parametric models, the explicit form is more complicated and is very different from ly (e. g.. for NNs, noise regularization in explicit form depends on the Jacobians wrt. parameters (Camuto
et al. ).20
5.2
OAR DROPOUT
Our OAR dropout is given by:
FoAR (g. 7) EOAR =F meeeeeoo))) (XX) -x(xx
where o is an clement-wise multiplication; p(v) X(v)/(A(z) + 1) (0. 1) is dropout
probability; (rop(p) is sampled from a scaled Bernoulli distribution (E 0 with probability p and { = 11(1 - p) with probability 1 p).
Given the definition of our OAR dropout, is easy to see that p x 1//(ax). This means, the dropout probability is p 0 in high-overlap regions (v(x) 1/4), and 1 in low-overlap regions (v(x)~ -).
Interestingly, both the regular (Wager et al. 2013 Bartlett et al., 2019) and OAR dropouts are significantly different from the lg-regularization, even for linear models g. We show with the following proposition.
Proposition (Explicit form of OAR dropout in linear g). For a linear model g(x) = BJ c, OAR dropout regularization has the following explicit form ACAR'
Coselg, 1)) & + Nour + B diag (oo())))AA())))
(8)
where & is given by the original error term from Eq. ② A(v) p(v)((1111)))))) - -(()). and diagl:] zeroes out all but the diagonal entries of a matrix.
Interpretation. Proposition 3 motivates our choice of p(v) as A(v)/(A(v) + 1) so that X(v) p(v)/(1 -p())). Also, we immediately see that, for linear models g. OAR is regularization but an overlap-dependent quadratic form for B. That is. our OAR dropout scales each B, prior to applying lg-regularization. Specifically, our OAR dropout in linear models is equivalent to ridge regression where each feature is scaled down proportionately to the product of
12-
the inverse overlap and its own second moment: X, = x, (El(A, ((X)) A(v(X) X3). For other parametric models, the explicit form of OAR dropout becomes more complex (c. g., the explicit form of the standard dropout in NNs has l2 (2-path regularizers and rescaling invariant sub-regularizers
(Mianjy & Arora 2019:
Wei et al. 2020)).
Implicit and explicit forms. Importantly, Propositions 2 and alsso show the effect of OAR applied on top of the retargeted learners if OAR is presented in the implicit form EOAR- For example, when multiplicative OAR noise regularization is used with R-IIW-lleaneers, they result in constant amount of regularization
low-overlap regions (i. e., (i T (X)) Am (v())) = 1/4 v(X) + 1/4 given the ground-Hng fround-truth -ruth nuisance functions). This suggests that, if we want to adaptively regularize retargeted learners with OAR noise regularization, we need to employ the squared multiplicative regularization function An2 (so that Elp( A, n (X)) Am (v(X))
in low-overlap regions).
5.3 DEBIASED OAR FOR PARAMETRIC MODELS
Here, we provide two debiased (one-step bias-corrected) versions, which we call dOAR noise regularization and dOAR dropout. Debiasing (van der Vaart 2000 Kennedy, 2024) is beneficial to remove the first-order errors from the estimated propensity score (())) = (z) (1- *(x)). Namely, our original OAR from Eq. (5) and (7) might be overly sensitive to the misspecificatiio of the overlap
weights (e. g., when the propensity score is badly estimated).
7

---

Published as a conference paper at ICLR 2026
Proposition 4 (Debiased OAR). Under the continuous differentiability of g(x; B. c). (1) debiased
OAR noise regularization and (ii) debiased OAR dropout are given by
Camulo n) = Cinxlg, n) E 1/ A 2) for (+),0E)
A;
(9)
(; -20(()) mo(X) - 9 + (e(o)( X,,
(10)
(x: -(X)(a mo - - e))) p(P(X)) [FP((()))
-20(X)(u1 - Velg)(X,E) IFPl((s)):
(11)
where COAR are from Eq. 5 and (7): Velgi is gradient wrt. E: and LF 2) is an efficient foceeencffuntioo function (see Appendix D.2 for further details). Furthermore, by construction, SiAAR is a
influence
Beymann-ortoggone risk.
Proof. For debiasinn, derived the efficient influence functions using chain rule together with reparameterizatioo and REINFORCE tricks. The full proof is in Appendix D
Importantly, after debiasing, our OAR recovers the property deyymnn-rtthogonalii (Chernozhukov et al. 2017 Foster & Syrgkanis 2023) when combined with the standard Neyman-orthogonal learners. Furthermore, we can show that, under some additional conditions, our OAR/dOAR are guaranteed to outperform the constant regularization (CR).
Proposition 5 (Excess prediction risk our OAR/dOAR dropout with linear second-stage model). The excess prediction risk of the DR-learner with the linear recond-stage model and dropout regular- ization has the following form:
119 - 9" 3"11i, = =E(B" xxTX) & tr (I(+ (2 (E+I) +8
T + +R 9),
(12)
biss temm
varianoe
where L'CR = Al for the CR, L'OAR = diag Exw) for the OAR/dOAR. Then, under (i) a conditional variance assumption (=conditional variance of the outcome is constant), the variance term for OAR/dOAR is less than or equal to the variance term of the CR. Also, under (ii) a twwoverlap-low- heterogeneity inductive bias (LOLH-IB), OAR/OOAR do not increase the bias term too much.
Proof. We used bias- vias-variance decomposition of the excess prediction risk for linear models. Then, we showed how assumptions (i)- help to reduce each term for our OAR/dOAR in comparison to the CR. The full proof is in Appendix D
We provide the full statement and the full proof of Proposition 5 in Appendix D Arguably, both assumptions of Proposition 5 are reasonable: (i) The conditional variance of the outcome, Var(Y the variance of the DR
while (ii) LOLH-IB is often assumed to simplify causal MI. (Curth & van der Schaar 202 Melnychuk al. 2026) (see Appendix [A). Importantly, P 5 Eq. (12) apply to any level of overlap. However, specifically for the low-overlap setting (i. e., with larger values of 1 1/v(z)), it is fair to assume (i) the conditional variance assumption, as the variance of the DR pseudo-outeome can be considered proportional to the inverse overlap.
A). becomes nearly constant comparing
5.4 NON-PARAMETRIC TARGET MODELS: OAR RKHS NORM
In the following, we introduce an instantiation of our OAR for a very general class of non-paran -anamarrettric models that belong a reproducing kernel Hilbert space k induced by a sum of arbitrary kernel K(--) -) and a constant kernel K (-.-)
Our OAR RKHS norm for a target model E HK+c can be instantiated as weighted kernel ridge regression (KRR) with modified, OAR-basee RKHS norm:
Cok (g. 7) = E AoxR = E +++++VVxu Xx
(13)
where given by the original error term from Eq. (2). and we assume that A(vvg E Hk for every g Hk+o (this assumption is required so that the modified RKHS norm is well-defined).
Here, the regularization function / A(v(z)) adaptively regularizes function gzx) and is known multiplier of the RKHS Hx (Szafranicc 20000 Paulsen & Raghupathi 2016). Then, under special conditions, the weighted KRR with OAR- based RKHS norm has well defined solution.
8

---

Published as a conference paper at ICLR 2026
Proposition 6 (Kernel ridge regression with an DAR-based RKHS norm). Let v Nlvig HK for every 9 Hk++ Then, the minimizer of the target risk g arg min y Hxx C (g, (g,n)) n)) is in
Ho+c and has an explicit solution.
We defer the full formulation of Proposition 6 and further discus- Furthermore,
100-
sions to Appendix C Proposition 7 in Appendix C shows a result similar to Proposition | for our OAR RKHS norm regulariza- tion. Specifically, we showed there that our OAR RKHS norm improves over the CR under the analogous conditions for the RKHS: (i) con- ditional variance assumption and (ii) low-overlap- horip--oo--eetee -ow-eeeroogyeetty inductive bias.
10-
04-
Yobe my Rlomee
148
Droputt Balceee
11018 sor Caaemm-Blarrr Bilerre
0.-
06-
00-
02-
10315 -V700o
5.5 IMPLEMENTATION DETAILS
00
The implementation of our OAR pro- ceeds in two stages (e. g., see Fig. 3 in Appendix D for the (a) paramet- ric OAR instantiations with NNs).
18
amagh.
OAR
DAB
Oeack
Twetfer
In stage 1, we estimate the nui- sance functions n.. For this, we fully-connected
Figure 2: Results for IHDP dataset experiments. Reported: median rPE 1 se over 100 runs.
fit cross-validate
NNs. Then, in stage 2, we fit target network using empirical versions of the target risks 7 in (a) and the KRR solution 9 in (b) to yield a target model (= CATE estimator). Furthermore, in both settings (a) and (b), we rescaled the regularization function X(v) (or f(vi) so that OAR can be comparable with the constant amount of regularization. For our dOAR, we additionally trimmed too large absolute values of the debiasing term C"; this helped to achieve a better stability of training the target models. We provide further implementation details in Appendix
Table 2: Results for 77 semi- synthetic ACIC 2016 experiments for the DR-lcarner. Reported: % of datasets, where our OAR/dOAR sig- ni nificantly outperforms CR (a 0.1 with 15 runs per dataset).
EXPERIMENTS
Setup. We follow prior literature (Curth & van der Schaar 2021b; Melnychuk et al.) 2023) and use several ssemi-)synthette datasets where both counterfactual outcomes [0] and Y [[] and ground-truth CATE are available. Specifi- cally, we used four datasets for benchmarking (see Appendix for details). For all four, report an out-sample root pre- cision in estimating heterogeneous effect (rPEHE or an improvement of our OAR over the baseline difference of the former (Ar)EHEou).
Reg.
Noise reg.
Dropout
A/p
0.05
0.3
Approach
OAR /Ple
14.29%
29.87%
dOARI Asox/ /oox/PP Plog?
7.79%
64.94%
OAR( (AR(Aan/Pull (aa)
31.17%
41.56%
Baselines. We compare all versions of our OAR and our debi-
dOAR(Aua FFa))
57.14%
70.13%
ased OAR (dOAR). As a baseline, we use only a comparable regularization strategy for meta-learners, namely, constant reg-
OAR(3 S.nne//Pue /Pnz)
27.27%
16. 16.88%
doAR(Am
76.62%
64.94%
ularization (CR) (Morzywolek et al., 2023). 3). In Appendix G we also report the results of other, not directly comparable
Higher better (improvement over the
baseline in <500% of runs in green)
baselines (e. g., trimming and balancing). Here, we compare how different amounts of regularization
learners: (i) DR-learner (Kennedy 2023). (ii) R-learner (Nie) & Wager 2021), and (iii) Wl--lrr (Fisher, 2024). For fair comparison, we rescaled our OAR/ dOAR so that they, on average, coincide with the CR values (see Appendix E)
work with three
IHDP dataset. The IHDP dataset (n = -672 + 75; dz 25) (Hill 2011 Shalit et al., 2017) is well-known to have severe overlap violations (Curth et al. 2021). Results. We show the results of the experiments with the IHDP dataset in Fig, 2 Therein, our OAR/dOAR are particularly effective for the De-learner and large regularization values. Notably, the best performance for every meta-learner and regularization type is achieved by some version of our OAR/ROAR.
9

---

Published as a conference paper at ICLR 2026
Table 3: Results for HC-MNIST experiments for OAR/dOAR (Am/ Dm). Reported: rPEHE ( (ArPEHEou)); mean + std over 30 runs.
Reg
Noku
Deegue
A/p
0.85
0:25
0.3
6.9
Learner
Agjrocc
(R coea0
02522 = 2000
0.44] 0:077
0,880
6:466 0.036
0032
0.11
0L025
0.743 = 0000 0939)
0736
0136 0,015)
0.033 0.015:
0742 03038 6:044)
0082 0,144)
0.201
0.025 0.0133
0.31 a 0.05 (-0821)
0/112:1
0/013 (-00529)
0.844 0.827 (-0.02))
6015
(-00055)
(UAIS
00001 (-0001))
0.002
9.755 0.05
0.70 4 0.000
0.64 0.007
6:720 4 0003
0.711 0.577
0.9
0.008
0.21
0011 0.034)
0,000
0.700 0.244
0015)
0.000 0.01
CR (3 connt
1.11 4 245
1.10 00335
LIPT I ILTH
1.117 0.71
1 0.281
IVW
DARL
1,099 a 237 (-011)
1,011
0.225 (-0,00))
1.964 1 (C215 (-1044)
1,311 1 0.29 (-00005)
1/01
0:211 (-10066
0,97 1 0.213
(-0 1160
1.105
= 0.239 0.016
1,058
0.217 0.044)
0,978 0.222 0.009
1,130 0:000)
1,100 + 0.293 0,000)
L
0.235
0.0960
Oracle
Tame: Nener (hor hed secend Best indirtrall Chang ovar foactre Traddcts Oaggrisoo Impromment goom, aprnca sernaning in miL 0T1S)
ACIC 2016 datasets. ACIC 2016 collection (Dorie et al. 2019 contains 77 semi- datasets (n = 4802, d, = 82) with varying overlap and CATE heterogeneity. Due to the high-dimensionality of covariates, we exclude RKHS norm regularization from the experiments. Results. Results for the DR-learner are in Table 2 as our OAR/dOAR were most effective in the combination with the DR-learner. Here, different versions of our OAR/GOAR lead to a high percentage of significant improvements over the CR. Furthermore, our dOAR often leads to a significant improvement in more than half of the datasets.
HC-MNIST dataset. Finally, we adopted a high- gihh-imennsonal HC-MNIST dataset (d, = 784 + 1) (Jesson et al. 2021), which naturally suffers from low overlap (due to the dimensionalityy. Results. Table provides the results for OAR/dOAR with the multiplicative regularization function (results for other regularization functions are in Appendix G.. Here, observe that our OAR/dOAR significantly improves the performance of the CR + TR-IR-/IVW-leanee best performance for every regularization value is always achieved by some version of our OAR/dOAR. This proves the effectiveness and scalability of our approach.
in the majority of cases. Notably, the
Additional results. In Appendix G we additionally report the results for the synthetic data from Fig. [1] Also, we report the results of other, not directly comparable baselines (e. g., trimming and balancing). There, we vary regularization hyperparameters responsible for addressing low overlap.
On the choice of the regularization function. Our results support the multiplicative regulariza- tion function as the most effective choice for our OAR. First, the proof of our Proposition 5 in Appendix suggests that, under the conditional variance assumption, the variance-optimal shape for adaptive regularization scales as fractional power of inverse overlap, namely A(v) xx nー This lies between the logarithmic dependence (looser penalization) and the multiplicative dependence (stronger penalization), making the latter a practical, more robust variant of the regularization function. This choice of a stronger regularization function can also be particularly relevant for non-linear target models (e. g., neural networks), whose effective variance typically grows with model complexity. Second, Corollary C in Appendix 2 shows that our OAR based on the multiplicative regularization function combined with the DR-learne is equivalent to the CR combined with the R-learner in kernel ridge regression settings. Given the well-studied effectiveness of the R-learner, this offers theoret- ical support for multiplicative regularization. Finally, our extensive empirical evaluation in Sec. [ and Appendix C consistently identifies the multiplicative regularization function as the strongest performer across diverse benchmarks. Collectively, these arguments justify our recommendation of the multiplicative regularization function as the default regularization strategy for our OAR.
On the best combination. We empirically found the combination of our OAR/OAAR noise regular- ization//ropout with DR-learner to be consistently good across all the benchmarks. The main reason for this is that our OAR/dOAR in combination with the DR-learner achieve just the right balance between the high-variability of the aseudo-outcome and the regularization strength (as suggested by the assumptions of Proposition 55). On the other hand, R- and IVW-learners in the combination with our OAR/dOAR might over-regularize the low-overlap areas, as the overlap already downscales the error term of the target loss.
Conclusion. In this paper, we introduced a novel approach for regularizing two-stage meta-learners: Overlap-Adaptive Regularization. Our OAR adaptively sets the regularization depending on the overlap so that low-overlap regions are regularized more. We showed that this approach is more effective than the existing constant regularization techniques. OAR performs best when low overlap coincides with low CATE heterogeneity (this can be seen an underlying inductive bias). Such an inductive bias is often meaningful in practice: in the absence of ground counterfactuals (i. e., in low-overlap versions), simpler models for the CATE may be preferred.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
Overlap-Adaptive Regularization (OAR) is designed to make conditional average treatment effect (CATE) meta-learners domains such as personalized medicine, where CATE estimates inform therapeutic choices, better behavior in low-overlap regions can translate into safer, more effective, and more equitable care decisions. Beyond healthcare, the technique may help policymakers or social-science researchers draw fairer conclusions from observational data by making CATE estimation more stable.
reliable when data exhibit poor treatment-control overlap. In high-stakes
REPRODUCIBILLIY STATEMENT
We have taken several measures to ensure the reproducibility of our work:
Algorithms. We provide full algorithmic descriptions, including pseudocode for adaptive regularization methods. Hyperparameters, update rules, and initialization strategies are explicitly detailed.
Experimental validation. We describe datasets, architectures, hyperparameters. and eval- uation procedures detail. We also released the code and experiment scripts facilitate verification.
The results from our paper can be fully reproduced using publicly available tools and the released supplementary materials.
Resources.
Thus, all results can be independently verified based on the text and accompanying resources.
LLMs USAGE STATEMENT
We used ChatGPT during the preparation of this paper in a limited way, primarily for language editing. All statements were verified independently by the authors. ChatGPT was used generate new research ideas, algorithms, or experiments.

---

Published as a conference paper at ICLR 2026
Acknowledgments. This paper is supported by the DAAD program "Konrad Zuse Schools of Excellence in Artificial Intelligence". sponsored by the Federal Ministry of Education and Research. S.F. acknowledges funding via Swiss National Science Foundation Grant 186932. This work has been supported by the German Federal Ministry of Education and Research (Grant: 011224082..
REFERENCES
Yaser S. Abu-Mostafa, Malik Magdon-Ismail, and Alsuan-Tier Lin. Learning from data, volume 4. AMLBook New York, 2012.
Ahmed Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task Gaussian processes. In Advances in Neural Information Processing Systems, 2017.
Ahmed M. Alaa and Mihaela van der Schaar. Bayesian nonparametric causal inference: Information rates and learning algorithms. IEEE Journal of Selected Topics in Signal Processing, 12:1031-1046, 2018.
Raman Arora, Peter Bartlett, Poorya Mianjy, and Nathan Srebro. Dropout: Explicit forms and capacity control. International Conference on Machine Learning. 2021.
Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin. Counterfactual representation learning with balancing weights. In International Conference on Artificial Intelligence and Statistics, 2021.
Susan Athey and Stefan Wager. Estimating treatment effects with causal forests: An application. Observational Studies, 5(2):37-51, 2019.
Susan Athey, Julie Tibshirani, and Stefan Wager. Generalized random forests. The Annals of Statistics, 47(2):1148 17788-117 1178, 2019.
Peter L. Bartlett, Nick Harvey, Christopher Liaw, and Abbas Mehrabian. Nearly-tight VC-dimension and pscudodimension bounds for piecewise linear neural networks. Journal of Machine Learning Research, 20(63):1-17, 2019.
Chris M. Bishop. Training with noise is equivalent to Tikhonov regularization. Neural Computation, 7(1):108-116, 1995.
Alexander Camuto, Matthew Willetts, Umut Simsekli, Stephen J. Roberts, and Chris C. Holmes. Explicit regularisation Gaussian noise injections. Advances in Neural Information Processing Systems, 2020.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Dufto, Christian Hansen, and Whitney Newey. Double/debiaseeNNeymn machine learning of treatment effects. American Economic Review, 107(5):261-265 2017.
Richard K. Crump, V. Joseph Hotz, Guido Imbens, and Oscar Mitnik. Moving the goalposts: Addressing limited overlap in the estimation of average treatment effects by changing the estimand, 2006.
Richard K. Crump, V. Joseph Hotz, Guido W. Imbens, and Oscar A. Mitnik. Dealing with limited overlap in estimation of average treatment effects. Biometrika, 96(1):187-199 2009.
Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogeneous treatment effect estimation. Advances in Neural Information Processing Systems, 202
Alicia Curth and Mihaela van der Schaar. Nomparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics, 2021b.
Alicia Curth and Mihaela van der Schaar. In search of insights, not magic bullets: Towards de- mystification of the model selection dilemma in heterogeneous treatment effect estimation. In International Conference Machine Learning, 2023.
12

---

Published as a conference paper at ICLR 2026
Alicia Curth, Ahmed M. Alaa, and Mihaela van der Schaar. Estimating structural target functions using machine learning and influence functions. arXiv preprint arXiv:2008. 200,00008.0..1, 2020.
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating CATE? A critical look at ML benchmarking practices in treatment effect estimation. In Advances in Neural Information Processing Systems, 2021.
Adji Bousso Dieng, Rajesh Ranganath, Jaan Altosaar, and David Blei. Noisin: Unbiased reg- ularization for recurrent neural networks. In International Conference Machine Learning, 2018.
Vincent Dorie, Jennifer Hill, Uri Shalit, Mare Scott, and Dan Cervone. Automated versus do-it- yourself methods for causal inference: Lessons learned from a data analysis competition. Statistical Science, 34(1):43-68, 2019.
Alexander D' Amour, Peng Ding, Avi Feller, Lihua Lei, and Jasject Sekhon. Overlap in observational studies with high-dimensional covariates. Journal of Econometrics, 221(2):644 654, 2021.
Stefan Feuerriegel, Dennis Frauen, Valentyn Melnychuk, Jonas Schweisthal, Konstantin Hess, Alicia Curth, Stefan Bauer, Niki Kilbertus, Isaac S. Kohane, and Mihaela van der Schaar. Causal machine learning for predicting treatment outcomes. Nature Medicine, 30(4):958-968, 2024.
Aaron Fisher. Inverse-variance weighting for estimation of heterogeneous treatment effects. In International Conference Machine Learning, 2024.
Dylan J. Foster and Vasilis Syrgkanis. Orthogonal statistical learning. The Annals of Statistics, 51(3): 879-908, 2023.
Dennis Frauen, Konstantin Hess, and Stefan Feuerriegel. Model-agnostic meta-learners for es- timating heterogeneous treatment effects over time. In International Conference Learning Representations, 2025.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2019.
Konstantin Hess, Dennis Frauen, Mihaela der Schaar, and Stefan Feuerriegel. Overlap- orthogonal meta-learner for treatment effect estimation over time, In International Conference on Learning Representations, 2026.
Jennifer L. Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Geoffrey E. Hinton, Nitish Srivastava, Alex Krizhevsky, llya Sutskever, and Ruslan R Salakhutdinov: Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv:1207.0580, 2012.
Andrew Jesson, Sören Mindermann, Yarin Gal, and Uri Shalit. Quantifying ignorance in individual- level causal-effect estimates under hidden confounding. In International Conference Machine Learning, 2021.
Fredrik D. Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International Conference on Machine Learning, 2016.
Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. Journal of Machine Learning Research, 23:7489-755 2022.
Edward H. Kennedy. Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2):308-30449 2023.
Edward H. Kennedy. Semiparametric doubly robust targeted double machine learning: review. Handbook of statistical methody for precision medicine, pp. 207-236, 2024.
13

---

Published as a conference paper at ICLR 2026
Sören R. Künzel, Jasjeet S. Sekhon, Peter J. Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10):4156-4155 2019.
Yann LeCun. The MNIST database of handwritten digits. http://yann.lecur lecu. http:///w..eeee..cmmm//ddmm///// 1998.
Daniel LeJeune, Hamid Javadi, and Richard Baraniuk. The implicit regularization of ordinary least squares ensembles. In International Conference on Artificial Intelligence and Statistics, 2020.
Fan Li, Kari Lock Morgan, and Alan M. Zaslavsky. Balancing covariates via propensity score weighting. Journal of the American Statistical Association, 1136621):390-400, 2018.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Alex Luedtke. Simplifying debiased inference via automatic differentiation and probabilistic program- ming. Journal the Royal Statistical Society Series B: Statistical Methodology. 88(1):113-32, 2026.
Roland A. Matsouaka, Yi Liu, and Yunji Zhou. Overlap, matching, or entropy weights: What are we weighting for? Communications Statistics Simulation and Computation, 54(7):1-20, 2024.
Kiyotoshi Matsuoka. Noise injection into inputs in back-propagation learning, IEEE Transactions on Systems, Man, and Cybernetics, 2263):436-440, 1992.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Normalizing flows for interventional density estimation. In International Conference Machine Learning, 2023.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Bounds espresentationninuuce confounding bias for treatment effect estimation. In International Conference on Learning Repre- sentations, 2024.
Valentyn Melnychuk, Dennis Frauen, Jonas Schweisthal, and Stefan Feuerriegel. Orthogonal rep- resentation learning for estimating causal quantities. In International Conference on Artificial Intelligence and Statistics, 2026.
Poorya Mianjy and Raman Arora. On dropout and nuclear norm regularization. In International Conference on Machine Learning, 2019.
Poorya Mianjy, Raman Arora, and Rene Vidal. On the implicit bias of dropout. In International Conference on Machine Learning, 2018.
Pawel Morzywolek, Johan Decruyenaere, and Stijn Vansteelandr. On general class of orthogonal learners for the estimation of heterogeneous treatment effects. arXiv preprint arXiv:2303. 12687, 2023.
Wenlong Mou, Yuchen Zhou, Jun Gao, and Liwei Wang. Dropout training, data-dependent regular- ization, and generalization bounds. In International Conference on Machine Learning, 2018.
Son Nguyen, Duong Nguyen, Khai Nguyen, Khoat Than, Hung Bui, and Nhat Ho. Structured dropout variational inference for Bayesian neural networks. In Advances in Neural Information Processing Systems, 2021.
Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108:299-319, 2021.
Kenneth R. Niswander. The collaborative perinatal study of the National Institute of Neurological Diseases and Stroke. The Woman and Their Pregnancies, 1972.
Vern I. Paulsen and Mrinal Raghupathi. An introduction to the theory of reproducing kernel Hilbert spaces, volume 152. Cambridge University Press, 2016.
Boris T. Polyak and Anatoli B. Juditsky. Acceleration of stochastic approximation by averaging.
SIAM Journal on Control and Optimization, 30(4):838-88 1992.
14

---

Published as a conference paper at ICLR 2026
James M. Robins and Andrea Rotnitzky. Semiparametric efficiency in multivariate regression models with missing data. Journal of the American Staristical Association, 90(429): 1995.
Jonas Rothfuss, Fabio Ferreira, Simon Boehm, Simon Walther, Maxim Ulrich, Tamim Asfour, and Andreas Krause. Noise regularization for conditional density estimation. arXiv preprint 107.v:1907.08882, 2019.
Donald B. Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of Educational Psychology, 66(5);688, 1974.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: General- ization bounds and algorithms. In International Conference Machine Learning. 2017.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(1):1229-1588 2014.
Ingo Steinwart and Andreas Christmann. Support Vector Machines. Heidelberg, Germany, 2008
Springer New York, NY,
Franciszek Hugon Szafraniec. The reproducing kernel Hilbert space and its multiplication operators. Complex Analysis and Related Topics, pp. 253-263, 2000.
Julie Tibshirani, Susan Athey, Rina Friedberg, Vitor Hadad, David Hirshberg, Luke Miner, Erik Sverdrup, Stefan Wager, and Marvin Wright. grf: Generalized random forests, 2018. URL https://CRAN 3 projectt www.org//peee... http///ccaaggrrrrrrr R package version 0.10.2.
Mark J. van der Laan. Statistical inference for variable importance. The International Journal of Biostatistics, 2(1), 2006.
Aad W. van der Vaart. Asymptotic statistics, volume 3. Cambridge University Press, Cambridge, United Kingdom, 2000.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1136523):1288-12222 2018.
Stefan Wager, Sida Wang, and Percy S. Liang. Dropout training as adaptive regularization. In Advances in Neural Information Processing Systems, 2013.
Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International Conference on Machine Learning, 2020.
Haobo Zhang, Yicheng Li, Weihao Lu, and Qian Lin. On the optimality of misspecified kernel ridge regression. In International Conference on Machine Learning, 2023.
Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou. How does mixup help with robustness and generalization? ? In International Conference on Learning Representations, 2021.
15

---

Published as a conference paper at ICLR 2026
A EXTENDED RELATED WORK
In the following, we briefly discuss plug-in learners and why they are limited in comparison to (two-stage) meta-learners (and, thus, not relevant baselines for our work). Later, we discuss existing works on inductive biases for CATE estimation and how they relate to the existing strategies to
tackle low overlap. Finally, we summarize the strategies address low overlap for both plug-in and meta-learners in T
Table Existing approaches for addressing low overlap in CATE estimation through regularization.
Most relevant methods are highlighted in yellow.
Effect of over-
Approach
Underlying learner
Underlying IBs
OART
regularization
Balancing Johansson et al.
NNs (plug-in)
LOLH-IB
X (f - DiM)
2016|Shalit al. 2017 Hassan-
pour & Greiner 2019
Johans-
500 ct 2022)
"Soft appronch"
(Curth &
NNs (plug-in) in)
S-IB
X + const)
X
van der Schaar 2021a
Const. reg. Morzywalek et al.
X.U.RA. Künzel
S-IB
/ (F + ATE)
X
x
2023
ct Curth van der
Schaar 202
Const. reg. Morzywolek et al.
DR Kennedy 2023
S-1B
/ (F + ATE)
2023
Const.
reg.
Retargeting
R Nie & Wager 2021
S-1B
LOLH-IB
4 (f WATE)
/
Morzywolek et al. 2023
IVW Fisher 2014
OAR (Our paper)
DR (Keuneyy 2023
LOLH-IB
(P - ATE)
(5)
OAR (Our paper) Retargeting
Nie & Wager
S-IB
LOLH-IB
(f WATE)
(V)
IVW Fisher 2024
Legend: model-agnostic (MA), Neymun-othogonal (NO), adapiiee regularization term (OART).
Plug-in learners. Plug-in learners aim at the conditional expected outcomes and yield the estimated
CATE the difference of the former. They can be either fully model-agnostic (e. g., S-/T-learner) model-specific (e. g., causal forest). Specific instantiations of plug-in learners include random forest
methods (Wager & Athey 2018 Tibshirani al., 2018 Athhe al., 2019 Athey & Wager 2019).
non-parametric kernel methods (Alaa & van der Schaar
2017 [ 8), and NN-based representation
learning methods (Johansson et al. 20 et al.] 2017: Hassanpour & Greiner 2019 Curth &
van der Schaar
2021b; Assaad et al. 2022 Johansson et al. 2022).
How plug-in learners deal with low overlap. In the low-overlap setting, plug-in estimators fail due to imprecise extrapolation wrt. counterfactual treatments (Jesson et al. 2021) (e. g., see Fig.
left). To tackle this, several regularization approaches have been proposed. For example, neural- based plug-in in learners can employ (i) balancing representations with empirical probability metrics
(Johansson et al. 2016), [hhllit et al. 2017: Johansson et al. 2022 Assaad et al. 2022 Melnychuk et 20266. Afternatively, one can use (ii) "xofi approuch" " of Curth & van der Schaar (2021a) which effectively forces the estimated conditional expected outcomes to be similar in low-overlap regions. Yet, both (i) and (ii) might have detrimental effect on the estimated CATE when too much regularization applied. For example, too much balancing leads to the estimation of a difference in means (DiM), also known as representation-induee induced confounding bias (Melnychuk al. 2024 2026); and the "soft approach" of Curth & van der Schaar (2021a) can force the estimated CATE to be constant. Therefore, we do not consider plug-in learners (and their regularization strategies) as relevant baselines.
Addressing low overlap through model class choice. An alternative the regularization approach for addressing low overlap is a choice of a model class / NN architecture. This approach was primarily studied for plug-in learners as it is tailored a specific model. For example, both estimated conditional expected outcomes can be forced to be similar, both (a) implicitly with an NN-based S-learner (= S-Net) (Curth& van der Schaar 2021b) and (b) explicitly with neural architecture design
Melnychuk et al. 2026) suggested a hypothetical way to incorporate balancing representations into a target model of meta-llearne In Sec. 4.22 we show that our instantiations of OAR are related to balancing of target models but, unlike balancing, are simpler to implement, as they do not require the evaluation of empirical distributional distances in the representation space.
16

---

Published as a conference paper at ICLR 2026
suggested in the "hard approach" Cf & van der Schaar (2021a). Furthermore. the lverlaerlap issue is implicitly addressed in random forests with
depth (Wager & Athey 2018
Tibshirani et al. 2018). However, in paper, focus on fully model-agnostic approaches to address low overlap that are based on regularizing target models and not on a model class choice
Inductive biases for CATE estimation. Regularization in ML is explicitly connected inductive biases: increased regularization prioritizes simpler models. An inductive bias (IB) can be thus defined as any (non-causal) assumption learning algorithm makes generalize beyond the training data (Abu-Mostala et al. 2012). In the context of CATE estimation, inductive biases are important due to the fundamental problem of causal inference (counterfactual outcomes are not observable, especially in low-overlap regions) and, thus, the impossibility of the exact data tiit--diiven model selection (Curth & van der Schaar 2023). the related work on CATE estimation, we outlined two main inductive biases: smoothness inductive bias (S-IB) and
inductive bias (LOLH-IB). S-IB assumes that the ground-truth CATE strictly simpler than both of the conditional expected outcomes (Curth & van der Schaar 2021a Morzywolek et al. 2023). By enforcing this inductive bias, we can improve low-overlap predictions by forcing lower CATE heterogeneity in the whole covariate space. LOLH-IB then extends S-IB further by assuming simpler models specifically in low- overlap regions (Melnychuk et al. 2026). In practice, both S-IB and LOLH-IB can be implemented in model-agnostic fashion via regularization. We summarize the connections between different regularization approaches and the underlying inductive biases in Table
'We acknowledge that this categorization is somewhat arbitrary, as some types of regularization might implicitly change the model class.
17

---

Published as a conference paper at ICLR 2026
B BACKGROUND MATERIALS
In the following, we provide background information on Neyman-orthogonality and two-stage meta-learners.
B. NETMAN-ORTTHOONNALIYY
We use the following additional notation: a S b means there exists C' 2 0 such that S C b, and
X, (rn) means X,/rn 2 O.
Definition 2 (Neyman-orthogonaal) (Chemozhukov et al. 2017 Foster & Syrgkanis 2023 Morzy
wolek et al. 2023)). A target risk C is called Neyman-orthogome if its pathwise cross derivative is zero:
D Do D(g" .1)le ." 7) = 0 for all E 9 and ij E H,
(14)
where D)F(J) = d F(f + th)l and F()((),, hg] = 8t, Mt F(f + t1h1 + +... tx (e,))le==.... =0= are pathwise derivatives (Foster & Syrgkanis 2023); g" = arg mingeg m): and 1) are the ground-tree nuisance functions.
dt
The definition of with respect to the misspecification of the nuisance functions.
informally means that target risk is first-order insensitive
B.2 TWO-STAGE META-I.EARNERS
To address the shortcomings of plug-in learners, two-stage meta-learners were proposed. These proceed in three steps as follows.
(i) First, one chooses target working model class G = (g(+) X - RR such as, for example, neural networks.
Then, (ii) the two-stage meta-learners define specific (original) target risk for g. Several possible target risks can be selected, and each option bears distinct interpretations and ramifications for population and finite-sample two-stage CATE estimation. For example, one can use regular MSE risk:
ni [m (oo)) (x(x)
+ F(g:F())).
(15)
weighte MSE risk:
C(g. n) E -xx
((x)(
(16)
where A(g; P(X)) is constant regularization term with a magnitude A. The latter (the overlap- weighted MSE risk) implements retargeting, but it only focuses on the overlapping regions of the population.
Finally, (iii) two-stage meta-learners minimize an empirical target risk (target loss) L(g, n) estimated from the observational sample and using the first-stage nuisance estimates i). When this empirical risk built from semi parametrrcalll efficient estimators, the resulting method is known as a Neyman- orthogonal learner (Robins & Rotnitzky 1995 Foster & Syrgkamis 2023).
Notably, the constant regularization term does not have detrimental effect on the CATE estimator. Notably, when too much regularization is applied, a non-regularized intercept of the target model yields ATE/WATE.
Remark (Over-regulatized leta-learners (Morzywolek et
2023 ). Consider a target model
class with a non-regularized intercept c. When À + co, the minimizer of Eq. (16) is given by WATE
g = arg min L(g. m) + c² ((x(X) (X) - (X))
(17)
E B(r(X)]
gEg
18

---

Published as a conference paper at ICLR 2026
C NON-PARAMETRIC TARGET MODELS: OAR RKHS NORM
Proposition 6 (Kernel ridge regression with an OAR-based RKHS norm). Let Vivig E Hx for every E HK+ Then, the minimizer of the target risk g = arg mill gEH 1. Fuu (g,1)] is in Hk+e and has the following form:
gx(x) (Tp, + Mx(e)) (p,k()
é' (X)
(X))], (18)
where (Tx, (x) (x)( is weighted covariance operator (Tp.K HK + Hkk, (s), = Eip N (x) X)O(Z, n)) is a weighted cross-covariance func- tional, n) = 6(2, n)) c² is centered psendo-outcome, and (Mx(rz))<>) = (((x)))(()) is bounded multiplication operator on (Mx (M) Hx
Hx)
Proof. See Appendix D.3
Proposition 6 Suuggests that under the special conditions, our OAR-based RKHS norm yields a KRR solution with varying regularization term A(v). For a finite-sample version of g", we refer to the following corollary.
Corollary 1. Consider that the assumptions (i)-(i)) of Proposition 6 hold and denote Kxx E
Raxa (x Ji.i,.... Kax E Rixn = -K(x,, j-m
jim.... R( E Ruxn
[p(a(), ((x())))==,... Ii A(v) E R=x== = () L and 2(1) E R
(6(2(), m)li=) I.... Then, a finite-sample KRR solution frou Proposition has the following,fome
((x) K
(D)
R(+)) it ( +
(19)
Also, Proposition&&showww shows that, although OAR-based RKHS-norm is generally undefined for linear kernel, it works well for more flexible, nnfnite-dimmesional kernels (c. g.. RBF and Matérn). In practice, assumption (i) can be satisfied by either assuming a sufficiently smooth v /A(v) (e. g., when the propensity score is smooth itself and bounded away from zero). by approximating VX(v) with some element I from Hx- This approximation
be done arbitrarily well with the infinite-dimensionall kernels they are dense in many smooth functional classes (e. RBF and Matérn are dense in compact class of continuously differentiable functions).
Finally, in the following corollary, we show the connection between KRR with retargeted learners (R-/IVW-learners) and OAR-based RKHS norm.
Corollary 2. solution of (i) the KRR with constant RKHS norm regularization with X = 1 for the
original risks of the retargeted learners (R-aVW-learrerr) coincides with a solution of (ü) the KRR with our DAR-based RKHS norm regularization with A(v(x)) = 1/v(x))fffteee for the original risk of the DR-learner, given the ground- truth nuisance functions Ty:
((x) = Kex (W Kxx + nl.)
W TT
T(n) +ē= Kæx (Kxx + +nA(v) T(n) +C.
(1)
(N)
(20)
where W(x) Raxn (m) (1 7 ((=()))) In and T(r) E Rax1 Le (x((0))
(z( I
Corollary 2 thus hints that our OAR-based RKHS norm is equivalent to retargeting with the constant regularization only in special (unnatural) case (i. e., when the ground-truth nuisance functions are known). That is, when R(T)) (n) are used instead of W (TT/TT(m) the equality (i) = (ii) does not hold anymore.
Based on Corollary we make another important observation for linear kernels K K(x, (2,2') = TT', namely that linear KRR can be formulated simultancously as parametric and non-parametric model. Interestingly, while Corollary still holds, the expression (ii) is, in general, not solution
the OAR-based KRR. This happens, the RKHS norm
VX(rg
is not defined for linear Hx
kernels when / A(v) is non-linear function. Consequently, for linear target models, the approach
19

---

Published as a conference paper at ICLR 2026
of retargeting cannot, in general (e. g., when the propensity score is not constant), be represented as version of our general OAR (Sec. 4.1).
Debiased OAR for RKHS norm. Unlike OAR for parametric models, debiasing OAR-basee RKHS norm is less intuitive. For example, the expected efficient influence function of the OAR- based RKHS norm cannot be expressed as RKHS norm itself. This can be seen after applying a Mercer representation theorem (Theorem 4.51 in Steinwart & Christmann (2008) implies that
((v)g
(((((()))9(X)(T (T)))) (x))
HK
E
IF
M )] EE
(x) # $|VEITE(AW)) A)]g
(21)
where the last equality does not hold as [[ x, A)] depends on A now). Therefore, we leave the debiasing of OAR-based RKHS norm for future work.
not a proper RKHS multiplier (it
Note on the squared multiplicative regularization. Our main motivation for the introduction of the squared multiplicative regularization was to counteract the weights of the R- and IVW-learness when noise regularization and dropout are used. Specifically, noise regularization and dropout their explicit form are down-scaled by p. In this way, the multiplicative regularization effectively results in a constant regularization, as [ip π (X)), Am (T())) = 1. Then, to preserve the weerap--adaptivie we introduced the squared multiplicative regularization. The nonparametric models (namely, kernel ridge regressions), on the other hand, have their regularizations in explicit form and without scaler p. Therefore, squared multiplicative regularization is not used with the RKHS norm.
Excess prediction risk. Finally, we show that, under some additional conditions, our OAR RKHS norm is guaranteed outperform the constant regularization (CR) (similarly to linear target models as described in Proposition 5
Proposition (Excess prediction risk of our OAR RKHS norm). Let MVV)g E HK for every g E Ho Then, the excess prediction risk of the DR-learner with the RKHS second-stage model and RKIS norm regularization has the following form:
1 [[TK + +1) TT (TK + +I) 'T To(Zzn),,Kl "g""""""'""""""""""""""""" +R f)) (22)
n
variance term
are (weighted) covariance operators (TK, T K HK + Hk): (TcRg)(z) = Ag(z) is a con- stant scaling operator for the CR; and (TonRg)(z) multiplication operator on (Mxu) HK - HHK for the OAR. Then, under (i) a conditional variance assumption (=conditional variance the oulcome is constant), the variance term for OAR is less than or equal to the variance term of the CR. Also, under (ii) a oow-overlap-oow-eeeeneeeeeed inductive bias (LOLH-IB), OAR does not increase the bias term too much.
where (TK = ((g(xx XX)) and (T (x))x))))x)))))))))))))))))))( K)
(Mx))g)(x)
A(v(z))g(z)) is bounded
Proof. See Appendix D.3
provide the full statement and the full proof of Proposition 7 in Appendix D
20

---

Published as a conference paper at ICLR 2026
D THEORETICAL RESULTS
D.1 GENERAL FRAMEWORK OF OAR
Proposition (Average regularization function as a distributional distance). The average amount of overlap-adaptin regularization is upper-hounded by the following f-divergences:
B[Am ((((v) & C. pmm. (P(X) [X A 0)) + Dím (F(X) P(X A- 1)) +1,
(23)
with C. 1/(4T))) and [m((1) 1/1² 1
E Atog (x(x)) = Clog + KL (P(X) P(X A == 0)) + KL (F(X) P(X A 1)),
(24)
with Clog = - (T
E[Am=(u(X))] I Cm2 Dr,, (F(X) || P(X A = 0)) v Ds (P(X) || P(X A- 1)) + (25)
fm=(t) = 1/t" -1, with Cm= 1/(16mnm and
where Taa P(A a), Dj an f-divergence Dy(Pr P2) S x)/P2(X z) dz; and KL is KL-divergence.
Proof. By the definitions of the regularization functions (Eq. 44). the following holds:
(P(X))²
B[Am (M())
E
(26)
4P(A
X) P(A X]
4T
PI A P( AA= 1)
) P(X A = 0)
PP(X)
P(X)
)
Cm
E
(27)
P(X A =
V Ja
P(X- 2)
p(x x)
Ca
P(X = A 07
+1 +1
(28)
P(X- 2) J. P(X = A-
1 x) +1
Cm / P/m (P(X) P(X A 0) +1 P/a II P(X A 1)) +1, (29)
where (*) holds due to Cauchy-Schward inequality, C. = 1/(4m)m1). and fm((1) = 1/12 1.
Analogously, it is easy to see that
[[Ama(u(X)) S Cm v Pr_.. (P(X) P(X A A) +1 P2 P(X A 1)) +1, (30)
where C.n 1/(1mmmm)), and (t) 1/t 1.
Similarly, we can show that the average logarithmic regularization function equals
E|Arog (x())] = - log(4) B| Po P(A = 0 X)] - E log ["A = 1 [)]
(31)
P(A 0 X)
PPA 1 X)
log(44om))
log
log
(32)
P(X)
P(X)
Clog + KT. (P(X) [P( A 0)) - (P(X)
P(X | A = 1)),
(33)
where Cleg = - log(11o01))
21

---

Published as a conference paper at ICLR 2026
D.2 PARAMETRIC TARGET MODELS: OAR NOISE REGULARIZATION & OAR DROPOUT
Proposition 2 (Explicit form of OAR noise regularization in linear g). For a linear model g(z) 3 + c, OAR noise regularization has the following explicit ACAR:
Loi 11) = AouR = 118 E(p(A,m(X)) Mv(X))).
(34)
where & is given by the original error term from Eq.
Proof. The implicit OAR noise regularization of linear target model has the following form:
Cose(g. (g. n)) B No (40 (x - -
(35)
B E [ (oo -c)
2
E-NNOO ((1 - - [T
(36)
E (.M (10703//AAAA,,,, ([(A,
E + 3 TTELE 5~ (ON0O./AAM0)) (o(A, 1
(37)
|181 EIp(A, (())) Varl X]].
(38)
where Var(g x = 2| x(v(x)).
Proposition (Explicit form of OAR dropout in linear g). For a linear model g(x) = =Bxx x OAR
noise regularization has the following explicit form AoAR:
Cos 1i) = E AauR + B3 diag (Ppl.a)-Acu)l)
(39)
where & given by the original error term from Eq. 2 ((v) = p(v) - p(v)), and diagl=] zeroes out all but the diagonal entries of a matrix.
Proof. The implicit OAR dropout of a linear target model has the following form:
CoE
Co (g. m) 1=0omm//ccn///444000000000000000
(40)
B b=-ompppoooop/04400002200000---0
E -ttp://wwooooo////ooooo000--00--
x)]]
(41)
E [ee-mmmt (X (X -xx]]]]]
+ (( Var (8T x|] = + E ((A, Var Σ 3,X,5 1 |x]]
(42)
de p(r(X)) (+,(((A, Σ 1- $ = )
de
-
O5
3m1
(43)
= = BT diag M()) B,
(44)
where the equality (*) holds as B [
Propusition 4 (Debiased OAR). Assume that the parametric model g(æ; B) is continuously differ- entiable wrt. v. Then, (i) debiased OAR noise regularization and (ii) debiased OAR dropout as
22

---

Published as a conference paper at ICLR 2026
follows:
L n) LinR( + { 1 = =2)=== for (+E, (45)
(x: V slmm) 2v(())(1 mx(X) o(X +1) (FANE)))))))))))))AAM.)))) X.A),
(46)
mo( 0))² 1-5 (F((((()); X,
(x; A::: Vegl: =mo()
pl((X))
(2) - po(X) x
.)) Velo((, (F((
(47)
Velgl is gradient wrt. E; and (F; are efficient influence functions of the regularization, functions. The latter are given as follows:
where CoAR are from Eq. and
6(X - 0) (A (2m))) 12
[F(mm(u())):: A)
(48)
P(X x
4w(x)²
8[X x]
((pm()(=));;
4(1 *(x)) (27(x) 1).
(49)
P(X = x)
8[X - #) (4 - xxx)) (2m(2) - 1)
(F (x()));
(50)
P(X = x)
f(2)
d(X =
xxx)) (1 (x)x))
IF [Flog (v()));X,A)
(51)
P(X x)
(1 - log( (2)))
6[X =] (A - (<)) (2n(x) - 1)
(x( = P(X =x)
(52)
8
6(X x) P(X 32v(a) (A- (A- xx)) (2 1),
IF(p (v(2)):X,
(53)
where 8 a Dirac delta function. Furthermore, by construction, Cdarn is a Neyman- risk.
Proof. We follow a standard technique for constructing Neyman Neyman-orthogonal risks (Foster & Syrgkanis
2023 by using one-step bors-corn with efficient influence functions (Kennedy 2024 Luedtke
2026
Ca(g. L(g, 17) + IF (C)
(54)
where C(g.1) 77) is an original target risk, La(g. 17) is Neyman-orthogonal risk, and (F(((g n); Z) is an efficient influence function of the original target risk.
To construct debiased (one-step bias-corrected) version of our OAR, we consider the OAR applied on top of the original target risk from Eq. (2):
COAR (g. n) = E Eg (TTT()))(()))))) (Mo()) -xxx)))(())))
Mo(x)
(55)
where & N(O, ((()))) or & (rop(p(((X))) and X = + X X o E. corre- spondingly, depending on the OAR version. Then, the efficient influence function of COAR (g. 1) is as follows
IF ()) / IE Es (c([(T()))(((( mx( g(xx)) P(X (F()=>)))((==)))))))) dz (56)
=)
Ee (M() - mo(x) o( xxxx)) LoAR(o,
Therefore, per Eq. (54), the debiased version of our OAR has following form:
CioAR(g, (g. (o + // IF Ex
(x)x)))()x)))) ) P(X =2)de] (57)
(xx
23

---

Published as a conference paper at ICLR 2026
The second term can then be found with product rule:
B
IF
(TT(( mo(X) 9 o(xx)) |( 2)
(58)
18
1F [F(((((())) Be( (x( o(xx)) ) z)
(59)
/ (T(()) TR Te (x) - mo(x) (x))- ol o(xe)) P( 2) dx)
+ E
=R
(A x xxx) w/ (m( (( - Ho(X) - a(xc))
(60)
+ R / (T())) TR The
- -
dz)
Hence, the debiased version of our OAR
CanAR (g. 71) = F. p( ((A, ((A, (O(() -xxx))))()))
-
(61)
TR J ((((()) TF (F (R( - - o(xx)) 2) 2)
where p(A, (X)) = (A- xXX) w/ (T(X)) + (T(())). Now we focus on the second term of Eq. 61 it differs depending the OAR version.
First, we consider OAR noise regularization. Let ~ N(00, 1) and consider a reparametrization trick & = A(v(X)), then
IF (B& (x(X) o(x+)) IF (E uo - o(x +E -(xx+-Xx(Xx))
(62)
2E (x) - mo(() - 1(( + (ee-((((xx))) (X) - (Fx Ho(x)]
(63)
2E. (1I(X)) - Ho(x)) - g(X + A(e(x)))) (F t - ((((M))
Given that (#) Mo(z); Z) = 3(x=z)
p(X=z) | (Y m=(x))) the debiased version of our
OAR becomes
L (g, m) = B p(A. x xx)) ( - (xX) - 9(X)))
(64)
p(A, (11(X)- (003)--ook po(x) - o(Xe)) (A-z()x)
xx(T(x))
(X (x- (x(X3))
+ 2E
p(A, x()) IE
+ E
E, (ww(x( (X) po(x) - 9(X + (F(X + (e-A((())11))) P(X = x) de
C*E(X:AE: velo))))
By completing a square, the latter target risk is equivalent in minimization to the following:
( [((((x) ( (A-x(X), (Y - -va(x))) + px(X) ma(X) -9(x)) "]
A 1))
E
Be ((A)
(Y-
v(X)
o(2.n)
E / E. x
(x P(X P(X x²) dr
(65)
where o(Z, is pscudo-outcome (Morzywolek et al. 2023 and, thus, the first term recovers
OAR applied to a Neyman-orthogonal target risk from Eq. (5)
Note that, to recover IVW-leamee = 1. This is reasonable choice when w(x) = v (as in the
Fisher 2024), we need to set w(n(X))
(X))
case of the IVW-Iearner) as E ((x(())) I = x(X)
æ 1. Also, is easy to see that,
((A,=(X))
after this modification, the IVW-learner is still Neyman-orthogonal.
24

---

Published as a conference paper at ICLR 2026
To derive C+ C+ Velgl:7), we use chain rule for IF (9 + + (F(x(X + e A(v(X)) = g' (X + (F
((x(X)))
(66)
where μF (X)((X)) (x(( is the efficient influence function of the regularization function (will be derived later).
Similarly, derive debiasing term for our OAR dropout. Let Bern(1 (1 -p(x(X))) that ç == (==((1-p((x)) p(x()))) and consider log-derivative trick:
IF [ (X) - (x) {(xx))))) IIF E.
(mx())- - -oo()) - o
plu(x))) (1-pl((x))) -1-X))((--))(
Xos
(67)
nP ( (no( ((x)))((()) 0 - plu(x)) P(E) de
oo
(68)
/ (mx(x) - (X) - 0 (1-p(x(x)))
Xos
)) LIF( (og//1 -((xx)] (()
2
(x)) (X) (X) - 9(X o E)) LF -xxx))
(M(X) (X) -9 (o (T-Ple(X)) )) TP( (I(o(I- (1-x(u(x))
Xoe
2E.
(69)
+ E.
(x( (X) -1 p(u(X))
Xoc
2
p(e(x)) +
plu(x))
1-c
(I((((())))
2
(m) mO( - 9(8 o)) (F()((X) - - mo(x))
(x(X)) PX( 0 (1-MM((T)))
Xoe
(I-Ple(x)) (F((((()))))
2E.
(70)
(1-x((())))
E2
(X) (X) (--(((X)) (le(x)) +
Xoe
e
1-e IF ((((((X11))
p(v(X))
2
(M)( XX( x( ()) (T((1(X)- - -o(x))
2
(xI( - (x( - g(x o E) (x (P((
(71)
+ Eg (Ha(X) - mx( 9(X E))² 1-E P(((
p(v()))
where LIF (p(x)()))) = (F(((((z)); X, A) is the efficient influence function of the regularization function (will be derived later). Now, we can complete the square, similarly to Eq. (65). which yields the following debiased target risk:
CRE (g, m) = T. Be p(A, (x)) ( ((((X)) ((A,+(X)) v(X)
(A- (A-m(x) (V - ma(x)) + (((x) vo(X) -(xx)) ']
m(2.0)
(72)
/ E. (e
E
where the second term is
c (X; Velgl))) = w(T(X)) (X) Ho(x) - g(X o))² 1- 1- (p((((()))
(73)
p(u(X))
-ww(x(X)) (x) (x x(X O) Veld (((((xX)))1(((((())))
Finally, we aim to derive IF (A((((z)); and IF (p(((())) X, A). For the multiplicative regu- larization function, namely:
Am (m())=)) 1/(4v(>)) and pm)((x)) = -4
(74)
25

---

Published as a conference paper at ICLR 2026
the efficient influence functions are
[T(((())
8{X 2) (A - xx))) (2T(2) 11)
TF(1((4(=)))11; 1;;,
(75)
4v(x)²
P(X = =x)
42(x)2
8{X x)
IF ((1-4v(z):X,A)
4 - {xx) (2T(x)) - 1).
(76)
P(X = x)
For the logarithmic regularization function, namely:
Atog (v(z)) = - (1)((4(x)) and Pm]((()) =
(77)
the efficient influence functions are
[F((())) 8{X 1) (A TT))) (2
IF ( (4g(4)())):XxA)
(78)
p(x)
P(X = x)
v(x)
[T(((()))
IF
X,A
(79)
- log((4((()))
(1 - log(4v(z)))²
v(x)
8{X x) (A- (x)) (1 2n (x))
(80)
P(X = x) (1 - log( (u 4u(x))))
Then, for the squared multiplicative regularization functtin, namely:
Am (v()) = 1/11v(xx)- and Pm=(v(x)) = - 16v(z))
(81)
the efficient influence functions are
IF [F(1/16(x)) (1:x,,) IT(v(x)) = 5{x - x} (A - -x(x)) (2T(x) 11
(82)
82(2)3
P(X = =z)
8v(z)3
8(X x} IP(X- = x)
- 1611-1
32 vxx) (A - (x)) (2T(F) - 1).
(83)
Notably, Dirac delta functions are later smoothed out with the integration fx. P(X = x) dr in the final formula of LooAR(g, 1)) Thus, they do not appear in the practical implementation.
Proposition 5 (Excess prediction risk of our OAR/dOAR dropout with linear second stage model). Let 9" (xx)) Br x denote the best linear predictor for the second-stage risk with oracle nui- sance functions (3" =arg ming E [[()),
3T x)²) ); and let ğ(x) =BT denote the finite-sample linear predictor based on CR/OAR/dOAR with the estimated nuisance functions (B = arg ming Le (g. o E {CR, OAR, dOAR)). Also, we consider the following reformula- tion of the DR
(Z 17) B*T X + o(Z,,) 1)) where o(Z, 11) is a non-lineat term with ElØ(Z, Z, m) X] = 0 (without the loss of generality).
Then, the excess prediction risk of the DR-learner with the linear second-stage model and dropout regularization has the following form:
ll - 9-" 1 = x x)] (E(E+T)) 22 + TTT" +R(n i)
(84)
bluzs temm
where L'CR- = Al for the CR. T'ORR- diag EX(ei] for the OAR/DOAR. Given this bias-variance decomposition, the following holds:
For the CR and our dOAR, the remainder term R(n, fi) only contains higher-order errors of the nuisance functions (thus. the CR and our dOAR are less sensitive to the nuisance functions' misspecificatinn. For example, the CR contains doubly-robuse terms 1iaa Hall 1all 1| - TI22.; our OAR contains doubly-robust terms 1| pall?, 12, || x1I., and
same-order propensity error || T li?.; and our dOAR contains both doubly-robust terms l/aa Ma 1al12.,11 ||
n1 1|,,, and higher-order propensity error || x: =1l..
26

---

Published as a conference paper at ICLR 2026
Under a conditional variance assumption (Varlò(Z.n) x d where o²
VarlY | X, A] is assumed to be constant), our OAR/DOAR reduces the variance term in comparison to the CR (given that OAR/dOAR is properly rescaled, i.e,, B(X(((X))) = X
see Appendix Ej. That is,
Ir [2 + (OAR) 20(zm)a " [ S ir [ (cn) 'D (
(85)
Under mild lov-overrap-loo-eeeeergeneeee (LOLH-IB) condition, OAR/dOAR does not increase the bias term too much. This means that the terms 3 OAR and [ only differ insignificantly. This the case, as the LOLH-IB assumes small values for $3 if some values of Xj lead to the low overlap.
Proof. Our proof follows in several steps.
1. Bias term. We start by defining an oracle regularized estimator, 30 = arg ming L (g. n). that relates to the oracle unregularized estimator 3 with shrinkage error b:
b= 8 3 = ( 13+1))113"
(86)
Then, the excess risk between 8 and 3 can then upper-bounded by the bias term:
I|g g" 122 E[(80T' X - x)
T(E+ '13" S 3 T "
(87)
Variance term. The finite-sample linear predictor with the estimated nuisance is given by the following:
( =
(88)
where a(î) is given by a finite-sample estimator of c(ñ) B(ø(Z, X) (the formula for 8 is analogous with c(1)). Then, the following holds asymptotically:
B - 30 ~ (=+ - c(n)).
(89)
Here, if 3 is based the dOAR, an additional second-order remainder has to be added R(1, i))
Also, the following holds due to the Neyman-orthooonnlliy::
- =((( R(n. n) and Corlé(n) = R(n, n)
(90)
Finally, the excess risk between 3 and 3 recovers our variance term:
1ا - 9 1 E[(3"X 8od x) tr [E( Covl -())()+))))()))+ R(n,
n
(91)
1 - tr - 212 (2tizm===+10 (E+1) + R(n, f). (92)
Now, we combine the bias and variance terms by decomposing 3 3' 3- 8 + 30 3 and formulate the final excess risk:
1 9 11 = =((
tr |Σ(Σ+T)-
+ R(n, n) ( (
(93)
4. To see why our OARAOAA improves the CR, under the conditional variance assumption, we show the following. Assuming the correlation matrix diagonal and has Var(X,) = (w...o.g.), the variance term has the following form:
da
((T)
Σ
m
(94)
(1+ 8,)2
27

---

Published as a conference paper at ICLR 2026
where my diaglP&<zm>l> /iag/En>/uls = Elo² Eloo/v(X) X3) and 8, diag [x(w)] B[A(2) X3]. If we then compare penalties with the same average strength C 8j = const), the minimal value of V(T) achieved when 8j x 1/3 (namely, a KKT solution). The latter can be achieved by choosing the regularization function A(v) as described in Definition (e. g. multiplicative, logarithmic, or squared multiplicative.
28

---

Published as a conference paper at ICLR 2026
D.3 NON-PARAMETRIC TARGET MODELS: OAR RKHS NORM
Proposition 6 (Kernel ridge regression with an OAR-based RKHS norm). Let /Nu)g E HK for every E Hk+c Then, the minimizer of the target risk g" = arg min ge f (g, (,11)] is in Hk+e and has the following form:
9" g'(x) (Tp,K + M (p,x(=)
(X)) (111x/X))
(95)
where (Tp.K))(o) = (()x is weighted covariance operator (Tp.k H. - H.k,, (s) = [ip(A, T (X)) (z, x))((, is a weighted cross-covariance func- tional, Φ = 6(Z, m) c¹ is a centered gheudo-outcome and (Mx(w)9) = M((()))((2) is bounded multiplication operator on
o (Mx(v) HK Hk).
Proof. The OAR-based KRR aims to minimize the following objective:
E (()))))))) + +|[Vacsss
g arg min
(96)
HK
which is equivalent to the minimization of the following objective:
E {()(A, x(x))²] + V(vij0 ||
9 c² + arg min gEHK
(97)
Hx
((z,1)= - c" is centered pseudo-outcome.
Under the strong overlap assumption (ii), ((() is bounded kernel multiplier and define
two self-adjoint multiplication operators (Szafraniec Paulsen & Raghupathi (016), M and Ma(e). that act from HK onto Hx:
(M ( ((((=))gg(x) and (Ma(s19) - M(v(x))a())
(98)
These two operators have the following property:
VAcejoll NK (M Vxcr 9, Vxiz g) (g M Axa
M /xes g) (g. Ma(s) mxx
(99)
where M is an adjoint operator.
Then, to find a minimizer of the centered OAR-based KRR objective in Eq. 977. we take path-wise (Gâteaux) derivative an arbitrary direction HK (Zhang et al. 2023):
Dy P.CChaRo m](h) dr (x)) s(X) ( + (g th, Ma(u th))
(100)
1=
=
a(x)) A(X] (g. Macs Waw + (h, Ma(e) g)
(101)
= 20(p(A, (X)) ((2.0) 2E (e[(((AT()) m()) a(x) M(X]
Ma(ss g)
(102)
2thh + 2/h, To.ki 2(h, Ma 9/
(103)
2(h, (T,, + Ma(1)9 - S,x)
(104)
where (T (x) = EEI(A, (Tp[A,((X))) (x)) X))(X) is a weighted covariance operator p Hx HR): S. (x) = Eip(A, n (X)) x))(Z.n) is weighted cross-covariance functional: and the equality (*) holds due to Mercer representation theorem (Theorem 4.51 in Steinwart & Christ
mann 2008 ). Namely, the Mercer representation theorem connects La dot product and the RKHS
dot product: E[g(XX h(x) = (h, T
Then, the centered OAR-based KRR objective in Eq. (97 optimized when for every h E Hk:
D CH (i", m)] =
g = (To. + Mxxe)) -'S,.
(105)
The latter then recovers the desired optimizer gº' 9 + c²
n
29

---

Published as a conference paper at ICLR 2026
Corollary 1. Consider that the assumptions (i)-(ii) of Proposition 6 hold and denote Kxx
RUxE = [x ,2 Ma.j=...... K E Rixn = s-...... R(
E R
[p(al), (=,,,.) I, A(V) E R = [)((((()) k=...... I and (()) E R
[6(2(), ,1)==,..... Then, a ffnite-samplle KRR solution from Proposition has the following
g(x) = X ( Kxx + nA(D))
R( ( +
(106)
Proof. The sample KRR solution immediately follows from Proposition 6 Specifically, we use a plug-in estimator of the weighted-covariance operator T R x plug-in estimator of the weighted cross-covariance operator Sp.x = R( o (f)) K and a plug-in estimator of the multiplication operator Ma(0) = A()). Then, the finite sample KRR solution is follows:
[((f) + -(())
((z) KrX
R(#) ( +
(107)
n2
KaX (R( (xx))xxx Kxx + mA(D)) RR(f) ( + à
(108)
Corollary 2. A solution of (i) the KRR with constant RKHS norm regularization with X = 1 for the
original risks of the retargeted learners (R-/IVW-learnee coincides with solution of (ii) the KRR with our OAR-basee RKHS norm regularization with M((()) = 1/v(x) for the original risk of the DR-learner, given the ground-truth nuisance functions Ty.
((x) = K.x (W(z) Kxx + nII,) 'W(T) T(m)
= =č= (Kxx +nA(v)) T(n) +c.
(109)
where W(x) Raxn (T(x(0) (1 - x(z(11))] o In and T(r) E Rax1 = (x - (z(x(() ]
Proof. Corollary follows from Corollary 1 and push-through identity:
(PQ 1) P P(QP +1)
(110)
where P and Q are conformable matrices, and is an identity matrix. Hence, by setting P = (
and Q Kxx, the following holds:
g(z) = KrX
Kxx
W(T) T(n) +8
(111)
(1)
= Kax - W((T)
1
Kxx Kxx- - II
T(n) +ē
(112)
n
KrX W((T)
1
(Kxx + nI, (W(=))-1) ((())
-1
T(n) +
(113)
n
n
KrX (Kxx +nA()) T(n) +c.
(114)
(1)
Proposition 7 (Excess prediction risk of our OAR RKHS norm). Let g" denote the best RKHS pre- dictor for the second-so second-stage stage risk with oracle nuisance functions (g" = arg min [ Z, 1) g(X))²) ); and let 9 denote the finite-sample RKHS predictor based on CR/OAR with the estimated muisance functions (g arg min H.
{CR, OAR)). Also, we consider the follow- ing reformulation of the DR psendo-outcome oi Z, n) 9" (X) 2(Z, 11), where ò(Z, 11) is RKHS approximation error term with (E)(
C (g. n)
X] = O. We assume c" w.l.o.g.
30

---

Published as a conference paper at ICLR 2026
Then, the excess prediction risk of the DR-learner with the RKHS second-stae stage model and RKHS norn regularization has the following form:
- 90 S tr [[TK + r) Tx (T< + +1) T0cz.mp,kk "g"",""""""" +)( (115)
bias term
variance
((x, are (weighted) covariance operators (TK, T.zzm)).K HK + Hk): (Tcng)(x) = >g(a) is a con- stant scaling operator for the CR; and (LaARg) = (M)))))(x)) {(v(z))g(z)) is bounded
where (TK9)(a) = Xg)(x) and ((iza)),x)(20
mutiplicatiion operator on (Mx(p) : HK - Hx) for the OAR.
Given this bias-variance decomposition, the following holds:
For the CR, the remainder term R(M, 9) only contains higher-order errors of the nuisance functions (thus, the CR is less sensitive the nuisance functions' misspecification)) Specifi- cally, the CR contains doubly-robust terms ||Aa Pa Pall2, æll... At the same time, our OAR contains doubly-robust terms Iiña - -Pallz, 11 т T122 and same-order propensity error 1| ----lll T1122
Under a conditional variance assumption (Varlo Z, x where o² VarlY | X, A] is assumed to be constant), our OAR/DOAR reduces the variance term in comparison to the CR (given that OAR/dOAR properly rescaled, i.e., B(M(((X)) = see Appendix Ep. That is,
-xxxx=xxxxxTx tr [TT - T'AR)) -TTk "TK(TK + --OAAA-'TTCZ,, "Tacz.n xl S tr [TK + +CC)) (TK(TK + [CRR (CR))"TTsZ.n)
(116)
Under mild dow.oveerap-loo--eeerggneeee (LOLH-IB) condition, OAR/dOAR does not increase the bias term too much. This means that the terms (g" T'A 9" Hx and (g* ,T,RTR 9" HH only differ insignificantly. This the case, as the LOLH-IB assumes a small norm for g" in the low-overlap regions.
Proof. Our proof proceeds in 4 steps, similarly to Proposition 5
Bias term. We start by defining oracle regularized estimator, 9º arg ginggeHii 2% (g, Then, the excess risk between g" and g" can be upper-bounded by the bias term:
11ge g" 11 [X(g"(X) - (x)) 119e -(( +(. 1)) - -1g" -(. m)li.
(117)
cth (g" m) - " .. (g" n) + ("" (+)
(118)
(1x) 2 Lg Nx + "O"Ig" I
(119)
where (+) holds as 96 is a minimizer of the corresponding loss, and (*x) holds as T is positive semi-definite operator.
Variance term. Similarly to Proposition 5 it can be shown that for the finite-sample KRR solution, the variance term is approximately
1 1|ا 1 [((( 96 (x)) = tr [TT +1)
(Tk -I
'T Øl ((2,)) + R(n, n)
(120)
Now, we combine the bias and variance terms by decomposing 9 - g" = gg- 90 + g" and formulate the final excess risk:
|| 9" 1 - 1 n
tr [(TK +D) -TT)- r) 'T 6(Z.)).. + "g*,Tg") + R(n, (121)
4. Under (i) the conditional-varianee assumption, the following holds approximately:
Tolz.)), T2/22 M T1/2
(122)
31

---

Published as a conference paper at ICLR 2026
Thus, the operator (TK + eigenmodes more effecttiveyy when I' = TOAR in comparison with I' = L'CR (given the same average regularization, namely oAR] = tr(['cR]). Hence, we obtain the desired inequality:
+ +1)- would suppresses the high-variance low-lverlap
Ir [[T + TOAR))'TKK' (Tx + (OAR) TTazzm)) kl S lr [(Tx + -Txx)'TTx( "Tx(TK + -cR) ' TT olZ,n),,KK R..
(123)
32

---

Published as a conference paper at ICLR 2026
E OAR IMPLEMENTATION DETAILS
E. RESCALING
In all the experiments, we performed the rescaling of our OAR so that it can be compared with the constant amount of regularization 0 (or (0):
{(v(x)) = À + B[)(((X))] (A(p(=))-E[A(())) ([A()(X))),
(124)
p(v(x)) = p + min
1-p (X)(I 1 (X)( (p(n(z)) (X)(x
p
(125)
where [0, 1] is adaptivity coefficient. Here, regularization. In our experiments, we set y = for (a) parametric target models and 7 = 0.9 for (b) non-parametrie target KRR (to bound away the RKHS norm regularization from zero). The rescaling is crucial, as now we can ensure that (v) (1) on average A, (2) varies depending overlap, and (3) lies in the admissible bounds (Az(v) and (u(() E (0, 1)).
leads full OAR, and y 0 is constant
Notably, after rescaling our OAR, we also need to adjust our debiased OAR (dOAR). Specifically, now we need to use (F(()((z)); X, A) instead of (F(((((z)) X, A) (and IIF (f(((x);); X, A) instead of (x)())())) X, A), respectively) in Eq. (TO)- TT. The influence functions of the rescaled OAR can
be found by using the chain rule (Kennedy) 2024 expressions:
Luedtke 2026) and are thus given by the following
Xx X((X)) (E)X) 1
(F({((())); X, 72
(126)
[x((v(X))]
(x)))
/ IF(A(U())): X, A)P(X =x) xr A(v(X)) E[(((X))),
(127)
if B|peTTN)]
7(1 (1-pp)
OF else,
E(((()))
(128)
IR(E)p(((X)) X, A) / TF(p(u())): X, A)P(X = 2) dr + p(u(X)) Elp(v(X))),
(129)
where IF (x)xxx)))) X, A) are (F((((()))) X, A) provided in Proposition
E.2 OTHER IMPLEMENTATION DETAILS
We implemented our OAR/dOAR in PyTorch and Pyro. It proceeds in two stages as follows (see all the details Algorithm T
Target network. J
Legend
Nadeenoo networks
Poan F
AJA AND Diaggei
FC
Fulyooerded Falkenk
Fc
Eyes
XLX, P
/odppta
Chesse coveeie Foom duilldiion
Peetecr wbiler
Stage Esymatinng nuisance lunsioons
Stage Filtin a targel network with
Qrideet
7 A1.
()) (debiased) DAR noise regulariatioon (0)) (detlased) DAR dropout
Figure 3: An overview of our OAR for neural network Our OAR/debiased OAR are used the second stage of the meta-learner to regularize the target network proportionally to the level of overlap (lower overlap leads to stronger regularization), Here, we instantiate OAR with noise injection for the middle layer of g: (i) OAR noise regularization and (ii) OAR dropout.
(a) parametric target model
33

---

Published as a conference paper at ICLR 2026
Stage 1. At the first stage, used neural networks (NNs) with fully-connected (FC) layers and exponential linear units (ELUs) as activation functions. Specifically, to estimate the propensity score, employed multi-layer perceptron (MLP) FC, consisting of L FC layers (with tunable number of hidden units in each layer). For the conditional expected outcomes, we used a TARNet (Shalit) et al., (017) FC consisting of representation sub-network FC and outcomes sub-networks FC (again, each sub-network has a tunable number of layers). We trained the propensity network FCr and the outcomes network FC with AdamW (Loshchilov Hutter 2019) and rlepochs 200 (epoccs = 20 for HC-MNIST dataset).
Stage (parametric target models). For second stage model, we used target MLP g with two sub-networks FC3 and FC, and ELU activations. Both FC& and FC, have a fixed number of hidden units, matching the hidden units of the TARNet FC, from the first stage. Again, for training, we employed AdamW (Loshchilov & Hutter 2019) with Nepochh = 200 (frepochs = 20 for HC-MNIST dataset). For all the experiments with the parametric models, we set the adaptivity coefficient y To further stabilize training of the target network, we (i) used exponential moving average (EMA) of model weights (Polyak & Juditsky, 1992) with a smoothing hyperparameter (K = 0.995); (ii trimmed too low propensity scores for both the pseudo-outcomes and our OAR/dOAR (0.05 - #XX) s 0.95); and (iii) clipped too large values of the bias-correction term with a threshold a = 1.0 (1C*) < & |091 S COAR(g f)))
Stage (non-parametric target models). We used KRR with radial basis function (RBF) kernel. An RBF bandwidth h is fixed individually for each dataset (see Table 5). For all the experiments with the non non- con-aaraettric models, we set the adaptivity coefficient y 0.9.
We use the same training data D for two stages of learning, as the (regularized) NNs belong to the Donsker class of estimators (van der Vaari 2000, Kennedy 2024).
Computational complexity. Importantly, the computation of the OAR/dOAR weights does not introduce lot of computational burden: Given that we have the estimated propensity scores, both OAR/dOAR are functions of the latter. The dOAR additionally requires the evaluation of the gradient wrt. the target model inputs. Yet, this operation also scales linearly wrt. the minibatch size. Therefore, both OAR/dOAR can be evaluated in linear time depending the minibatch size.
34

---

Published as a conference paper at ICLR 2026
Algorithm Pseudocode of our OAR/IOAR with meta-leamers
1: Input: Training dataset D; OARMOAA version E {+-E, of, H); average regularization strength À >
0/p (0, 1)
2: Stage Estimate nuisance functions (Pe, (l,
3: Fit propensity network FC,, (MLP) by minimizing BCE loss, L. Pn, (BCE(((C)))))
4: Fit outcomes network FC,, (TARNet) by minimizing an MSE loss, LMSE Pr. ((Y (X),
5: Output: Nuisance functions estimators f = (FC 0), F F). (z))
6: Stage 2: Fit a target model = arg min Lan
9 arg min DEO (innn(g,
7:
A(P(X)) ... (see Eq. 4 p(p(x)) + X(LX))(((((X)) +
8:
I [ I #(X) _ 0.95) E[A(r(X)) < (x)
Trimming indicator
9:
10:
Elp(u(X)) - Pa(((X) P(P(X))/Pa(((X))
11:
X(x(X) A 7 EAAC(X
(X((x)) E[X(((X))I) (see Eq. 124
RRsscaling
12:
p(P(X) + p ((X) min
E(p(2(X0) 1- (p(((x)) Blp(u(X)))) (sce Eq. 125
13:
if Target model == MLP then for i 0 Repochs n/byl do
b Parametric target models
14:
15:
Draw minibatch = AX (. ((XX) of size by from D T (X)
16:
8 N (0. ((((()))
<~ Droo((((((())
B Noise regularization dropout
18:
9(X) + FC- +5)
9(X) + (()
19:
Lanu(g, Pb ((A,
$12, - -xx))*
20:
if dOAR then
21:
C + P P(X) V (see Eq. 10 11 and Eq. (126). T28 LiNAR ?i) + LoxR g. i) + C 1[ S a & 10 S (o i))
22:
23:
end if
24:
Gradient & EMA update of the target network g wrt. Lana(g, /// LL (g. ii) end for ((x) < (z),
25:
27:
else Target model == KRR then
b Non-parametric target models
28:
if TAR then
29:
g(x) K.x (R(#) Kxx (A)) "R(()) (i) (see Eq.
else
31:
Undefined (see discussion in Appendix
32:
end if
33:
end
Output: CATE estimator ğ
35

---

Published as a conference paper at ICLR 2026
E.3 HYPERPARAMETER TUNING
We performed hyperparameter tuning of the first-stage models based on five-fold cross-validation using the training subset. For the second stage, we used fixed hyperparameters for all the experiments. as an exact hyperparameter search is not possible for target CATE models solely with the observational data (Curth & van der Schaar 2023). Table provides all the details on hyperparameter tuning, For reproducibility, we made tuned hyperparameters available in our GitHub
Table 5: Hyperparameter tuning for our OAR/OAA with meta-learners.
Stage
Model
Hyperparameter
Range/ Value
Learning rale
0.001, 0.005, 0.01 32, 64, 128
Minibatch size, by
Weight decay
0.0, 0.001, 0.01, 0.1
Hidden layers in FC+
L
Propensity network (MI.P
Hidden units in FC,
Rd,, 1.5 Rd 2 Rd,
Tuning stralegy
random grid search with 50 runs
Tuning criterion Optimizer
factual BCE loss
AdamW
Stage 1
Learning rate Minibatch size, by
0.001, 0.005, 0.01 32, 64, 128 Rd., 1.5 R R., 2 d. R d,, 1.5 R d, 2 de
Hidden units in FC
Dimensionality of d, dø
Outcomes network (TARNet)
Hidden units in FC
Rd 1.5 R dg. 2 R de 0.0, 0.001, 0.01, 0.1 random grid search with 50 runs factual MSF loss AdamW
Weight decay
Tuning strategy
Tuning criterion Optimizer
Learning rate Minibatch size, by
0.005
64
EMA of model weights, &
0.995
Target network (MLP)
Hidden units in Tuning strategy
Hidden units in FC,
Stage 2
tuning
Optimizer
Adam W
Adaptivity coefficient, Y
1
RBF bandwidth,
Target KRR
Adaptivity coefficient, 7
0.9
(synthetic data, IHDP dataset, ACIC 2016 datasets),
(HC-MNIST dataset)
R (synthetic data), (IHDP dataset), R 0.25 (ACIC 2016 datasets, HC-MNIST dataset)
(synthetic data), h 2 (ACIC 2016 datasets), h = (IHDP datasct)
httpps://gittuub.com//aa
997/0AR
36

---

Published as a conference paper at ICLR 2026
F DATASET DETAILS
F.1 SYNTHETIC DATA
We adapted the synthetic dataset from Melnychuk et al. (2023) where the amount of overlap can be varied, Specifically, we took the original generative mechanisms for the covariate X and the treatment but modified the data-generating process of the outcome Y:
Mixture 1) 1 (6,N 1)),
L, -UA < log (=(x)/(1 - =(x))
A =
T(z) M
N(x00113)
UA ~ Logistic(O, 1);
0, otherwise
Y~ N(3cos(3X2 2X + 0.5) 2.5 2X 0.5),
(130)
where N(x; μ4, o²) is density of normal distribution (H,, o²) with a mean μ and standard deviation σ; and a parameter E R regulates the amount overlap (b implics perfect overlap). In our synthetic experiments, set b = 2.
F.2 IHDP DATASET
The Infant Health and Development Program (IHDP) dataset (Hill 2011 Shalit et al. 2017) is a standard semi-synthetic benchmark for assessing treatment effect estimators. It comes with 100 predefined train-test splits, each containing Aerrin 672, fleett 75, and dg 25. The ground-truth CAPOs are then given by an exponential function (no (2)) and linear function (#1 (z)) Notably, the IHDP dataset has a well-known drawback of severe lack of overlap, which causes instability for
approaches that depend on propensity score re-weighting (Curth van der Schaar 2021b Curth
et al. 2021
F.3
ACIC 2016 DATASETS
The ACIC 2016 benchmark (Dorie al. 2019) builds its covariates from the extensive Collaborative Perinatal Project on developmental disorders (Niswander 1972). Its datasets differ in (i) the number of ground-truth confounders, (ii) the degree of covariate overlap, and (iii) the smoothness and the functional form of the CAPOs. In total, ACIC 2016 supplies 77 different data dene-rangg processes, each paired with 100 identically sized samples. After one-ho hot encoding categorical variables, every sample contains
4, 802 observations and dx 82 features.
F.4 HC-MNIST DATASET
The HC-MNIST benchmark was proposed as high-dimensional, semi-synthetic dataset (Jesson et al. 202T), derived from the original MNIST digit images (LeCun 1998). contains Alaain 60, 000 training images and Reest = 10, 000 test images. HC-MNIST compresses each high-resolution image into single latent coordinate. Ø, so that the potential outcomes are complex functions of both the image's & mean pixel intensity and its digit label. Treatment assignment is determined by this one-dimensional summary together with an additional latent (synthetic) confounder U, which we treat as an observed covariate. Hence, HC-MNIST is characterized by the following data-generating process:
U Bem(0.5),
X MNIST-image()))
Φ = (clip
PN, -Pe
1.4, 1441.44 1.4
Min.)
Mix. Min.
1.4 (4-(4-14)
(131)
a(0; I*) F 0iimmoid(077000000 1-y
B():T*) = o 7.75
F
+ 11-r.
pro
A Bem
a + B(o:1")
u
y N(2A 1) + (2.1 1) 22in 1)Φ) 2(20 1)(1 + 0.5Ø), 1).
where cis label the digit from the sampled image X; n. is the average intensity of the sampled image; Pe and a. the mean and standard deviation of the average intensities of the images with
37

---

Published as a conference paper at ICLR 2026
the label C; and Min. = -2 + 10 Max.- -2 + 10 (c + 1). The parameter I defines what factor influences the treatment assignment to a larger extent, ie.. the additional confounder or the one-dimensional summary. We set I* exp(1). For further details, we refer to Jesson et al. (2021).
38

---

Published as a conference paper at ICLR 2026
G ADDITIONAL EXPERIMENTAL RESULTS
G.1 SYNTHETIC DATA
We adapted fully-synthetic dataset (dg 1) from Melnychuk al. (2023) where the amount of overlap can be varied. Here, the ground truth CATE is 0, yet both conditional expected outcomes are highly non-linear. We simulated low-overlap setting with Auaiin 250 (see Fig. D.
Results. Results are shown in Fig. 4 We see that our OAR/OOAR noise regularization improves the performance of the CR + DR-Leurer. At the same time, both OAR/dOAR dropout and RKHS norm improve the performance of (almost) all the learners + CR. This hints at more flexible nature of dropout and RKHS norm, as they depend on the covariates and, thus, are more adaptive.
Noise reg. DR-learuer
Dropout
DR-learner
RKHS norm
DR-learner
1.4
1.4
1.4
iPHHE
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.6
10
10
0.25
0.30
0.75
10
10P
Noise reg. R- R-Jearner
Dropout learrer
RKHS norm leaurer
:PEHE.st
0.8
0.8
0.8
0.7
0.7
0.7
0.6
0.6
0.6
10
10
10a
10
0.25
0.30
0.75
10-2
10"
Noise reg. IVW-learner
Dropout + IVW-learner
RKHS norm + Wea-ree
2.5
2.5
2.5
2.0
2.0
2.0
1.5
1.5
1.5
1.0
1.0
1.0
10
-2
10
10"
10a
0.25
0.50
0.75
10" Average reg. strength, À
10-2
Average reg. strength,
Average reg. strength,
Approach
OAR type
w/o reg, (A/p=0) 0)
OAR JOAR
Oracle
Arn (iom
CR (A/p - const)
Nog/Dan
Figure 4: Results for synthetic experiments. Reported: rPEHE mean + se over 40 runs. Lower is better.
39

---

Published as a conference paper at ICLR 2026
G.2 HC-MNIST DATASET
In Table 6, we demonstrate the remaining 8 experiments for the HC-MNIST dataset. Therein, our OAR/dOAR improves over CR DR--R-IIV-learmers in the majority of combinations, often significantly. Furthermore, the best performance for almost every regularization value is achieved by our OAR/dOAR.
Table 6:
Results for HC-MNIST experiments for OAR/OOAR Aoog Plog))
and
OAR/dOAR /m=)) Reported: rPEHE (Ar) mean std over 30 runs.
A/5
005
0.35
0.1
0.5
11.5
Leumes
Approuch
0,028
0.007
0U000
06666
0,022
0.704
(-0105)
0,08 (-0008)
1IR
40AR)26 Boa)
0.004)
0006- 0000))
0667
D0300 0.055
0.70
0008 0032)
0.091 (.077)
0.002
00002
0.299
0.007
0713
0.00 (-0.077)
6:110
0008 (+0.006)
0331
0.08 (-0004)
0.206
(-0.000))
0.701
0.025
0.203
0.066 (-0099
0.00 (-01000)
0.71
(+0.00)
0.700 4 0.000 (-0035)
0.066 18171
0.009 +0075
0.00
/AR
0.022 (-0000)
0.015
0.87
R
=.033
0700
0.000 (.0044)
0576 DO0G
0668
0.000 0.221
0
0.024 0.006)
0,68
0.015 0.004)
GARI /
(0.01-
8,99
(-145
0666
D448 (+0.032)
0317 00022 (407003)
0,688
0.000
(0.00)
0.84 01000-
JOAR/A (uu))
0x94 = (-1U00)
R896 1: (-8.08)
0.731 à 0.341 (44.059)
0,600 1 0.811 (-48035)
0
0,013 (-UU25)
(L680 & (-00E)
/
0.242 0.05)
0.299 0.066
0,83 0.193 0.994)
1,09 0.266 (10000)
0.342 (.23))
1.97 0.2200 0.72)
IVW
ADARU /
0.236 (-0.01))
0.218 (-0021)
6.94 0.186 (-0.083)
1070 0.577 (-0019)
= 0:311 (-0.045)
1,055
GARI (8w3)
1,890 0.211 (-11023)
1,000
0.226
(-0001)
1,05% + 0451 (+0.066)
1/887 à 01255 (-0000)
0009 0.212 (-0
01000
1-9 327)
JDAR/ up (9,))
0000 0.346 (-1.083
1,07
0.2 (-0000)
1.116 8.799 (49.149)
1,066 0.251 (-0003))
(-
0880 (-240
Lower beter holl socond best undetlizeli Change over Saacine is brockcts Injpniiiani improwetsent groce, alpzilioant worseiing 0005)
G.3 ADDITIONAL BASELINES
Baselines. In the following, we provide an absolute comparison between our OAR/dOAR approach and other existing approaches that tackle low overlap (see Sec. 2). Here, we included trimming of the IPTW weights (with the threshold E {0.05, 0.1, 0.2) }). Note that OAR/JOAR also uses default amount of trimming t = 0.05 to stabilize the training. Also, we added balancing representations with different empirical probability metrics (Johansson al., 2016 Shalit al. (017)) with o (0.5, namely Wasserstein metric (WM) and maximum mean discrepancy (MMD). For a fair comparison, implemented balancing for the target models (as it was suggested by Melnychuk ct al. (0266)).
Results. Results for synthetic and the HC-MNIST datasets are shown in Tables 7 and | respectively. For the synthetic data, our OAR/dOAR with the multiplicative regularization function in combination with the DR-learner outperforms other approaches. For the HC-MNIST dataset, our OAR/dOAR with the multiplicative regularization function outperforms all the other baselines. The main reason that other baselines fall short of our approach in the HC-MNIST that (i) trimming simply discards datapoints with low overlap; and (ii) balancing representations becomes highly unstable with the high-dimensional data (as gets increasingly harder to estimate empirical probability metrics).
'Notably, we do not report the results for OAR/dQAR RKHS norm, as kernel ridge regression does not scale
well with the dimensionality of the covariates (here, d2 784 1).
40

---

Published as a conference paper at ICLR 2026
Table 7: Results for synthetic experiments for (AR//OOR(2m/// and other baselines that tackle low overlap. Reported: rPEHE mean + std over 40 runs.
DR- learner
R-learner
Trimming (f 0.05)
1.306 + 0.375
0.837 + 0.861
Trimming (1 0.1)
1.339 + 0.374
0.837 + 0.861
Trimmiin (1 0.2)
1.353 1: 0.376
0.837 1: 0.861
Balancing (a 0.5; MMD)
1.032 + 0.406
0.658 + 0.037
Balancing 015; WM)
1.222 0.370
0.658 + 0.037
Balancing (a 5.0; MMD)
0.878 1: 0.289
0.658 1: 0.037
Balancing (a 5.0; WM)
1.156 + 0.359
0.658 + 0.037
OAR Dropout (Pm: p 0.5)
0.815 + 0.203
0.713 + 0.249
4OAR Dropout (Pmn: P 0.5)
0.828 + 0.211
0.691 + 0.072
OAR Noise (A.,. A = 1.0)
0.853 + 0.211
0.689 + 0.099
dOAR Noise reg. (A. À 1.0)
0.856 1 0.207
0.737 1: 0.172
Oracle
0.652
Lower better (best in bold, socond best underlined)
Table 8: Results for HC-MNIST experiments for OAR/OAAR and other baselines that tackle low overlap. Reported: rPE mean + std over 30 runs.
DR-leamner
R-learner
Trimming (t 0.05)
0.754 + 0.040
0.731 + 0.029
Trimming (f 0.1)
0.736 0.019
0.731 + 0.029
Trimming (f 0.2)
0.714 0.010
0.731 1 0.029
Balancing (a 0.5: MMD)
0.721 = 0.027
0.780 + 0.045
Balancing (a 0.5; WM)
0.908 + 0.167
0.946 + 0.033
Balancing (a 5.0; MMD)
0.842 1 0.230
1,210 1 0.007
Balancing (a 5.0; WM)
1.171 + 0.177
1.124 + 0.031
OAR Dropout (Pmi p 0.3)
0.713 1 0.032
0.696 1 0.013
dOAR Dropout (P..ip 0.3)
0.705 1 0.031
0.687 1 0.013
OAR Noise reg. (AmiiA 0.1)
0.726 1 0.036
0.696 1 0.009
dOAR Noise reg. (Ami À 0.1)
0.712 1 0.033
0.695 1 0.010
Oracle
0.513
Lower better (best bold, second underlined)
G.4 RUNTIME
Dhoveddroovidess the duration of one the IHDP dataset. Each run includes stages: in stage 1, we fit nuisance functions; and, in stage 2, we fit all three meta-lcarners, namely, DR----IVW-eeareers (see Algorithm| DT Here, we compared the constant regularization strategy (CR) with our OAR and dOAR w/ the multiplicative regularization function, \m/Pm. (runtimes for the logarithmic and squared multiplicative regularization functions are analogous). We observe that our OAR has almost the same runtime as the constant regularization. On the other hand, our dOAR has slightly longer training times that are attributed to the calculation of the gradient Ve[g) in the debiasing term, see Eq.
Table 9: Total runtime (in minutes) for different regularization strategies. Reported: mean + std over 100 runs. Lower is better. Experiments were carried out on 2 GPUs (NVIDIA A100-PCIE- 40GB) with Intel Xeon Silver 4316 CPUs 2.30GHz.
Reg. Approach
Noise reg.
Dropout
RKHS
CR (A/p const)
1.85 + 0.03
1.88 + 0.04
1.40 + 0.21
OAR (Am//ma)
1.87 + 0.04
1.89 + 0.03
1.35 + 0.25
dOAR (Ama/ Pm))
2.96 0.09
3.86 0.08
41


Reviews:
[{"review_id": "IwruvtPlvR", "reviewer": "Reviewer_Yknh", "summary": "This paper addresses the challenge of conditional average treatment effect (CATE) estimation under low-overlap conditions. The authors propose Overlap-Adaptive Regularization (OAR), which adjusts the regularization strength according to estimated overlap weights. The paper also introduces a debiased version that preserves Neyman-orthogonality.", "questions": "1.\tHow sensitive is OAR/dOAR to the choice of the regularization function? Are there guidelines for choosing one over the others in practice?\n\n\n2.\tSince the method relies on estimated overlap weights, how robust is the approach to mis-specification or noise in the propensity score? Does the debiased version fully mitigate this sensitivity?\n\n3.\tHow does per-instance adaptive regularization scale with large sample sizes or high-dimensional covariates, especially in neural network settings?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.\tThe paper addresses a well-known and practically important issue in CATE estimation—poor performance under low overlap conditions.\n\n2.\tIntroducing overlap-adaptive regularization is a straightforward and intuitive idea that appears novel in the context of CATE estimation.\n\n3.\tThe proposed OAR framework can be applied to various existing meta-learners (DR-, R-, and IVW-learners) and accommodates both parametric and non-parametric models.\n\n4.\tThe inclusion of a debiased version that preserves Neyman-orthogonality enhances theoretical soundness and robustness to nuisance estimation errors.", "weaknesses": "1.\tThe method’s performance is sensitive to the accuracy of estimated overlap weights, which may be unstable in practice.\n\n2.\tTheoretical analysis focuses mainly on orthogonality preservation; further insights into convergence or generalization behavior under varying overlap levels could add depth.\n\n3.\tAdaptive regularization requires computing overlap-dependent weights for each sample, which may increase overhead for large-scale or high-dimensional settings.\n\n4.\tA structural issue is that Section 3 contains material that could be moved to the appendix, while several important technical details and proofs are placed in the appendix instead of the main text, which may hinder readability and comprehension.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ad4XA441VU", "reviewer": "Reviewer_B5nJ", "summary": "This paper proposes overlap-adaptive regularizers for the second stage of meta-learners, gives parametric instantiations plus and RKHS variant, adds a debiased version.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The concept is interesting, addresses a gap in overlapping issue in causal inference.\n- Definitions and the family of $\\lambda$ functions are precise.\n- Provide some solid theory for parametric part. RKHS part is useful but assumes strong overlap.", "weaknesses": "- The paper says the noise/dropout (implicit) and explicit forms are “equivalent,” but the exact equivalence is proved only for linear target classes and generally holds as a first-order approximation for NNs. This is stated too broadly.\n- Proposition 6 requires that $P(\\epsilon < \\pi(X) < 1-\\epsilon)=1$, which makes the explicit solution neat but excludes the core part the paper aims to help (for regions with very low overlap).\n- Several IFs scale like $1/\\nu(x)^k$ can blow up when it's approaching 0. Please discuss the integrability/robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "drom6ZZdAa", "reviewer": "Reviewer_uzmE", "summary": "The paper proposes and studies methods for regularizing causal meta-learners so that, in a particular sense, the strength of the regularization explicitly adapts to the degree of overlap. More precisely, the regularization is stronger in regions of the covariate space in which there is limited overlap (i.e., the propensity score is close to zero or one). Some existing regularization methods like dropout have been shown to implicitly adapt to overlap in certain settings (Wager, Wang, and Liang, 2013), however the present work appears to be the first paper in the context of causal meta-learners to propose explicitly tying the degree of regularization to the degree of overlap differentially over the covariate space.\n\nThe authors consider existing regularization methods namely: dropout, noising of the regressors, and kernel ridge regularization. They propose overlap-adaptive versions of these techniques. In order to regularize adaptively, the degree of overlap, as measured by the conditional variance of the treatment, must be estimated in a first stage. This conditional variance function constitutes a high-dimensional nuisance parameter and the authors suggest a first-order bias correction so that the corrected objective function is Neyman Orthogonal to this new nuisance parameter in addition to the other nuisance parameters.\n\nThe authors examine the performance of their methods relative to benchmarks on a range of synthetic and semi-synthetic datasets. The adaptive approach leads to significant improvements on some of these benchmarks and generally depends on a) the degree of regularization, b) which regularization method is used, and c) which meta-learning objective is regularized. The most consistent improvement appears in the context of the HC-MNIST data. At least in my reading, no one particular version of their approach (i.e., dropout vs. noising, choice of base meta-learner, debiased vs. not debiased) seems to dominate the others.", "questions": "More a suggestion than a question, but I wonder whether there may be some way to get the best of both worlds. That is to allow for some compromise between constant and adaptive regularization. For example, suppose we let $\\lambda(v):= (4v)^{-\\gamma}-\\gamma$ where $\\gamma\\geq0$ is an additional regularization parameter to be selected by the researcher. Then in the limiting case of $\\gamma= 0$ we recover constant regularization and with $\\gamma=1$ you recover $\\lambda_n(v)$.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "I found the paper to be clear and well-explained. I find the proposed methods intuitively very reasonable, and the debiasing of the adaptively regularized objective makes the problem quite non-trivial. The discussion of the interpretation of the regularization methods in linear models seemed to me to be quite helpful. The empirical analysis of the performance of the methods is quite thorough.", "weaknesses": "While the method leads to significant improvements over baselines in some cases, the results are somewhat mixed. This is likely inevitable because, as the authors note, adaptively regularizing is only likely to be effective when the CATE is smooth in around regions with limited overlap. Nonetheless, given the mixed results, empirical practitioners may find it rather difficult to choose whether or not to regularize adaptively and additionally, which adaptive method to use. I wonder if there could be some more guidance on this front?\n\nFigure 2 seems to suggest that performance is monotonically improving with the degree of regularization almost uniformly for all methods. The same is true for the results in Table 3. For the synthetic data this is not surprising given that the true CATE is constant (equal to zero) so one would ideally one would regularize as much as possible. In order to make the comparison meaningful, the authors compare methods for a given average regularization strength, but I do not see why this is the correct comparison to make: in practice practitioners would presumably choose the regularization parameters by a data driven method like cross-validation, and may get very different parameters for the different methods. Personally, I think the constant CATE in the simulated data makes it unsuitable for comparing the methods, and for the HC-MNIST data, it might be helpful to see how performance compares when the regularization parameters are chosen empirically or to compare performance under the best choice of regularization parameter for a given method (which would require a comparison for small enough regularization parameters such that performance is no longer monotonically improving.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ciw1lVbeBj", "reviewer": "Reviewer_EKxX", "summary": "This paper introduces an overlap-adaptive regularization approach for estimating Conditional Average Treatment Effects (CATE). The key idea is to adjust the regularization strength based on the degree of overlap: regions with strong overlap receive lower regularization, while regions with weak overlap are penalized more heavily to mitigate instability. Additionally, the authors propose a debiased version of the estimator to address bias arising from estimation errors in propensity scores.", "questions": "I am curious about the different performance of the noise-injection and dropout methods shown in Figure 2. These two regularization strategies appear to behave quite differently under constant regularization. Specifically, for noise injection, the conventional regularization (CR) performs relatively well, whereas for dropout, CR performs poorly.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The idea of overlap-adaptive regularization appears to be a reasonable strategy for CATE estimation. The authors introduce two practical implementation techniques—noise injection and dropout—to facilitate ease of application in real-world settings. They provide theoretical analysis under a linear model framework. Numerical experiments are conducted to illustrate the performance of the proposed estimator.", "weaknesses": "1. The theoretical analysis provided in the paper appears inadequate. The results for the linear model can be considered preliminary, as the simplicity of the linear setting limits the relevance of regularization—making it less critical in practice. It would be more valuable to establish theoretical guarantees for nonparametric models, which are more representative of real-world applications. Furthermore, the numerical experiments primarily involve neural networks and kernel-based methods. While Appendix C presents theoretical results for RKHS, these should ideally be included in the main paper to highlight their importance. Moreover, Appendix C primarily focuses on deriving the closed-form solution for the RKHS estimator, without providing any discussion of the statistical properties of the final estimator. This omission limits the theoretical contribution and leaves readers uncertain about the estimator’s theoretical improvement over the traditional CR RKHS estimator. \n\n\n2. The simulation results are based on only 10 runs, which raises concerns. Moreover, the improvement of the proposed overlap-adaptive regularization (OAR) over CR does not appear to be statistically significant when accounting for the reported standard errors. Similarly, for the synthetic data example.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)