Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
52,198 characters · 12 sections · 31 citation commands
Residual Balancing for Non-Linear Outcome Models in High Dimensions
In observational studies, practitioners often rely on the unconfoundedness assumption, which posits that treatment assignments are as good as random given observed covariates. This assumption is at the core of many popular methods for causal inference, including regression adjustment, matching, propensity score weighting, and doubly robust methods rosenbaum1983central, imbens2015causal, hirano2003efficient. These methods, when applied in low-dimensional settings, allow for the estimation of Average Treatment Effects (ATEs) with attractive statistical properties, including asymptotic normality and efficiency.
However, modern datasets often feature a high-dimensional structure, where the number of covariates \( p \) can exceed the sample size \( n \). High-dimensional covariates arise both from richer data collection processes and deliberate expansions of low-dimensional features (e.g., interactions, polynomial expansions, fixed effects). While high-dimensional data can mitigate issues of unmeasured confounding by capturing more covariate variation, they also pose significant challenges for causal inference. In particular, classical approaches that rely on \( p \)-fixed asymptotics may no longer yield valid inferences, and new techniques are required to handle the biases introduced by high-dimensional regularization methods. This has motivated a new line of research focused on developing robust methods for causal inference in high-dimensional environments.
Recent developments in high-dimensional causal inference have adapted classical methods to this regime. One prominent approach is the Approximate Residual Balancing (ARB) framework introduced by athey2018approximate, which combines regularized regression adjustments with balancing weights. ARB demonstrates that, in high-dimensional linear models, it is possible to achieve \(\sqrt{n}\)-consistent and asymptotically normal inference for the ATE by constructing approximately balancing weights and correcting for residual confounding. This framework eliminates the need for consistent propensity score estimation and focuses directly on addressing linear biases in the outcome model.
While ARB provides strong theoretical guarantees and practical utility for linear models, many real-world applications involve nonlinear relationships between covariates and outcomes. Generalized Linear Models (GLMs) are a natural extension, allowing for nonlinear link functions \( \psi(\cdot) \) that capture more complex outcome structures. In such settings, the standard ARB framework is insufficient, as it does not account for the non-linearities introduced by the link function and its derivatives. Recognizing this limitation, athey2018approximate noted the extension to GLMs as a crucial direction for future research.
This paper directly addresses that challenge. We develop a principled extension of the ARB methodology to nonlinear outcome models. We show that a naive application of ARB fails in the nonlinear case because it ignores the curvature of the link function. Our key insight is that controlling for estimation bias requires a second-order correction. This leads to a new optimization problem for constructing balancing weights that simultaneously controls for imbalances in both the covariates and their second moments, weighted by the first and second derivatives of the link function, respectively.
Our main contributions are:
Our work builds on several streams of literature in high-dimensional statistics and causal inference.
\paragraph{High-Dimensional Causal Inference.} The primary challenge in high-dimensional causal inference is obtaining estimators that are not only consistent but also asymptotically normal, enabling the construction of valid confidence intervals. Two main paradigms have emerged to address this.
The first is the rate double robustness paradigm, where an estimator relies on both sparsity in the outcome regression and propensity score model, remaining consistent if at least one is correctly specified. The debiased or double machine learning (DML) framework is a leading example, providing a general recipe for constructing estimators insensitive to first-order errors in these nuisance components belloni2014inference, farrell2015robust, chernozhukov2017double. While methods relying on sparsity double robustness like those in bradic2019sparsitydoublerobustinference, ning2020hdcbps can accommodate a nonlinear propensity score model (e.g., a GLM for a binary treatment), their framework is fundamentally different from ours. Theirs is an approach that requires estimating both an outcome model and a propensity score model. By contrast, we purposely avoid propensity estimation and achieve orthogonality only with respect to the outcome bias via link-sensitive balancing. In nonlinear links, the Taylor remainder contains a $\psi^{\prime \prime}$ weighted quadratic term that does not vanish under first-order balance alone; consequently, we introduce a second-order guardrail (or impose a beta-min regime with support recovery) to render the remainder $o_p\left(n^{-1 / 2}\right)$. Methodologically, hdCBPS enforces balance through PS estimating equations, whereas our weights directly balance cross-group moments tied to the link's derivatives (e.g., $X^{\top} \psi^{\prime}\left(X \hat{\beta}_c\right)$ ), reflecting our outcome-centric residual-balancing perspective.
Our work is located in the second paradigm, which achieves consistency by focusing on sparsity in a single nuisance component (model double robustness) but under stronger sparsity restrictions. We provide valid inference under a sparse outcome model without requiring any specification or estimation of the propensity score model.
\paragraph{Balancing and Calibration Weights.} A parallel line of research constructs weights to directly balance covariates, without relying on a correctly specified propensity model. Entropy balancing calibrates weights to exactly match prespecified moments via a KL projection hainmueller2012entropy. Stable balancing weights minimize dispersion subject to exact balance constraints, controlling variance while enforcing mean balance zubizarreta2015sbw. Empirical balancing calibration weighting achieves semiparametric efficiency within a calibration family by solving a global moment‑calibration problem chan2016ebcw. Kernel balancing extends balance to rich function classes by matching RKHS means hazlett2020kernel. Generalized optimal matching (and KOM) unifies matching and weighting as optimization problems that trade off imbalance against variance kallus2020gom - finding weights that minimize more general measures of distributional discrepancy. Propensity‑based balancing approaches estimate the score to satisfy balance moments, e.g., CBPS imai2014cbps and its extension to continuous treatments, CBGPS fong2018cbgps. Finally, li2018balancing formalize “balancing weights” that target estimands such as ATE/ATT/overlap via propensity‑based reweighting tuned to covariate balance. Relative to these methods, our estimator is a targeted balancing scheme where the moment set is link‑sensitive—we balance cross‑group moments of $X^\top\psi'(X\hat\beta_c)$ (and a second‑order $\psi''$ term) because these are the objects that control the remainder in high‑dimensional nonlinear plug‑ins. The quadratic penalty on $\|\gamma\|_2$ plays the same role as variance/dispersion control in SBW and GOM/KOM, but our constraints and moments are specialized to the nonlinear transport of control outcomes to the treated group.
\paragraph{Inference for High-Dimensional GLMs.} Our methodology is closely related to the literature on parameter estimation and inference in high-dimensional Generalized Linear Models (GLMs). Foundational work by negahban_2012 and vandeGeer established consistency and rates of convergence for Lasso-penalized $M$-estimators, which include GLMs, under restricted strong convexity conditions. More recently, a line of work has focused on constructing confidence intervals for individual parameters in high-dimensional GLMs. Methods like the debiased Lasso JMLR:v15:javanmard14a, javanmard-montanari.debiaslasso construct approximate inverses of the Hessian to remove regularization bias. While these methods provide valid inference for coefficients, our goal is different: we target an aggregate causal parameter (ATT) and, following ARB, debias the first-stage outcome model by balancing the residual structure relevant for the estimand. This perspective connects to augmented minimax linear estimation, which chooses weights to approximate a Riesz representer subject to variance control hirshberg2021aml, but here the representer is cross‑group and depends on the link’s derivatives.
Our estimator can thus be viewed as a form of debiased nonlinear Lasso tailored to the causal estimand of interest. In hirshberg2018debiasedinferenceaveragepartial, the same plug‑in plus weighted‑residual template is used for average partial effects in a single-index model without treatment assignment, with weights obtained on a single distribution via an augmented minimax linear program. By contrast, we handle a two-sample transport problem (treated to controls), include first-order link‑sensitive cross‑group balance, and introduce an explicit second‑order $\psi''$ guardrail to control the nonlinear remainder—features that are unnecessary in their single‑distribution APE setting but essential here.
\paragraph{Higher-Order Corrections in Semiparametric Inference.} A crucial feature of our method is the use of a second-order correction in the balancing procedure. This idea has deep roots in the semiparametric efficiency literature. Robins_2008, for instance, developed higher-order influence functions to construct estimators for nonlinear functionals that are robust to estimation errors in nuisance parameters. The need for such corrections arises when the parameter of interest is not pathwise differentiable, so that first-order approximations are insufficient to remove asymptotic bias. In the machine learning literature, mackey2018orthogonalmachinelearningpower also explored the limitations of first-order orthogonality and the potential need for higher-order adjustments. Our work operationalizes this concept in the context of residual balancing, demonstrating that a second-order Taylor expansion of the bias term leads to an essential correction for achieving valid inference where simpler, first-order balancing would fail, a correction that is essential for inference but requires a stronger sparsity condition on the outcome model. \\
The remainder of the paper is organized as follows. Section 2 lays out our statistical setup and notation. Section 3 develops our main theoretical results for the known-link setting; we derive a key bias decomposition, introduce our second-order balancing procedure, and establish the asymptotic normality of our estimator. In Section 4, we extend this framework to the more flexible single-index model where the link function is unknown. Section 5 evaluates the method's performance through simulation studies and a real-world application. Finally, Section 6 concludes with a summary of our findings.
Our goal is to estimate average treatment effects (ATEs) in the potential outcomes framework. For each unit in a population, there is a pair of potential outcomes, \((Y_i(0), Y_i(1))\), where \( Y_i(1) \) represents the outcome under treatment, and \( Y_i(0) \) represents the outcome under control. Each unit is assigned to either treatment or control, with the treatment assignment denoted by \( W_i \in \{0, 1\} \). Additionally, each unit is characterized by a vector of covariates \( X_i \in \mathbb{R}^p \), where \( p \) can be large, possibly larger than the sample size \( n \).
For a random sample of size \( n \), we observe the triplet \( (X_i, W_i, Y_i^\text{obs}) \), where: \[ Y_i^obs = Y_i(W_i) =
\] is the realized outcome corresponding to the treatment assignment. Let \( n_t \) and \( n_c \) denote the number of treated and control units, respectively. The feature matrices corresponding to treated and control units are denoted by \( X_t \) and \( X_c \).
In this work, we focus on estimating the average treatment effect for the treated (ATT): \[ \tau = \frac{1}{n_t} \sum_{i : W_i = 1} \mathbb{E}[Y_i(1) - Y_i(0) \mid X_i]. \]
Throughout this paper, we adopt the notation: $\|A\|_{\infty }:=\max _{i, j}\left|A_{i j}\right|$ for entrywise $\max ;\|A\|_{\infty \rightarrow \infty}:= \max _i \sum_j\left|A_{i j}\right|$ for the matrix $\ell_{\infty}$ operator norm; and $\|v\|_1,\|v\|_2,\|v\|_{\infty}$ for vector norms. We also make the following assumptions:
We allow for nonlinear links with unbounded second and third derivatives, but the growth must be carefully controlled relative to the statistical complexity of the problem. Specifically, the second derivative's growth is restricted to be at most logarithmic in the number of covariates $p$, a mild condition ensuring the Taylor remainders in our analysis remain asymptotically negligible.
Under these assumptions, the ATT can be expressed as: \[ \tau = \mu_t - \mu_c, \quad \text{where} \quad \mu_t = \frac{1}{n_t} \sum_{i : W_i = 1} \psi(X_i^\top \beta_t), \quad \mu_c = \frac{1}{n_t} \sum_{i : W_i = 1} \psi(X_i^\top \beta_c) \] Estimating \( \mu_t \) is straightforward using the sample mean of observed treated outcomes: \[ \hat{\mu}_t = \bar{Y}_t =\frac{1}{n_t} \sum_{i : W_i = 1} Y_i^\text{obs}. \]
In contrast, estimating \( \mu_c \) is more challenging, especially in high-dimensional settings where \( p > n \). Addressing this challenge is the primary focus of this paper. Under the Approximate Residual Balancing (ARB) framework, the estimation of \( \mu_c \) is decomposed into two stages. First, a model parameterized by \( \beta_c \) is fitted to capture the strongest signals in the outcome model. This is done using regularized regression methods, such as the lasso, which are well-suited for high-dimensional settings. Second, numerical rebalancing is applied to adjust for residual confounding. Specifically, balancing weights \( \gamma_i \) are computed to align the covariate distributions between treated and control units, correcting for the residual differences in the outcomes not explained by the first-stage model.
The resulting ARB estimator for \( \mu_c \) can be expressed as: \[ \hat{\mu}_c = \frac{1}{n_t} \sum_{i : W_i = 1} \psi(X_i^\top \hat{\beta}_c) + \sum_{i : W_i = 0} \gamma_i \left(Y_i^\text{obs} - \psi(X_i^\top \hat{\beta}_c)\right), \] The first term \( \psi\left(X_i^\top \hat{\beta}_c\right) \) uses the fitted model to estimate the main effect for the control group, and the second term applies rebalancing weights \( \gamma_i \) to the residuals \( Y_i^\text{obs} - \psi(X_i^\top \hat{\beta}_c) \), ensuring that any leftover signal not captured by the main model is adjusted directly.
This two-stage process ensures that strong predictive signals are captured while residual confounding is addressed numerically, enabling robust estimation of \( \mu_c \) in high-dimensional nonlinear models. This estimator extends the ARB methodology to accommodate nonlinear link functions \( \psi \), for example in GLMs. Another interpretation for this estimator is a form of debiased nonlinear lasso JMLR:v15:javanmard14a, javanmard-montanari.debiaslasso, hirshberg2018debiasedinferenceaveragepartial.
Our proposed procedure is then: \paragraph{Step 0: Sample splitting.} We divide the dataset into two folds: one for estimating the outcome model ( $\hat{\beta}_c$ ) and the other for estimating the balancing weights ( $\gamma$ ). This splitting simplifies the theoretical analysis by ensuring independence between estimation errors.
\paragraph{Step 1: Fit the outcome model using nonlinear lasso.} Estimate \(\beta_c\) in the nonlinear model\footnote{If we have a GLM, we can instead use regularized MLE to better suit the form of exponential families.} by solving: \[ \hat{\beta}_c = \arg \min_{\beta} \left\{ \sum_{i : W_i = 0} \left(\psi(X_i^\top \beta) - Y_i^\text{obs}\right)^2 + \lambda \|\beta\|_1 \right\}, \] where \(\psi(\cdot)\) is the link function, \(Y_i^\text{obs}\) represents the observed outcomes, and \(\lambda\) is the regularization parameter.
\paragraph{Step 2: Solve the balancing problem.} Define $W_c(\hat\beta_c) = \operatorname{diag}(\psi^{\prime}(X_c \hat\beta_c))\in\mathbb{R}^{n_c\times n_c}$. Moreover, $W_t^\prime(\hat\beta_c)=\operatorname{diag}\left(\psi^{\prime \prime}\left(X_t \hat{\beta}_c\right)\right)$, and $V_i = \psi^{\prime \prime}\left(X_{c, i}^{\top} \hat{\beta_c}\right) X_{c, i} \otimes X_{c, i}$.
Compute approximately balancing weights \(\gamma\) by solving:
where $\psi^\prime(X_t\hat\beta_c)\in\mathbb{R}^n$ is interpreted as component-wise derivative. \paragraph{Step 3: Estimate the average treatment effect \(\tau\).} Compute the ATE as: \[ \hat{\tau} = \bar{Y}_t - \left\{ \frac{1}{n_t} \sum_{i:W_i=1} \psi(X_i^\top \hat{\beta}_c) + \sum_{i : W_i = 0} \gamma_i \left(Y_i^\text{obs} - \psi(X_i^\top \hat{\beta}_c)\right)\right\}, \] where \(\bar{Y}_t\) is the average observed outcome for the treated group.
\paragraph{Step 4: Cross-fitting.} Estimate the outcome model and the balancing weights in the other respective fold and aggregate the resulting ATT. \\
While we focus on the ATT, this framework can be adapted to estimate the Average Treatment Effect on the Untreated (ATU) or the overall Average Treatment Effect (ATE) by defining the target population and constructing analogous balancing weights. For instance, the latter can be estimated as \[ \hat{\tau}_{ATE} = \left\{ \frac{1}{n} \sum_{i} \psi(X_i^\top \hat{\beta}_t) + \sum_{i : W_i = 1} \gamma_{i,t} \left(Y_i^\text{obs} - \psi(X_i^\top \hat{\beta}_t)\right)\right\} - \left\{ \frac{1}{n} \sum_{i} \psi(X_i^\top \hat{\beta}_c) + \sum_{i : W_i = 0} \gamma_{i,c} \left(Y_i^\text{obs} - \psi(X_i^\top \hat{\beta}_c)\right)\right\} \] where $\gamma_{t}$ is given by the solution to
and we have $W_t(\hat\beta_t) = \operatorname{diag}(\psi^{\prime}(X_t \hat\beta_t))\in\mathbb{R}^{n_t\times n_t}$, $W^\prime(\hat\beta_t)=\operatorname{diag}\left(\psi^{\prime \prime}\left(X \hat{\beta}_t\right)\right)$, and $V_{i,t} = \psi^{\prime \prime}\left(X_{t, i}^{\top} \hat{\beta_t}\right) X_{t, i} \otimes X_{t, i}$.
In the next section, we study the theoretical properties of this procedure designed to obtain asymptotic normality for the ATT.
To analyze the performance of the ARB estimator for nonlinear models, we rely on the following bias-variance decomposition for the ARB estimate \( \hat{\mu}_c \):
This decomposition separates the bias and variance contributions to the error of the estimator \( \hat{\mu}_c \). The first term measures the effect of residual imbalances in covariate distributions, weighted by the error \( \hat{\beta}_c - \beta_c \) in estimating the outcome model. The variance term captures the contribution of the intrinsic noise \( \varepsilon_i \) to the estimator, weighted by the balancing weights \( \gamma \).
The bias-variance trade-off is achieved by carefully selecting the weights \( \gamma \) to balance the two terms in the decomposition. Moreover, we want \( \gamma \) and \( \hat{\beta}_c \) satisfying the following rate conditions:
where \( k \) is the sparsity level of \( \beta_c \).
The balancing weights \( \gamma \) are chosen to control the bias due to residual confounding by aligning the treated and control covariate distributions. Simultaneously, the regularized estimator \( \hat{\beta}_c \) for the outcome model must achieve the desired \( \ell_1 \)-rate, ensuring that the bias remains small even in high dimensions.
This decomposition and the associated conditions form the foundation for analyzing the performance of the ARB estimator and deriving theoretical guarantees. By ensuring that both conditions are met, our approach achieves robust and efficient estimation of \( \mu_c \) in high-dimensional nonlinear models.
Balancing weights \( \gamma \) are constructed to minimize bias and variance while guaranteeing a desired rate condition:
This ensures that the treated covariate distribution is well-aligned with the control covariate distribution in terms of the nonlinear weights \( \psi'(\cdot) \). Observe that the problem ((ref)) relies on an oracle $\beta_c$, which is unknown. A simple plug-in is not enough, since we would have to account for the estimation error. To make the optimization problem in ((ref)) implementable, we further linearize \( \psi'(X_i^\top \beta) \) around \( X_i^\top\hat{\beta_c} \), introducing second-order derivative terms\footnote{This further expansion to higher-order terms is in line with the strand of literature that uses higher-order corrections to get residuals that decay with faster rates. See Robins_2008, mackey2018orthogonalmachinelearningpower.} \( \psi''(\cdot) \) that will impose additional constraints on the balancing weights. Theorem (ref) makes explicit this decomposition.
Thus, we need to balance a `variance' term as well (to correct the bias from using an estimation for $\beta$ in the balancing problem) but not as aggressively as the first-order term. The balancing problem now becomes:
where the parameter $\eta$ reflects how aggressively we should balance the second-order correction term. To solve this, we propose reformulating the optimization problem for the balancing weights \( \gamma \) from its penalized form into a direct constrained optimization problem:
and note that we dropped the normalization constraints. We impose the hard constraint $\left\|\frac{1}{n_t}X_t^\top\psi^\prime(X_t\hat\beta_c)-X_c^\top W_c(\hat\beta_c)\tilde\gamma\right\|_\infty\leqslant K_1\sqrt{\frac{\log(p)}{n_c}}$ to achieve the right rate for $\hat\mu_c$ in light of the decomposition in ((ref)). Moreover, since the term in the LHS of equation ((ref)) is multiplied by $\|\hat\beta_c-\beta_c\|_1^2$ in the fundamental decomposition, it only suffices to be of constant order\footnote{This also justifies why we don't have higher-order correction terms in the balancing problem.}. In other words, a solution to this problem guarantees that
Our next result guarantees that the covariate balancing enjoys fast enough rates by proving that the optimization problem above is feasible for a wide class of designs. This is made concrete in the next assumption.
We will condition on $X_c$ when applying Hanson-Wright; the random vector $Q_i$ is independent of $X_c$ and has iid sub-Gaussian entries with $\psi_2$-norm $\varsigma$ . The scalars $\psi^{\prime}\left(X_i \beta\right)$ act as fixed weights in the quadratic forms.
This assumption ensures that the control group covariate matrix \( X_c \) follows a structured design, where the rows are independent sub-Gaussian random vectors scaled by the covariance matrix \( \Sigma_c^{1/2} \). The sub-Gaussian property controls the tail behavior of covariates, while bounds on the variances (\( \Sigma_c \)) and population means (\( m_t(\beta) \)) guarantee that covariate distributions remain well-behaved.
We assume the following standard assumptions on the rates for the nonlinear lasso.
This assumption is mild and is met in a variety of designs; we record two canonical examples.
Finally, we establish the following asymptotic normality result for the ATT. We make the following assumption on the variance structure
Theorem (ref) yields an asymptotically linear representation with influence function \(\varphi_i=\gamma_i\,\varepsilon_i\) over the control sample: \[ \hat\mu_c-\mu_c = \sum_{i:W_i=0}\gamma_i\,\varepsilon_i \;+\; o_p(s_n), \qquad s_n^2 := \operatorname{Var}\!\Big(\sum_{i:W_i=0}\gamma_i\,\varepsilon_i \,\Big|\,\mathcal F_n\Big) = \sum_{i:W_i=0}\gamma_i^2\,\sigma_i^2, \] where \(\sigma_i^2:=\operatorname{Var}(\varepsilon_i\mid X_i)\) (so \(s_n^2=\sigma^2\|\gamma\|_2^2\) under homoskedasticity).
The dispersion constraint \(\|\gamma\|_\infty\le n_c^{-2/3}\) together with \(\|\gamma\|_2^2\asymp n_c^{-1}\) ensures the Lyapunov/Lindeberg condition via \(\|\gamma\|_\infty/\|\gamma\|_2\to 0\) and rules out superefficiency by preventing a few units from dominating, which implies \(s_n^2\asymp n_c^{-1}\) when the \(\sigma_i^2\) are bounded away from \(0\) and \(\infty\). For inference under heteroskedasticity we use the robust estimator \(\widehat s_n^{\,2}=\sum_{i:W_i=0}\gamma_i^2\,\hat\varepsilon_i^2\), computed on the evaluation fold, which satisfies \(\widehat s_n^{\,2}/s_n^{2}\to_p 1\).
We do not target the semiparametric efficiency bound for the full unconfoundedness model because we do not orthogonalize with respect to the propensity score. In exchange, we avoid propensity estimation and retain transparent balance diagnostics.
Note that cross‑fitting is sufficient to guarantee conditional independence between \(\gamma\) and \(\varepsilon\); it is not necessary if one establishes an equicontinuity bound for the map \(\beta\mapsto\gamma(\beta)\). Consider a further linear approximation $$ \gamma_i(\hat{\beta_c}) \varepsilon_i(0)=\underbrace{\gamma_i\left(\beta_c\right) \varepsilon_i(0)}_{\text {main term }}+\underbrace{\left(\gamma_i(\hat{\beta_c})-\gamma_i\left(\beta_c\right)\right) \varepsilon_i(0)}_{\text {remainder }} $$ we can establish \[\sum_{i\,:\,W_i=0}\left[\gamma_i(\hat{\beta_c})-\gamma_i\left(\beta_c\right)\right] \varepsilon_i(0)=\left(\hat{\beta_c}-\beta_c\right)^{\top} \sum_{i\,:\,W_i=0} \nabla \gamma_i(\tilde{\beta_c}) \varepsilon_i(0)=o_p(1)\] , by imposing an “equicontinuity condition” on $\nabla\gamma_i(\cdot)$ to control the remainder term. The same application of Lyapunov's CLT on $\sum_{i\,:\,W_i=0}\gamma_i(\beta_c)\varepsilon_i(0)$, will yield the result given that $\gamma_i(\beta_c) \perp \varepsilon_i(0)$.
We can relax Assumption (ref), allowing for unknown link functions:
Since $\psi$ is unknown, we estimate it nonparametrically along the scalar index $Z=X^\top\beta_c$. To decouple errors, we use a three-fold scheme.
\paragraph{Step 0: Three-fold splitting.} Partition the control sample into disjoint folds $A,B,C$ :
\paragraph{Step 1: Fit the outcome model.} On $A$, fit a high-dimensional single-index estimator for $\beta_c$ (GLM penalized MLE (ref) or PV-linear lasso with whitening (ref)), yielding \[ \|\hat\beta_c-\beta_c\|_2=O_p\!\Big(\sqrt{\tfrac{k\log p}{n_A}}\Big) \]
\paragraph{Step 2: Nonparametric derivatives.} On $B$, compute $\hat Z_i=X_i^\top\hat\beta_c$ (fixed given $A$) and estimate $\widehat\psi',\widehat\psi''$ by a one-dimensional smoother of $Y$ versus $\hat Z$ (e.g. local polynomials Tsybakov2009).
Since we have assumed $\psi\in\mathcal{H}^s$, $s\ge 3$, we have the $L_2$ rates stone_1980
\paragraph{Step 3: Find balancing weights (with $\ell_2$ balance).} On $C$, define \[ \widehat W_c=\operatorname{diag}\!\big(\widehat\psi'(X_c\hat\beta_c)\big),\quad \widehat W_t'=\operatorname{diag}\!\big(\widehat\psi''(X_t\hat\beta_c)\big),\quad \widehat V_i=\widehat\psi''(X_{c,i}^\top\hat\beta_c)\,X_{c,i}\otimes X_{c,i}. \] Let\footnote{The factor $B_X$ arises from lifting pointwise derivative errors to vector gaps via Cauchy–Schwarz: for instance, \[ \left\|\frac{1}{n_t}\sum_{i:W_i=1}X_i\{\widehat\psi'(Z_i)-\psi'(Z_i)\}\right\|_2 \;\leqslant\; \sqrt{\frac{1}{n_t}\sum_{i}\|X_i\|_2^2}\,\|\widehat\psi'-\psi'\|_{L_2} \;\lesssim\; B_X\,r_{1,n}. \] and the last inequality follows with high probability given that the first factor concentrates around $\left\{\mathbb{E}\|X\|_2^2\right\}^{1 / 2}=\sqrt{\operatorname{tr}\left(\Sigma_c\right)}$. In practice we should consider a consistent estimate of the variance.} $B_X:=\sqrt{\operatorname{tr}(\Sigma_c)}$. The balancing weights $\gamma$ are chosen by
The additive terms $c_1 B_X r_{1,n}$ and $c_2 B_X^2 r_{2,n}$ absorb the derivative‑estimation error when passing from the oracle constraints ($\psi',\psi''$) to the feasible ones ($\widehat\psi',\widehat\psi''$). \paragraph{Step 4: Estimate the average treatment effect \(\tau\).} Compute the ATE as: \[ \hat{\tau} = \bar{Y}_t - \left\{ \frac{1}{n_t} \sum_{i:W_i=1} \widehat\psi(X_i^\top \hat{\beta}_c) + \sum_{i : W_i = 0} \gamma_i \left(Y_i^\text{obs} - \widehat\psi(X_i^\top \hat{\beta}_c)\right)\right\}, \] where \(\bar{Y}_t\) is the average observed outcome for the treated group.
\paragraph{Step 5: Cross-fitting.} Rotate folds to aggregate. \\
For the estimator $\hat{\tau}$ to be asymptotically normal, the main term in the error decomposition must be the stochastic component $\left|\gamma^{\top} \varepsilon_c\right|$, while all other bias and remainder terms must be asymptotically negligible. The standard deviation of the main term is $\|\gamma\|_2 \sigma \asymp n^{-1 / 2}$. Therefore, we require all remainder terms to vanish at a rate faster than $n^{-1 / 2}$, i.e., they must be $o_p\left(n^{-1 / 2}\right)$.
This leads to the following sufficient conditions, which formalize the trade-off between the model's sparsity, its dimensionality, and the smoothness of the unknown link function.
Condition (a) is a standard strong sparsity condition required for high-dimensional inference. Conditions (b) and (c) are new and represent the statistical price of not knowing the link function. They govern how large the dimension $p$ can be relative to the sample size $n$. Conditions (b) and (c) are the most stringent. It shows that for a fixed level of smoothness, the dimension $p$ cannot grow too quickly relative to $n$ for the theory to hold. This is the direct consequence of needing to control the error from the second-derivative estimation term, which involves a $p^2$-dimensional object $(X\otimes X)$.
In this paper, we addressed an open problem posed by athey2018approximate by extending the Approximate Residual Balancing (ARB) framework to handle high-dimensional nonlinear outcome models. We demonstrated that a naive application of ARB fails in this setting due to bias introduced by the link function's curvature. Our central contribution is the identification of a second-order balancing principle: to achieve valid inference, one must balance not only the covariates weighted by the link's first derivative but also their outer products weighted by the second derivative.
Based on this insight, we developed a practical new estimator and a novel optimization problem for constructing the required balancing weights. We provided rigorous theoretical guarantees, establishing that our estimator is $\sqrt{n}$-consistent and asymptotically normal under standard high-dimensional assumptions. This result holds without requiring any estimation of the propensity score, preserving the key advantage of the ARB paradigm. Furthermore, we extended our framework to the more flexible semi-parametric single-index model, showing that asymptotic normality can be retained at the cost of a stricter trade-off between model sparsity and the link function's smoothness.
In conclusion, this work provides a rigorous and practical solution to a key limitation of the original ARB framework. Our findings help bridge the gap between high-dimensional causal theory and the types of nonlinear models frequently used in empirical research.