The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.
48,012 characters
\runningauthor{Arbour, Parikh, Niknam, Stuart, Rudolph, Feller}
\twocolumn[
\aistatstitle{Regularizing Extrapolation in Causal Inference}
\aistatsauthor{David Arbour\footnotemark[1] \And Harsh Parikh\footnotemark[1] \And Bijan Niknam}
\aistatsaddress{Adobe Research \And Yale University \And Johns Hopkins University}
\aistatsauthor{Elizabeth Stuart \And Kara Rudolph \And Avi Feller}
\aistatsaddress{Johns Hopkins University \And Columbia University \And University of California, Berkeley}
]
\footnotetext[1]{co-first authors; alphabetical order}
\begin{abstract}
Many causal inference and machine learning estimators are linear smoothers, where the prediction is a weighted average of training outcomes. Whether weights are constrained to be non-negative creates a key tradeoff: non-negative weights (e.g., inverse propensity weighting, random forests) limit extrapolation but can worsen covariate imbalance, while unconstrained weights (e.g., OLS, kernel ridge regression) improve balance but increase dependence on parametric assumptions. We propose a unified framework that directly penalizes extrapolation via a soft constraint on negative weights, replacing the standard hard non-negativity restriction. We derive a worst-case error bound and introduce a novel ``bias-bias-variance'' tradeoff among distributional imbalance, model misspecification, and estimator variance; this tradeoff is especially pronounced in high dimensions with poor positivity. We develop a convex optimization procedure that regularizes this bound and outline how to use the extrapolation penalty as a sensitivity analysis for parametric assumptions. We demonstrate our approach on synthetic data and a real-world application generalizing randomized trial estimates to a target population.
\end{abstract}
\section{INTRODUCTION}
A core challenge in observational causal inference and domain adaptation is to adjust data distributions so that features are comparable across distinct groups, such as control and treated arms or source and target populations \citep{imbens2015causal, farahani2021brief}. Weighting estimators and linear smoothers, in which the prediction is a weighted average of training outcomes, are widely used for such adjustment; examples include implicit weighting estimators like ordinary least squares (OLS) and random forests and explicit weighting approaches like inverse propensity score weighting \citep{li2013propensity} and importance sampling \citep{thomas2017importance}.
An important divide among weighting estimators is whether weights are constrained to be non-negative, such as in traditional IPW, matching \citep{stuart2010}, the synthetic control method \citep{abadie2010synthetic}, and stable balancing weights \citep{zubizarreta2015stable, ben2021balancing}, as well as in the weighting component of popular doubly robust estimators like double machine learning \citep{chernozhukov2018double}.
This constraint limits extrapolation and dependence on parametric modeling assumptions, but typically at the cost of worse feature imbalance between re-weighted groups. This imbalance is especially pronounced in high-dimensional settings, when the curse of dimensionality means that positivity is less likely to hold, leading to further bias \citep{damour2021overlap}.
By contrast, linear smoothers like OLS and kernel ridge regression allow for arbitrarily negative weights \citep{robins2007comment}, which can improve feature imbalance but at the cost of greater model dependence and higher estimator variance.
Finally, augmented estimators that combine outcome modeling with explicit weighting strategies can be viewed as performing controlled extrapolation, balancing model dependence against feature imbalance. Pure weighting and pure outcome modeling thus represent the two extremes of no versus uncontrolled extrapolation.
In this paper, we leverage this geometric perspective to establish a general framework for systematically controlling extrapolation.
In particular, we propose a unified approach that directly penalizes the level of extrapolation, replacing the current practice of a hard non-negativity constraint with a soft constraint and corresponding hyperparameter.
Unlike prior research on extrapolation in machine learning that emphasizes predictions beyond observed covariate support, we conceptualize extrapolation through unit weights, a particularly natural framework for handling high-dimensional covariates \citep{ben2021augmented}.
Specifically, our contributions are:
\begin{itemize}[leftmargin=*]
\item \textit{Bias-bias-variance tradeoff.} We propose a framework quantifying a ``bias-bias-variance" tradeoff, decomposing error into bias from distributional imbalance, bias from outcome model misspecification, and estimator variance. This captures key tradeoffs encountered in common causal inference and distribution shift scenarios.
\item \textit{Error bound and constrained optimization.} We derive an error bound based on worst-case H{\"o}lder continuity deviations from linearity. We present an optimization approach to minimize this bound, explicitly controlling tradeoffs between biases. We characterize the finite-sample variance through our error bounds.
\item \textit{Sensitivity analysis framework.} We introduce a sensitivity analysis methodology integrated into our optimization framework, enabling systematic evaluation of distributional imbalance and outcome model misspecification impacts. We illustrate this using synthetic data and a practical application involving the transportation of causal estimates to a novel target population.
\end{itemize}
\subsection{Related work}
Extrapolation and generalization are core topics in causal inference and machine learning. Recent surveys by \citet{degtiar2023review} and \citet{johansson2022generalization} provide comprehensive overviews on generalizability and transportability methods.
\paragraph{Extrapolation and the synthetic control method.} Extrapolating far from the support of the data is a longstanding concern in statistics and the social sciences especially; see \citet{king2006dangers} for a seminal discussion of possible dangers of unchecked extrapolation. Methods that limit extrapolation are common; the synthetic control method \citep{abadie2010synthetic} is a particularly prominent example. \citet{doudchenko2016balancing} discuss the non-negativity constraint in this context, and explore possible regularization. Most relevant to our approach, \citet{ben2021augmented} developed the augmented synthetic control method, which combines outcome modeling with constrained weights to reduce bias while controlling extrapolation.
\paragraph{Extrapolation in machine learning.}
Within machine learning, there has been substantial recent progress on approaches for addressing extrapolation. \citet{shen2024engression} introduced engression, a framework that views extrapolation through the lens of distributional regression, enabling principled uncertainty quantification outside the training distribution. \citet{kong2024towards} developed a causal lens for understanding extrapolation, establishing theoretical connections between causal structure and extrapolation. \citet{netanyahu2023learning} proposed a transductive approach for learning to extrapolate, leveraging unlabeled test points to guide the extrapolation process. \citet{dong2022first} provided foundational analysis toward understanding the extrapolation of nonlinear models to unseen domains, establishing bounds on extrapolation error. Finally, \citet{pfister2024extrapolation} developed extrapolation-aware nonparametric statistical inference methods, with formal guarantees on validity beyond the support of training data.
Unlike this recent literature, we approach extrapolation from a weighting perspective, which offers particular advantages in high-dimensional settings. Rather than focusing on predictions outside the covariate support, we frame extrapolation in terms of the properties of unit weights, providing a natural parameterization for high-dimensional settings \citep{ben2021augmented}. This perspective allows us to directly quantify and regularize the degree of extrapolation without relying on complex directional derivatives or high-dimensional density estimation.
\paragraph{Positivity violations and shifting the target.} Our discussion is closely related to the literature on positivity violations in causal inference.
\citet{crump2006moving}, \citet{li2018balancing}, and \citet{parikh2024we} all proposed to avoid issues due to positivity violations by shifting the estimand to regions with greater overlap. By contrast, our approach directly incorporates the severity of positivity violations into the weight estimation process.
\paragraph{Weighting representations.} A growing literature highlights the connections between various causal estimators through their weighting representations \citep{chattopadhyay2023implied}. \citet{knaus2024treatment} provided a unified framework for viewing treatment effect estimators as weighted outcomes. \citet{bruns2023augmented} showed that augmented balancing weights can be interpreted as a form of linear regression. \citet{lin2022regression} examined regression-adjusted imputation estimators through their weighting properties. Our framework builds on these insights by explicitly parameterizing the degree of extrapolation through weight regularization, providing a continuum of estimators that navigate the bias-variance tradeoff.
\section{PRELIMINARIES}
\subsection{Setup and notation}
To ease exposition, we focus on the causal inference problem of estimating the missing control potential outcome for the Average Treatment Effect on the Treated (ATT). As we note below, however, these results hold for general linear estimands as well as for domain adaptation in ML \citep{johansson2022generalization}.
For each unit $i \in [n]$, we observe the tuple $(X_i, Y_i, Z_i)$, with covariates $X_i \in \mathcal{X}$, outcome $Y_i \in \mathbb{R}$, and binary treatment $Z_i \in \{0,1\}$.
Invoking SUTVA, let $Y_i(0)$ and $Y_i(1)$ denote the control and treated potential outcomes, respectively, for unit $i$.
Our estimand of interest is the ATT, $\mathbb{E}[Y_i(1) - Y_i(0) \mid Z_i = 1]$; we discuss generalizations below. Since we observe $Y(1)$ for the treated group, the key challenge is to estimate the missing control potential outcome mean, $\mathbb{E}[Y_i(0) \mid Z_i = 1]$. Finally, define the density ratio $dQ/dP(x)$, where $Q$ and $P$ denote the populations of units assigned to treatment and control, respectively.
We need the following key assumptions for nonparametric identification:
\begin{enumerate}[label = A.\arabic*., leftmargin=*]
\item (Exchangeability) $\mathbb{E}[Y(0) \mid X, Z = 1] = \mathbb{E}[Y(0) \mid X, Z = 0]$ \label{a: cond_ign}
\item (Population overlap) $dQ/dP(x) < \infty$ for all $x \in \mathcal{X}$\label{a: pos}
\end{enumerate}
In our setup, we consider situations when the population overlap assumption \ref{a: pos} might be violated. In that case, researchers can instead
rely on parametric assumptions on $\mu({\boldsymbol{x}}) = \mathbb{E}[Y(0) \mid X_i = {\boldsymbol{x}}]$, such as linearity, to identify and estimate the expected outcomes.
Finally, following \citet{chattopadhyay2023implied}, we will focus on estimating the mean at a target \emph{covariate profile}, ${\boldsymbol{x}}^\star \in \mathcal{X}$, corresponding to our estimand of interest. For the ATT, this profile is simply the mean of the treated population, ${\boldsymbol{x}}^\star = \mathbb{E}[X \mid Z = 1]$, where $\mu({\boldsymbol{x}}^\star) = \mathbb{E}[Y(0) \mid Z = 1]$.
\paragraph{Linear in features.}
Since we are focused on linear smoothers, we consider models that are linear in \textit{some} features, but which could be complex functions of the underlying covariates. This is an extremely large model class that ranges from simple linear models to the last layer embedding from a pre-trained large language model. For our setup, we let ${\boldsymbol{x}}$ be the features in the representation implied by the parametric model, rather than simply the raw covariates.
We further assume:
\begin{assumption}
$\mu$ is H\"older continuous such that $|\mu({\boldsymbol{x}}) - \mu({\boldsymbol{x}}')| \leq a \cdot \|{\boldsymbol{x}} - {\boldsymbol{x}}'\|^\alpha$, with $a > 0$ and $\alpha > 0$.
\end{assumption}
Parameterizing $\mu$ in terms of its H\"older constants is useful for characterizing departures from linearity that directly affect the estimation error bound.
\paragraph{General linear estimands.} We can leverage recent work on the Riesz representer \citep{chernozhukov2022automatic} to immediately generalize our results to any linear functional of the data. Following \citet{bruns2023augmented}, for each unit $i \in [n]$, we observe the tuple $(X_i, Y_i, Z_i)$, with covariates $X_i \in \mathcal{X}$, outcome $Y_i \in \mathbb{R}$, and treatment $Z_i \in \mathcal{Z}$. The target functional is then
$\mathbb{E}[h(X_i, Z_i, m)]$
for a function $h \in L_2$, where $m(x,z) = \mathbb{E}[Y_i \mid X_i = x, Z_i = z]$.
Many common problems in causal inference and domain adaptation are special cases of this setup, including counterfactual quantities like the average derivative and the expected policy-specific outcome.
Finally, define a feature map $\phi : \mathcal{X} \times \mathcal{Z} \to \mathbb{R}^d$; the target \textit{feature profile} is then $\phi^*(x,z) = \mathbb{E}[ h(X, Z, \phi) ]$. Our results below apply by replacing the simple covariate profile ${\boldsymbol{x}}^\star$ with the much more general feature profile $\phi^\star(x,z)$.
\subsection{Weighting form of causal inference estimators}
\label{sec:weighting_intro}
Our focus is on \textit{weighting estimators} or \textit{linear smoothers} \citep{buja1989linear} of the form:
\begin{align*}
\hat{\mu}({\boldsymbol{x}}^\star) = \sum_{i = 1}^n w_{\star \leftarrow i}({\boldsymbol{x}}_i) Y_i,
\end{align*}
with weights $w_{\star \leftarrow i}({\boldsymbol{x}}_i)$, where $\star \leftarrow i$ emphasizes that the weights can depend both on the source covariates ${\boldsymbol{x}}_i$ and the target covariates ${\boldsymbol{x}}^\star$ \citep{lin2022regression}. When there is no ambiguity, we suppress the dependence on the covariates ${\boldsymbol{x}}_i$ and the target ${\boldsymbol{x}}^\star$.
A broad class of estimators have this form.
See \citet{knaus2024treatment} for a comprehensive discussion of the weighting form for common causal inference estimators. We highlight several special cases here, with a focus on whether the implied weights are constrained to be non-negative.
\paragraph{Explicit weighting estimators.} The first class of methods estimate the density ratio $\widehat{dQ/dP}(x)$, either directly or indirectly.
\begin{itemize}[leftmargin=*]
\item
\emph{Traditional Inverse Propensity Score Weighting.} In standard IPW \citep{rosenbaum1987model}, researchers first estimate a propensity score, $e({\boldsymbol{x}}) = \mathbb{P}[Z_i = 1 \mid {\boldsymbol{X}}_i = {\boldsymbol{x}}]$ via a binary classifier like logistic regression, and then plug into a known functional form for $dQ/dP(x)$. For the ATT, $\hat{w}({\boldsymbol{x}}) = \hat{e}({\boldsymbol{x}})/(1-\hat{e}({\boldsymbol{x}}))$; since $\hat{e}({\boldsymbol{x}}) \in (0,1)$, $\hat{w}_i({\boldsymbol{x}}) > 0$ for all $i$.
\item
\emph{Balancing weights, synthetic control, and matching.} An alternative weighting approach instead directly estimates $dQ/dP(x)$ via constrained optimization \citep{ben2021balancing}. For example, consider the minimum variance weights that control imbalance in ${\boldsymbol{x}}$ between $P$ and $Q$:
\begin{equation}\label{eq:balancing_weight_optim}
\hat{\mathbf{w}} \in \arg \min_{\mathbf{w} \in \mathcal{W}} \left\| \sum_{i=1}^n w_i {\boldsymbol{X}}_i - {\boldsymbol{x}}^\star\right\|^2_p + \lambda\|{\boldsymbol{w}}\|^2_2,
\end{equation}
where $\|\cdot\|_p$ is the $p$ vector norm and where $\mathcal{W}$ are possible constraints on the weights. \emph{Stable balancing weights} \citep{zubizarreta2015stable} and the \emph{Synthetic Control Method} \citep{abadie2010synthetic} are special cases where $\mathcal{W}$ is the simplex ($w_i \geq 0$, $\sum w_i = 1$) and the imbalance norm is $p = \infty$ and $p = 2$, respectively. \emph{Matching} is a special case where the weights are also constrained to be discrete.
\item
\emph{Riesz regression.} A final weighting approach, also known as automatic estimation of the Riesz representer \citep{chernozhukov2022debiased} also finds weights via Problem \eqref{eq:balancing_weight_optim}, albeit \textit{without} imposing the constraint that weights are non-negative. For example, minimum distance lasso Riesz regression in \citet{chernozhukov2022debiased} solves Equation \eqref{eq:balancing_weight_optim} with $\mathcal{W} = \mathbb{R}^n$ and $p = \infty$.
\end{itemize}
\paragraph{Linear smoothers and implicit weighting estimators.} A wide range of popular outcome models are linear smoothers that implicitly estimate weights $w$, including (kernel ridge) regression, $k$-nearest neighbors, random forests, xgboost, and many implementations of neural networks; see \citet{lin2022regression, curth2024random}. We highlight two prominent examples with and without a non-negativity constraint.
\begin{itemize}[leftmargin=*]
\item \emph{(Kernel) ridge regression.} For features ${\boldsymbol{X}}$, the implied ridge regression weights are:
$$
w_{\star \leftarrow i} = {{\boldsymbol{x}}^\star}^\top ({\boldsymbol{X}}^\top{\boldsymbol{X}} + \lambda \mathbb{I})^{-1} {\boldsymbol{X}}_i,
$$
where $\lambda$ is a regularization parameter; ordinary least squares (OLS) as a special case when $\lambda = 0$. Kernel ridge regression is instead based on the implied kernel features $\phi({\boldsymbol{x}})$; see \citet{bruns2023augmented}, \citet{hirshberg2019minimax}. As \citet{bruns2023augmented} discuss, the ridge regression weights are equivalent to solving optimization problem \eqref{eq:balancing_weight_optim} with the imbalance norm set to $p = 2$ and with $\mathcal{W} = \mathbb{R}^n$, which does \emph{not} include a non-negativity constraint.
\item \emph{Random forests.} As \citet{athey2019generalized} discuss in the context of causal inference, (honest) random forests is a locally adaptive linear smoother with \textit{non-negative} weights:
$$
\hat{w}_{\star \leftarrow i} = \frac{1}{B} \sum_{b=1}^{B} \frac{\mathbb{I}\{{\boldsymbol{x}}^\star \in L_b({\boldsymbol{x}})\}}{|L_b({\boldsymbol{x}})|},
$$
where $L_b$ is the set of units that share a leaf node with the target ${\boldsymbol{x}}^\star$ and $b = 1, \ldots, B$ index the trees.
\end{itemize}
\paragraph{Augmented and hybrid estimators.} Finally, augmented or hybrid estimators combine initial weights ${\boldsymbol{w}}^0$ and outcome model $\hat{m}$:
\begin{align*}
\hat\mu^{dr}({\boldsymbol{x}}^\star) &= \sum_{i=1}^N \hat{w}^0_i Y_i + \left( \hat{m}({\boldsymbol{x}}^\star) - \sum_{i=1}^N w^0_i \hat{m}({\boldsymbol{x}}_i) \right) \\&= \hat{m}({\boldsymbol{x}}^\star) + \sum_{i=1}^N \hat{w}^0_i (Y_i - \hat{m}({\boldsymbol{x}}_i)).
\end{align*}
When $\hat{m}$ is a linear smoother, then $\hat\mu^{dr}({\boldsymbol{x}})$ also has a weighting representation.
Let $\hat{m}({\boldsymbol{x}}^\star) = \sum \hat{\omega}_i({\boldsymbol{x}}) Y_i$ for a weighting function $\hat{\omega}: \mathbb{R}^d \to \mathbb{R}^n$. Following \citet{ben2021augmented}:
\begin{align*}
\hat{\mu}^{dr}({\boldsymbol{x}}^\star)
&= \sum_{i=1}^N \left(\hat{w}^0_i + \hat{w}^{\text{adj}}_i\right)Y_i \\& \mbox{ where } \;\;\hat{w}^{\text{adj}}_i \equiv \hat{\omega}_i({\boldsymbol{x}}^\star) - \sum_{j=1}^n \hat{w}^0_j \hat{\omega}_i({\boldsymbol{x}}_j).
\end{align*}
For example, when the outcome model is ridge regression, the implied weights for the doubly robust estimator have the following form:
$$
\hat{w}^{dr}_i = \hat{w}^0_i + ({\boldsymbol{x}}^\star - {\boldsymbol{x}}'\hat{{\boldsymbol{w}}}^0)' ({\boldsymbol{x}}'{\boldsymbol{x}} + \lambda \mathbb{I})^{-1} {\boldsymbol{x}}_i.
$$
Importantly, even if the initial weights ${\boldsymbol{w}}^0$ are constrained to be non-negative, such as in traditional IPW, the implied doubly robust weights ${\boldsymbol{w}}^{dr}$ could be negative. In fact, the combined weights can be negative even if both the initial weights ${\boldsymbol{w}}^0$ and the outcome model-implied weights $\hat{\omega}$ are non-negative.
There are many examples of combined estimators of this form: standard Augmented IPW \citep{chattopadhyay2023implied}, bias correction for inexact matching \citep{lin2021nn_matching}, augmented synthetic control method \citep{ben2021augmented}, and regression-adjusted imputation estimators more broadly \citep{lin2022regression}. Finally, both debiased machine learning \citep{chernozhukov2018double} and \emph{automatic} debiased machine learning \citep{chernozhukov2022automatic} have this form. The former constrains the initial weights to be non-negative; the latter does not.
\section{REGULARIZING WORST-CASE EXTRAPOLATION BIAS}
Our goal is to bound the estimation error $\left| \mu({\boldsymbol{x}}^\star) - \sum_{i=1}^n w_i Y_i \right|.$
We begin by building intuition for our approach in three steps.
\paragraph{Reflection representation.}
Under linearity, $\mu(-{\boldsymbol{x}}_i) = -\mu({\boldsymbol{x}}_i)$, so a negative weight $w_i < 0$ on ${\boldsymbol{x}}_i$ yields $w_i \mu({\boldsymbol{x}}_i) = |w_i| \mu(-{\boldsymbol{x}}_i)$; in other words, a negative weight $w_i$ is equivalent to applying a positive weight $|w_i|$ to the reflected point $-{\boldsymbol{x}}_i$. We can then construct a ``reflected'' estimator, denoted by $\ddagger$, which reflects points with negative weights around the origin:
\begin{align*}
\hat{\mu}^\ddagger ({\boldsymbol{x}}^\star) &= \sum_{i=1}^n w_i\mathbbm{1}(w_i \geq 0)
\mu({\boldsymbol{X}}_i) + |w_i|~\mathbbm{1}(w_i < 0) \mu(-{\boldsymbol{X}}_i) \\
&= \sum_{i=1}^n |w_i| \mu({{\boldsymbol{X}}}^\ddagger_i), \qquad {{\boldsymbol{X}}}^\ddagger_i = \begin{cases}
{\boldsymbol{X}}_i, & w_i \geq 0 \\
-{\boldsymbol{X}}_i, & w_i < 0
\end{cases},
\end{align*}
where $\hat\mu({\boldsymbol{x}}^\star) = \hat\mu^\ddagger({\boldsymbol{x}}^\star)$ if $\mu$ is an odd function, and where $w_i {\boldsymbol{X}}_i = |w_i| {\boldsymbol{X}}^\ddagger_i$ for all $i$.
\paragraph{Measuring parametric model violations.}
The difference between $\hat\mu({\boldsymbol{x}}^\star)$ and $\hat\mu^\ddagger({\boldsymbol{x}}^\star)$ measures the degree to which the assumed parametric model is violated. Specifically, we can write $\mu(-{\boldsymbol{x}}_i) = \delta({\boldsymbol{x}}_i) - \mu({\boldsymbol{x}}_i)$, where $\delta({\boldsymbol{x}}_i) \equiv \mu(-{\boldsymbol{x}}_i) + \mu({\boldsymbol{x}}_i)$. If $\mu$ is an odd function (e.g., $\mu$ is linear through the origin), then $\delta({\boldsymbol{x}}_i) = 0$ for all $i$. Thus, $\delta({\boldsymbol{x}}_i)$ is a point-specific measure of the extent to which the true outcome function violates the assumed parametric model; throughout, our discussion of ``nonlinearity'' should be understood as referring to such violations.
We use this representation to decompose the estimator $\hat{\mu}({\boldsymbol{x}}^\star)$:
\begin{align*}
&\hat{\mu}({\boldsymbol{x}}^\star) = \sum_{i=1}^n w_i Y_i
\;\;=\;\; \sum_{i = 1}^n \;w_i( \mu({\boldsymbol{X}}_i) + \epsilon_i) \\
&= \sum_{i=1}^n w_i\mathbbm{1}(w_i \geq 0)
\mu({\boldsymbol{X}}_i) + \\
&\qquad \quad |w_i|~\mathbbm{1}(w_i < 0)
\left(\mu(-{\boldsymbol{X}}_i) - \delta({\boldsymbol{X}}_i)\right) + w_i\epsilon_i \\
&= \underbrace{\sum_{i=1}^n |w_i| \mu({{\boldsymbol{X}}}^\ddagger_i)}_{\hat\mu^\ddagger({\boldsymbol{x}}^\star)} +
\underbrace{\sum_{i=1}^n w_i \mathbbm{1}(w_i < 0) \delta({\boldsymbol{X}}_i)}_{\text{model violation}} + \underbrace{\sum_{i=1}^n w_i \epsilon_i}_{\text{noise}}.
\end{align*}
\paragraph{Error bound.}
Although $\delta({\boldsymbol{X}})$ is unknown, we can bound it via H\"older continuity: $|\delta({\boldsymbol{x}})| = |\mu({\boldsymbol{x}}) + \mu(-{\boldsymbol{x}})| \leq |\mu({\boldsymbol{x}}) - \mu(\mathbf{0})| + |\mu(-{\boldsymbol{x}}) - \mu(\mathbf{0})| \leq 2a\|{\boldsymbol{x}}\|^\alpha$, where the last step uses $\mu(\mathbf{0}) = 0$, which holds after centering.\footnote{Replace $Y_i$ with $Y_i - \hat{\mu}(\mathbf{0})$, where $\hat{\mu}(\mathbf{0})$ is the fitted intercept.}
The resulting error bound is therefore
\begin{multline} \label{eq:error_bound}
\left|\mu({\boldsymbol{x}}^\star) - \hat{\mu}({\boldsymbol{x}}^\star)\right| \leq
\underbrace{\left|\sum_{i=1}^n |w_i|\mu({{\boldsymbol{X}}}^\ddagger_i) - \mu({\boldsymbol{x}}^\star)\right|}_{\text{error in } \hat\mu^\ddagger({\boldsymbol{x}}^\star)} \\
+ \underbrace{2a \sum_{i=1}^n |w_i|\mathbbm{1}(w_i < 0) \|{\boldsymbol{X}}_i\|^\alpha}_{\text{error due to model violation}}
+ \underbrace{\left| \sum_{i=1}^n w_i\epsilon_i \right|}_{\text{noise}} .
\end{multline}
The first term directly depends on the imbalance between the target point ${\boldsymbol{x}}^\star$ and the re-weighted (reflected) training points $|{\boldsymbol{w}}|'{\boldsymbol{X}}^\ddagger$. The second term captures additional error due to model violation, corresponding to the $\delta({\boldsymbol{X}})$ term above; this is the key new term that our framework regularizes. The third term is the noise.
\subsection{Characterizing asymmetry-induced bias}
Thus far we have presented a conservative nonparametric bound. We now provide a slightly refined characterization by noting that the extent of the bias induced by negative weights is driven by the asymmetry in $\mu$. We do so by considering the decomposition of $\mu$ into its even and odd components, i.e., $\mu(x) = \mu_e(x) + \mu_o(x)$. By the definition of odd functions, we have $-\mu_o(x) = \mu_o(-x)$; we can then bound the worst-case risk of $\hat{\boldsymbol{w}}$ using the assumed H{\"o}lder constants $a$ and $\alpha$ and isolate the effect of the even component.
The formal statement is given below in Proposition~\ref{prop: err}; the proof is given in Appendix~\ref{sec: proof_err_bnd}.
\begin{proposition}
\label{prop: err}
Let $\hat{\mu}(x^*) = \sum_{i=1}^n \hat{w}_i Y_i$ be the estimate of $\mu(x^*)$ with weights estimated via Equation~\eqref{eq: optimize} (defined below). Given $Y_i = \mu(X_i) + \epsilon_i$ where $\epsilon_i$ are independent random variables with $\mathbb{E}[\epsilon_i] = 0$ and finite second moment $\sigma^2 = \mathbb{E}[\epsilon_i^2]$, and $\mu$ is H\"older continuous with constants $a$ and $\alpha$. If $\epsilon_i$ are sub-Gaussian\footnote{We assume mean zero sub-Gaussian noise, analogous results can be obtained with this assumption replaced by bounded noise.} with parameter $\sigma$, then with probability at least $1-\delta$,
\begin{multline} \label{eq:prop_err_bound}
|\mu(x^*) - \hat{\mu}(x^*)| \leq \left|\mu(x^*) - \sum_{i=1}^n |\hat{w}_i|\,\mu({\boldsymbol{X}}_i^\ddagger)\right| \\
+ 2\sum_{i=1}^n \lvert\hat{w}_i\rvert\,\mathbf{1}(\hat{w}_i < 0)\,a\|{\boldsymbol{X}}_i\|^\alpha + \sigma \|\hat{w}\|_2 \sqrt{2\log(2/\delta)}
\end{multline}
where ${\boldsymbol{X}}_i^\ddagger = \mathrm{sign}(\hat{w}_i){\boldsymbol{X}}_i$ as defined above. The first term is the imbalance of the reflected estimator $\hat\mu^\ddagger(x^*) = \sum_i |\hat{w}_i|\mu({\boldsymbol{X}}_i^\ddagger)$; the second is the model violation bias due to negative weights (identical to the second term of \eqref{eq:error_bound}); the third is the noise.
\end{proposition}
The proof in Appendix~\ref{sec: proof_err_bnd} proceeds via an even-odd decomposition that shows that negative weights introduce additional bias only through the \emph{even} component of $\mu$; the odd part is absorbed exactly into the first term via the reflection ${\boldsymbol{X}}_i^\ddagger$. The noise term in Proposition~\ref{prop: err} requires a sub-Gaussian assumption.
Since $\mu_e$ is unidentifiable from a single dataset, we construct a conservative worst-case form that does not require access to $\mu_e$. For completeness, Proposition~\ref{prop: erremp} in Appendix~\ref{sec: proof_err_bnd} provides an empirical analog that approximates $\mu_e$ via nearest-neighbor matching when the data are approximately symmetric.
Finally, following \citet{chattopadhyay2023implied}, we define \textit{negative influence} as the fraction of total weight on units with negative weights, $\sum_i\mathbf{1}[w_i < 0] |w_i| / \sum_i |w_i|$. This is a useful summary of the extent to which the estimate relies on extrapolation.
\subsection{Proposed Estimator}
We now propose an estimator to learn weights ${\boldsymbol{w}}$ that directly control the error bound in Equation \eqref{eq:error_bound}. To do so, we modify the standard balancing weights optimization problem in Equation \eqref{eq:balancing_weight_optim} by using the Lagrangian form of the non-negativity constraint, rather than the hard constraint. Thus, the combined estimator minimizes the error bound by controlling three terms: covariate imbalance, dispersion of the weights, and level of extrapolation:
\begin{align}\label{eq: optimize}
\hat{\mathbf{w}} \in \arg \min_{\mathbf{w}} \;\;
& \underbrace{\left\| \textstyle\sum_{i=1}^n w_i {\boldsymbol{X}}_i - {\boldsymbol{x}}^\star\right\|^2_2}_{(a)\;\text{imbalance}}
\;+\; \underbrace{\lambda \|\mathbf{w}\|^2_2}_{(b)\;\text{variance}} \notag \\
&+ \;\underbrace{\gamma \textstyle\sum_{i=1}^n \mathbbm{1}\!(w_i < 0)\, \lvert w_i\rvert \, \|{\boldsymbol{X}}_i\|^\alpha}_{(c)\;\text{extrapolation}}
\end{align}
where
\begin{itemize}[leftmargin=*]
\item Term (a): Enforces balance between the target point \( {\boldsymbol{x}}^\star \) and the re-weighted training points \( \{{\boldsymbol{X}}_1, \dots, {\boldsymbol{X}}_n\} \), recalling that $w_i {\boldsymbol{X}}_i = |w_i| {\boldsymbol{X}}^\ddagger_i$ for all $i$. We focus on $p = 2$, but this generalizes to $p = \infty$. This corresponds to the first term of \eqref{eq:error_bound}.
\item Term (b): Regularizes the dispersion of the weights \( \mathbf{w} \), controlling the noise term via $\|\hat{\mathbf{w}}\|_2$ in Proposition~\ref{prop: err}.
\item Term (c): Penalizes model violation bias: the penalty $\gamma \sum_i \mathbf{1}(w_i<0)|w_i|\|{\boldsymbol{X}}_i\|^\alpha$ is proportional to the second term of \eqref{eq:error_bound}, with $\gamma$ scaling the sensitivity to parametric model violations.\footnote{In practice, $\alpha = 1$ corresponds to Lipschitz continuity; larger $\alpha$ assumes smoother departures from the parametric model and penalizes extrapolation less aggressively.}
\end{itemize}
Compared to the standard balancing weights problem \eqref{eq:balancing_weight_optim}, which trades off only imbalance and variance, the new objective \eqref{eq: optimize} introduces term (c) to control extrapolation. When the target lies outside the convex hull of the training points, achieving balance requires some weights to be negative, which increases both $\|{\boldsymbol{w}}\|_2$ and reliance on parametric assumptions. For $\gamma = 0$, Equation~\eqref{eq: optimize} recovers unconstrained balancing weights; at the other extreme, $\gamma \to \infty$ is equivalent to a hard non-negativity constraint. Increasing $\gamma$ reduces extrapolation bias and $\|{\boldsymbol{w}}\|_2$ but worsens imbalance in term (a).
\textbf{Regularizing existing estimators.} Since many causal estimators have a weighting representation (Section~\ref{sec:weighting_intro}), we can regularize extrapolation in any baseline estimator with implied weights ${\boldsymbol{w}}'$ by solving
$$
\hat{\mathbf{w}} \in \arg \min_{\mathbf{w}} \|\mathbf{w} - \mathbf{w}' \|_2^2 + \gamma \textstyle\sum_{i=1}^n \mathbbm{1}(w_i < 0)\, \lvert w_i\rvert \, \|{\boldsymbol{X}}_i\|^\alpha.
$$
For example, the augmented synthetic control method \citep{ben2021augmented} first solves with non-negative weights, then augments with a ridge outcome model that implicitly introduces negative adjustment weights; the formulation above instead directly controls the degree of negativity through $\gamma$.
\textbf{Convexity.} Despite the indicator function in term (c), the optimization problem \eqref{eq: optimize} is strongly convex and admits a unique global minimizer. To see this, note that $\mathbf{1}(w_i < 0)|w_i| = \max(0, -w_i)$, which is convex as the pointwise maximum of two affine functions. Term (a) is a squared norm of an affine function of ${\boldsymbol{w}}$, hence convex, and term (b) is strongly convex with parameter $2\lambda$. More generally, term (c) can be written using an $\ell_p$ norm over the vector of per-unit penalties. When $p = 1$ (as written above and in Proposition~\ref{prop:regularized_bound}), introducing slack variables $s_i \geq -w_i$, $s_i \geq 0$ reduces the problem to a quadratic program (QP), which can be solved exactly in polynomial time with standard solvers. When $p = 2$, the problem becomes a second-order cone program (SOCP), which is likewise solvable in polynomial time.
Finally, we can specialize the error bound for our proposed estimator:
\begin{proposition}[Regularized Bound]
\label{prop:regularized_bound}
Let $\hat{w}$ solve Equation~\eqref{eq: optimize} and $\hat{\mu}(x^*) = \sum_{i=1}^n \hat{w}_i Y_i$ where $Y_i = \mu(X_i) + \epsilon_i$. Under the assumptions of Proposition~\ref{prop: err}, with probability at least $1 - \delta$:
\begin{multline*}
|\mu(x^*) - \hat{\mu}(x^*)| \leq \left|\mu(x^*) - \sum_{i=1}^n |\hat{w}_i|\,\mu({\boldsymbol{X}}_i^\ddagger)\right| \\
+ 2\sum_{i=1}^n |\hat{w}_i|\,\mathbf{1}(\hat{w}_i < 0)\,a\|{\boldsymbol{X}}_i\|^\alpha + \sigma\|{\boldsymbol{x}}^*\|_2\sqrt{\log(2/\delta)/\lambda}
\end{multline*}
where ${\boldsymbol{X}}_i^\ddagger = \mathrm{sign}(\hat{w}_i){\boldsymbol{X}}_i$ as in Proposition~\ref{prop: err}.
\end{proposition}
The proof is provided in the appendix. The variance term $\sigma\|{\boldsymbol{x}}^*\|_2\sqrt{\log(2/\delta)/\lambda}$ does not depend on $\gamma$: increasing $\gamma$ further constrains the feasible set and cannot inflate $\|\hat{w}\|_2$. This means regularizing extrapolation reduces the bias terms without incurring additional variance cost. The full bias-bias-variance tradeoff is demonstrated empirically in Sections~\ref{sec:synth} and~\ref{sec: moud}.
\subsection{Practical guidance: $\gamma$ as a sensitivity parameter}
We argue that $\gamma$ should be treated as a sensitivity parameter rather than a tuning parameter, and encourage researchers to examine the full set of estimates it spans. We recommend the following procedure:
\begin{enumerate}[leftmargin=*]
\item Fix $\lambda$ via cross-validation on the standard balancing weights problem (i.e., with $\gamma = 0$).
\item Sweep $\gamma$ over a grid from $0$ to a value $\gamma_{\max}$ at which all weights become non-negative.
\item For each $\gamma$, record the point estimate, covariate balance (RMSE), and negative influence.
\item Examine how estimates change across this range to assess sensitivity to parametric assumptions.
\end{enumerate}
If the goal is a single point estimate, we can instead choose $\gamma^*$ in the spirit of Lepski's method, selecting the largest $\gamma$ for which the change in the point estimate remains below a researcher-defined cutoff.
\section{SYNTHETIC DATA STUDY} \label{sec:synth}
We evaluate our approach using synthetic data with both linear and nonlinear data generating processes (DGPs) where the target point lies outside the convex hull of training points (Figure~\ref{fig: convex_hull}), creating a challenging extrapolation scenario with limited sample size ($n = 10$ training units, $n/p = 5$). We also consider a high-dimensional setting ($p = 5000$, $n/p = 0.2$) using the Friedman DGP. Full descriptions and additional figures are provided in Appendix~\ref{sec: simulation}.
\begin{figure}[t]
\centering
\includegraphics[width=0.30\textwidth]{figures/sim_convex_hull.png}
\caption{Convex hull of source (training) and target units. The target point lies outside the convex hull, requiring extrapolation.}
\label{fig: convex_hull}
\end{figure}
We consider two DGPs:
\begin{align*}
\text{Linear:} \;\; \mu({\boldsymbol{X}}) &= \beta^\top {\boldsymbol{X}}, \\
\text{Nonlinear:} \;\; \mu({\boldsymbol{X}}) &= 2X_1^2 + X_2 + X_1 X_2.
\end{align*}
For the linear DGP, where the parametric assumption holds, estimation error increases monotonically as we regularize extrapolation (i.e., increase $\gamma$): relying on correct parametric assumptions yields optimal estimates. However, for the nonlinear DGP (Figure~\ref{fig:sim_nonlinear_main}), the quadratic and interaction terms violate the linearity assumption. The results illustrate the bias-bias-variance tradeoff predicted by our theory: small amounts of extrapolation remain beneficial due to the linear component, but excessive extrapolation leads to high error rates due to violations of the parametric model. The high-dimensional experiment further demonstrates that sweeping over $\gamma$ smoothly interpolates between parametric regression-like behavior (unconstrained extrapolation) and IPW-like behavior (no extrapolation).
\begin{figure}[t]
\centering
\includegraphics[width=0.75\linewidth]{figures/sim_res_nl_Error.png}
\caption{Estimation error (MSE) for the nonlinear DGP as a function of $\gamma$. The U-shaped curve illustrates the bias-bias-variance tradeoff: small $\gamma$ allows beneficial extrapolation, while large $\gamma$ incurs bias from poor balance.}
\label{fig:sim_nonlinear_main}
\end{figure}
\section{GENERALIZING MEDICATION FOR OPIOID USE DISORDER TRIAL EVIDENCE} \label{sec: moud}
We now apply our framework to the problem of generalizing causal estimates from a randomized trial to a target population. Appendix~\ref{sec: transport_assumptions} states the formal identification assumptions and maps this setting onto the general framework of Section~\ref{sec:weighting_intro}.
The Starting Treatment With Agonist Replacement Therapies (START) trial ($S_i = 1$), initiated in 2006, was a multi-center study comparing buprenorphine versus methadone in treating opioid use disorder \citep{saxon2013buprenorphine, hser2014treatment}. The trial enrolled 1,271 participants, who were randomized in a 2:1 ratio to receive either buprenorphine or methadone. Methadone was found to have higher rates of patient retention in treatment compared to buprenorphine \citep{hser2014treatment}. Our analysis focuses on the outcome of relapse to regular opioid use within 24 weeks of medication assignment, defined as non-study opioid use for four consecutive weeks or daily use for seven consecutive days.
\citet{parikh2024we} identified that Hispanic women with a pre-treatment history of amphetamine and benzodiazepine use were underrepresented in the START trial relative to the target population ($S_i = 0$), highlighting a practical violation of the positivity assumption~\ref{a: positivity_transport}. In this study, we estimate the target average treatment effect (TATE), $\tau = \mathbb{E}[Y_i(1) - Y_i(0) \mid S_i = 0]$, for this underrepresented subgroup using our proposed framework alongside standard linear regression, gradient boosting regression (GBR), inverse probability weighting (IPW), and double machine learning estimators.
The target sample is drawn from the 2015–2017 Treatment Episode Dataset - Admissions (TEDS-A), which includes data on individuals entering publicly funded substance use treatment programs across 48 states (excluding Oregon and Georgia) and the District of Columbia. Our analysis focuses on Hispanic women with a pre-treatment history of amphetamine and benzodiazepine use.
We code methadone as $Z=1$ and buprenorphine as $Z=0$, with $Y=1$ representing relapse. Pretreatment covariates include age, race, biological sex, and substance use history (amphetamine, benzodiazepines, cannabis, and intravenous drug use) measured at the initiation of medication for opioid use disorder (MOUD) treatment. For each treatment arm $z \in \{0,1\}$, we estimate $\mu_z({\boldsymbol{x}}^\star) = \mathbb{E}[Y_i \mid X_i = {\boldsymbol{x}}^\star, S_i = 1, Z_i = z]$ at the target profile ${\boldsymbol{x}}^\star = \mathbb{E}[X_i \mid S_i = 0]$ using the trial participants assigned to arm $z$ as source units. The estimated TATE is then $\hat{\tau} = \hat{\mu}_1({\boldsymbol{x}}^\star) - \hat{\mu}_0({\boldsymbol{x}}^\star)$.
\begin{figure}
\centering
\includegraphics[width=\linewidth]{figures/moud_result_TATE.png}\\
\includegraphics[width=0.6\linewidth]{figures/moud_result_method.png}
\caption{Target Average Treatment Effects for the Target Sample for Hispanic Females who have a history of Amphetamine and Benzodiazepine use in TEDS-A population. Each hue corresponds to a value of $\lambda$ and the x-axis corresponds to different values of $\gamma$ (on log scale).}
\label{fig:moud_ate}
\end{figure}
\begin{figure}
\centering
\includegraphics[width=\linewidth]{figures/balance_moud_results_0.png}\\
\includegraphics[width=\linewidth]{figures/balance_moud_results_1.png}\\
\caption{Balance between the trial and the target samples measured as the root mean squared error (RMSE) for different values of $\gamma$ and $\lambda$.}
\label{fig:moud_diagnostics_balance}
\end{figure}
\begin{figure}
\centering
\includegraphics[width=\linewidth]{figures/neg_inf_moud_results_0.png}\\
\includegraphics[width=\linewidth]{figures/neg_inf_moud_results_1.png}
\caption{Negative influence, defined as the contribution of negative weights in estimation, for different values of $\gamma$ and $\lambda$.}
\label{fig:moud_diagnostics_neg_inf}
\end{figure}
We then apply our proposed framework to this problem. By varying \(\gamma\) from 0.01 to 10, we examine how treatment effect estimates shift with increasing regularization of negative weights. Without regularization, the point estimates converge to those from linear regression. As regularization intensifies, however, the estimates smoothly shift towards zero and occasionally change sign from negative to positive for smaller values of \(\lambda\). This sensitivity underscores the influence of assumptions on the point estimates. While increasing \(\gamma\) reduces negative influence (Figure~\ref{fig:moud_diagnostics_neg_inf}), it worsens covariate balance, as reflected in higher RMSE values (Figure~\ref{fig:moud_diagnostics_balance}). Thus, our framework highlights a trade-off between minimizing reliance on parametric assumptions and achieving optimal covariate balance. Applied researchers should therefore interpret treatment effect estimates for this under-represented subgroup with caution given the sensitivity to modeling assumptions. As \citet{parikh2024we} emphasized, collecting more representative trial data is critical to credibly estimate treatment effects.
\section{CONCLUSION}
This work proposes a framework for regularizing extrapolation in causal inference by replacing hard non-negativity constraints with soft penalties on negative weights. Our theoretical error bounds show a fundamental ``bias-bias-variance'' tradeoff between distributional imbalance, model misspecification, and estimator variance, decomposing extrapolation bias through a novel reflection perspective. Empirically, synthetic data experiments confirm that controlled extrapolation smoothly interpolates between fully constrained and unconstrained approaches. A real-world medication trial illustrates how sweeping over the regularization parameter provides a practical sensitivity analysis for transportability estimates under positivity violations.
\textbf{Limitations and Future Work.} Our approach focuses on weighting-type estimators and relies on H\"{o}lder continuity and conditional ignorability, which may not hold in practice. Operationalizing sensitivity analysis for unmeasured confounding is a critical next step; existing proposals for balancing weights \citep{soriano2023interpretable} do not directly apply to our framework, and adapting such methods is an important direction. More broadly, future work should extend the bias-bias-variance tradeoff analysis to more flexible estimator classes and weaker continuity assumptions. A key open question is to characterize data-generating processes under which soft-constrained extrapolation ($\gamma > 0$) provably improves MSE relative to both unconstrained ($\gamma = 0$) and fully constrained ($\gamma \to \infty$) estimators, for example when the density ratio is large near ${\boldsymbol{x}}^\star$. Finally, our theoretical results assume sub-Gaussian noise, though analogous bounds follow under bounded outcomes via Hoeffding-type inequalities.
\section*{ACKNOWLEDGMENTS}
The authors would like to thank the reviewers, the area chair, and the program chair of AISTATS 2026 for their constructive input to help improve the paper. Harsh Parikh, Kara Rudolph, and Elizabeth Stuart would like to acknowledge that this work was funded by NIH NIDA R01DA056407.
\bibliographystyle{abbrvnat}
\bibliography{references}
\section*{CHECKLIST}
\begin{enumerate}
\item For all models and algorithms presented, check if you include:
\begin{enumerate}
\item A clear description of the mathematical setting, assumptions, algorithm, and/or model. [\textbf{Yes}]
\item An analysis of the properties and complexity (time, space, sample size) of any algorithm. [\textbf{Yes}]
\item (Optional) Anonymized source code, with specification of all dependencies, including external libraries. [\textbf{No}]
\end{enumerate}
\item For any theoretical claim, check if you include:
\begin{enumerate}
\item Statements of the full set of assumptions of all theoretical results. [\textbf{Yes}]
\item Complete proofs of all theoretical results. [\textbf{Yes}]
\item Clear explanations of any assumptions. [\textbf{Yes}]
\end{enumerate}
\item For all figures and tables that present empirical results, check if you include:
\begin{enumerate}
\item The code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL). [\textbf{Yes}]
\item All the training details (e.g., data splits, hyperparameters, how they were chosen). [\textbf{Yes}]
\item A clear definition of the specific measure or statistics and error bars (e.g., with respect to the random seed after running experiments multiple times). [\textbf{Yes}]
\item A description of the computing infrastructure used. (e.g., type of GPUs, internal cluster, or cloud provider). [\textbf{Yes}]
\end{enumerate}
\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets, check if you include:
\begin{enumerate}
\item Citations of the creator If your work uses existing assets. [\textbf{Yes}]
\item The license information of the assets, if applicable. [\textbf{Not Applicable}]
\item New assets either in the supplemental material or as a URL, if applicable. [\textbf{Not Applicable}]
\item Information about consent from data providers/curators. [\textbf{Not Applicable}]
\item Discussion of sensible content if applicable, e.g., personally identifiable information or offensive content. [\textbf{Not Applicable}]
\end{enumerate}
\item If you used crowdsourcing or conducted research with human subjects, check if you include:
\begin{enumerate}
\item The full text of instructions given to participants and screenshots. [\textbf{Not Applicable}]
\item Descriptions of potential participant risks, with links to Institutional Review Board (IRB) approvals if applicable. [\textbf{Not Applicable}]
\item The estimated hourly wage paid to participants and the total amount spent on participant compensation. [\textbf{Not Applicable}]
\end{enumerate}
\end{enumerate}
\clearpage