Outcome-Aware Spectral Feature Learning for Instrumental Variable Regression

AI summary

Overview

Research area: Nonparametric instrumental variable (NPIV) regression and causal effect estimation with hidden confounders, sitting at the intersection of statistical machine learning, operator/spectral learning, and deep representation learning.

Technical level: Advanced. The paper works in Hilbert spaces, uses singular value decompositions of compact conditional-expectation operators, perturbation arguments, and non-asymptotic concentration bounds.

Scope (one sentence): The paper proposes Augmented Spectral Feature Learning, an outcome-aware spectral feature method for NPIV that learns features from an augmented operator incorporating the outcome, and provides theoretical guarantees plus synthetic, semi-synthetic, and reinforcement-learning experiments.

What This Paper Is About

In IV regression one wants the causal effect of a treatment X on outcome Y when an unobserved confounder U affects both, using an instrument Z that is related to X but uncorrelated with U. The NPIV model assumes Y = h₀(X) + U with E[U | Z] = 0, and reformulates the problem as the linear inverse problem 𝒯h₀ = r₀, where r₀ = E[Y | Z] and 𝒯 maps h to E[h(X) | Z].

Existing spectral methods such as SpecIV (Sun et al., 2025) learn neural network features by approximating the top singular subspaces of 𝒯. These features are "outcome-agnostic": they capture the dominant treatment-instrument relationship only, and, as Meunier et al. (2026) showed, the approach degrades when h₀ is poorly aligned with the top singular functions of 𝒯. This paper's goal is to make the feature-learning process outcome-aware so it remains effective under such spectral misalignment.

Key Contributions

1. Diagnosis and fix for outcome-agnostic spectral features. The authors identify the misalignment limitation of existing spectral NPIV methods and propose Augmented Spectral Feature Learning, introducing an augmented operator 𝒯_δ that incorporates outcome information, together with a new contrastive loss for learning task-specific spectral features.

2. Comprehensive theory. They provide a non-asymptotic, high-probability generalization error bound for the resulting 2SLS estimator (agnostic to the choice of representation), plus an analysis characterizing when the augmented approach stays robust to the spectral misalignment that defeats previous methods.

3. Empirical validation on challenging problems. They validate on synthetic and semi-synthetic examples, including a new and more challenging version of the dSprites IV benchmark (Xu et al., 2021), plus an Off-Policy Evaluation (OPE) experiment in reinforcement learning (Chen et al., 2022) with dynamically changing environments.

4. Operator-level characterization. Proposition 4.1 shows minimizing the augmented loss is equivalent to finding the best rank-d approximation of 𝒯_δ, connecting the learning objective to a truncated SVD of a compact augmented operator.

Main Findings

• Outcome-agnostic features fail under misalignment: Spectral features spanning the top singular subspaces of 𝒯 can be uninformative for the final task when h₀ lies outside the dominant subspace, so the method can fail even though it is optimal when h₀ is well-aligned with the top singular functions.

• An augmented operator solves it in principle: Defining 𝒯_δ: L₂(X) × ℝ → L₂(Z), (h, a) ↦ 𝒯h + a·δ·r₀, written as 𝒯_δ = [𝒯 | δr₀] = 𝒯[I_{L₂(X)} | δh₀], the added "column" amplifies components of h₀ that would otherwise sit in the low-singular-value region of 𝒯, improving alignment with the top spectral features of the augmented operator.

• Exact optimality characterization: The augmented loss satisfies ℒ_δ^(d)(θ, ω) ≥ −‖𝒯_δ^(d)‖²_HS, with equality if and only if the learned operator Ψ_θ^(d)[Φ_θ^(d)* | ω] equals 𝒯_δ^(d) (Proposition 4.1).

• The regularizer is a prediction error: δ^(−2)ℛ_δ^(d)(θ) equals, up to a constant independent of θ, the mean squared error of the best linear predictor of Y from ψ_θ^(d)(Z). Minimizing the original profile loss requires backpropagating through an inverse covariance matrix, so the authors instead minimize a jointly convex-in-ω equivalent loss, whose minimum over ω is ω_θ^(d) = δ C_{ψ_θ^(d)}^{−1} E[Y ψ_θ^(d)(Z)].

• High-probability 2SLS bound: Under Assumptions 2–4 (representation boundedness with constant ρ ≥ 1; positive measure of ill-posedness c_{φ_θ^(d), ψ_θ^(d)} = σ_d(C_{Z,θ}^{−1/2} C_{ZX,θ} C_{X,θ}^{−1/2}); sub-Gaussian noise U and approximation error (h₀ − h_θ)(X)), if n ≥ 16 d ρ² log²(4d/τ) c^{−2} then with probability at least 1 − τ, ‖ĥ_θ − h₀‖{L₂(X)} ≤ C(‖h₀ − h_θ‖{L₂(X)} + (1/c)√(d/n)√(σ_U² + ρ²/n) log(4/τ)). This separates a deterministic approximation error from a statistical error and improves on Chen and Christensen (2018, Theorem B.1) by holding in high probability under a sub-Gaussian assumption.

• Spectral gap controls the error: With the singular components of 𝒯 partitioned into a d-dimensional signal subspace (N̄) and a noise subspace (N̲), and with s_d = V̄_d V̄_d* h₀ and q_d = V̲_d V̲_d* h₀, the projection difference decomposes as ‖Π_{Ū_d} − Π_{Ψ̂}‖op ≤ ‖Π{Ū_d} − Π_{Ψ_}‖op + ‖Π{Ψ_} − Π_{Ψ̂}‖op; the first term is bounded by δ‖q_d‖{L₂(X)} / γ_d(δ), where γ_d(δ) = ‖[Λ̄_d(I + δ² ᾱ ᾱᵀ)^{1/2}]^{−1}‖_op^{−1} − ‖Λ̲_d‖_op, and the second term is the optimality gap ℰ_d(θ̂_m, ω̂_m, δ) = ‖𝒯_δ^(d) − Ψ_θ^(d)[Φ_θ^(d)* | ω]‖_op.

• "Good" scenario: When s_d is spanned by the top-d singular functions of 𝒯 (N̄ = {1, …, d}), γ_d(0) = λ_d − λ_{d+1} is positive and the method works efficiently even with δ = 0, giving ‖h₀ − ĥ_θ‖{L₂} ≲ ‖q_d‖{L₂} + ℰ_d(θ̂_m, ω̂_m, 0)/λ_d + (1/λ_d)√(d/n) log τ^{−1}. The paper states the corresponding "bad" scenario analysis follows from the same decomposition.

• Numerical results: The abstract states the approach is validated on challenging benchmarks, and the contributions list synthetic and semi-synthetic examples plus a dSprites IV benchmark that is new and more challenging than the version of Xu et al. (2021), and an OPE experiment. No numerical performance figures, dataset sizes, or specific benchmark scores appear in the provided paper content; they are reported in the appendix and experiment section, which were truncated.

Methodology in Plain English

The starting point is that the instrument-to-treatment relationship can be summarized by the dominant "directions" (singular functions) of the operator 𝒯, and SpecIV learns neural features that approximate them. The problem is that the directions that best explain treatment from the instrument are not necessarily the directions that best explain the outcome. The authors' fix is to add a small amount of outcome information into the operator before taking its spectral decomposition: they build an augmented operator 𝒯_δ = [𝒯 | δr₀], where r₀ = E[Y | Z] is the best prediction of the outcome from the instrument, and δ is a hyperparameter controlling how much weight that extra information gets.

Practically, they add a penalty term to the SpecIV objective that rewards learned instrument features for linearly predicting Y. Because the direct penalty involves inverting a covariance matrix, which is numerically unstable to backpropagate through, they rewrite the objective with an extra learnable vector ω and minimize a form that is convex in ω for fixed features. They use two independent datasets: one of size m for learning features and one of size n for estimating the structural function, avoiding the need to reuse samples for both stages. The theoretical analysis then asks how well the learned features approximate the signal subspace of the original operator, which requires bounding how much the augmented operator's top singular subspaces move away from those of the original operator, and how much the signal-to-noise singular gap protects against that movement.

Why This Matters

For causal inference research, the paper shows that a well-known family of spectral NPIV methods has a structural blind spot, and that a small, principled modification — injecting outcome information into the operator being decomposed — repairs it without discarding the spectral framework or its computational advantages. The high-probability 2SLS guarantee also strengthens the standard style of analysis in this area, which typically only gives guarantees in expectation.

Real-world applications named or implied in the paper:

• Economics of education: estimating the effect of years of schooling on wages when unobserved factors such as innate ability or family background confound the relationship, using proximity to college as an instrument (the Card, 1993 setting).
• Off-policy evaluation in reinforcement learning: value function estimation can be framed as an IV problem (Hu et al., 2025), and the paper includes an OPE experiment showing the approach remains robust in dynamically changing environments.
• Learning evolution operators in scientific domains: molecular dynamics and climate science (Turri et al., 2025), where spectral decompositions of conditional operators are useful.
• Any setting where a conditional operator must be decomposed but the downstream task is known, since the method applies beyond IV regression.

Industry relevance: firms and agencies that need credible causal estimates from observational data with unmeasured confounders — for example, pricing, advertising, education and labor policy evaluation — depend on methods that remain reliable when the standard "dominant direction" assumption fails. Making feature learning task-aware rather than purely instrument-driven is a general pattern that transfers to any two-stage learned estimator.

Future Directions

• General rank-K perturbation theory: the authors note that extending the augmentation to multiple functions of Y, such as higher conditional moments E[Y^k | Z], keeps the same learning objective and truncated-SVD optimality characterization, but a full analysis of general rank-K perturbations requires further development of the perturbation framework and is explicitly left for future work.
• Validation beyond the "good" regime and the truncated experiments: the paper sets up a good/bad scenario dichotomy, and the bad-scenario consequences are stated to follow from the decomposition in Theorem C.6; the numerical results reported in the provided content are qualitative, so larger-scale quantitative comparisons are open.
• Scope of the augmentation: choosing δ, the number of augmented directions, and which outcome-derived functions to include is not resolved by the theory presented in this truncated content.
• Comparison in norms and assumptions: the paper contrasts its approach with Bruns-Smith (2025), whose strong-norm convergence holds under more stringent restrictions on the ill-posedness of the inverse problems; closing or characterizing that gap is a natural next question.

Target Audience

Researchers and graduate students in statistical machine learning, causal inference, and econometrics who work on instrumental variable regression, spectral/operator learning, or two-stage estimators with learned representations. It will also interest reinforcement learning researchers working on off-policy evaluation and scientists who learn spectral decompositions of conditional or evolution operators. Readers unfamiliar with Hilbert space operators, singular value decompositions of conditional expectation operators, or ill-posed inverse problems will find the theory section demanding.

Authors’ abstract

We address the problem of causal effect estimation in the presence of hidden confounders using nonparametric instrumental variable (IV) regression. An established approach is to use estimators based on learned spectral features, that is, features spanning the top singular subspaces of the operator linking treatments to instruments. While powerful, such features are agnostic to the outcome variable. Consequently, the method can fail when the true causal function is poorly represented by these dominant singular functions. To mitigate, we introduce Augmented Spectral Feature Learning, a framework that makes the feature learning process outcome-aware. Our method learns features by minimizing a novel contrastive loss derived from an augmented operator that incorporates information from the outcome. By learning these task-specific features, our approach remains effective even under spectral misalignment. We provide a theoretical analysis of this framework and validate our approach on challenging benchmarks.