Technical Summary: Outcome-Aware Spectral Feature Learning for Instrumental Variable Regression

Problem Statement

The paper addresses the challenge of estimating causal effects in the presence of hidden confounders using Nonparametric Instrumental Variable (NPIV) regression. The core model assumes a structural relationship $Y = h_0(X) + U$ with $E[U|Z]=0$, where $X$ is the treatment, $Y$ is the outcome, $Z$ is the instrument, and $U$ is an unobserved confounder. The goal is to recover the structural function $h_0$ from i.i.d. samples of $(Y, X, Z)$.

This problem is formulated as a linear inverse problem $T h_0 = r_0$, where $T$ is the conditional expectation operator mapping $h(X)$ to $E[h(X)|Z]$. While state-of-the-art methods like SpecIV (Sun et al., 2025) learn spectral features by approximating the top singular subspaces of $T$, the authors identify a fundamental limitation: these features are outcome-agnostic. They capture the dominant modes of the relationship between $X$ and $Z$ but may fail to represent the structural function $h_0$ if $h_0$ lies outside the dominant singular subspace of $T$ (a scenario termed "spectral misalignment"). In such cases, standard spectral methods degrade significantly.

Methodology: Augmented Spectral Feature Learning

To mitigate spectral misalignment, the authors propose Augmented Spectral Feature Learning, a framework that makes the feature learning process outcome-aware.

1. Augmented Operator and Loss Function

The method introduces an augmented operator $T_\delta$ that incorporates information from the outcome $Y$. Specifically, $T_\delta$ is defined as a rank-one perturbation of $T$:
$$T_\delta: L^2(X) \times \mathbb{R} \to L^2(Z), \quad (h, a) \mapsto Th + a \cdot \delta \cdot r_0$$
where $r_0 = E[Y|Z]$ and $\delta$ is a hyperparameter controlling the strength of the outcome signal.

The learning objective is modified to minimize a novel contrastive loss $L^{(d)}_\delta(\theta, \omega)$, which combines the standard spectral loss $L^{(d)}_0$ (approximating $T$) with a regularization term derived from the outcome:
$$L^{(d)}_\delta(\theta, \omega) = L^{(d)}_0(\theta) - 2\delta E[Y \psi^{(d)}_\theta(Z)]^\top \omega + \omega^\top C_{\psi^{(d)}_\theta} \omega$$
Here, $\psi^{(d)}_\theta$ and $\phi^{(d)}_\theta$ are neural network feature maps for $Z$ and $X$, and $\omega$ is an auxiliary variable. Minimizing this loss is theoretically equivalent to finding the best rank-$d$ approximation of the augmented operator $T_\delta$.

2. Theoretical Mechanism

By augmenting $T$ with the outcome signal $r_0$, the method amplifies the singular components of $T$ that are relevant to $h_0$. Even if $h_0$ aligns with a "bad" (low singular value) direction of $T$, the perturbation can shift the dominant singular subspace of $T_\delta$ to align with $h_0$. This allows the learned features to capture the signal necessary for predicting $Y$, even when standard spectral features would miss it.

3. Estimation Procedure

The framework integrates with the Two-Stage Least Squares (2SLS) estimator.

• Stage 1: Learn features $\phi^{(d)}_\theta$ and $\psi^{(d)}_\theta$ by minimizing the empirical augmented loss on a dataset $\tilde{D}_m$.
• Stage 2: Estimate the structural function $h_0$ using the learned features in a standard 2SLS regression on a separate dataset $D_n$.

Key Contributions

1. Identification of Limitation: The authors formally identify that existing spectral methods for NPIV are outcome-agnostic, leading to performance degradation under spectral misalignment.
2. Augmented Framework: They propose Augmented Spectral Feature Learning, introducing the operator $T_\delta$ and a corresponding contrastive loss that incorporates outcome information into the feature learning process.
3. Theoretical Guarantees:
   • A non-asymptotic, high-probability error bound for the 2SLS estimator is provided, separating approximation error from statistical error.
   • The analysis demonstrates that the method remains robust in "bad" scenarios (spectral misalignment) where standard SpecIV fails. Specifically, if the signal-to-noise ratio is sufficient, the method can recover $h_0$ with a single feature even when $h_0$ aligns with a low-singular-value component of $T$.
4. Empirical Validation: The approach is validated on synthetic data, semi-synthetic dSprites benchmarks (including a new, more challenging structural function), and Off-Policy Evaluation (OPE) tasks in reinforcement learning.

Experimental Results

• Synthetic Data: In scenarios with poor spectral alignment (where $h_0$ aligns with low singular values of $T$), increasing $\delta$ significantly reduces the Mean Squared Error (MSE) compared to standard SpecIV ($\delta=0$).
• dSprites Benchmark:
  • On the standard "well-aligned" benchmark, the method shows modest improvements.
  • On a new "misaligned" benchmark ($h_{new}$), where the structural function depends on rotation (a low-singular-value feature), vanilla SpecIV severely underperforms. The outcome-aware method, with tuned $\delta$, matches or slightly exceeds the performance of Deep Feature IV (DFIV), a strong competitor.
• Off-Policy Evaluation (OPE): In reinforcement learning settings (CartPole, Mountain Car, Catch), the method demonstrates robustness. While no single method dominated all tasks, Augmented SpecIV showed superior adaptability to spectral misalignment compared to standard SpecIV, particularly in the CartPole environment.

Significance and Claims

The paper claims that Augmented Spectral Feature Learning addresses a fundamental limitation of current spectral methods by ensuring that learned representations are not only predictive of the treatment but also relevant to the causal outcome. The authors emphasize that this approach provides robustness in regimes where standard methods fail due to spectral misalignment.

The work is positioned as a step toward more adaptive NPIV estimation. The authors note that while their current approach relies on a rank-one augmentation, extending the framework to higher-rank perturbations (e.g., using multiple functions of $Y$) remains a promising direction for future research. The paper concludes that their method broadens the applicability of spectral feature-based NPIV, particularly in complex settings like off-policy evaluation where the target function may shift dynamically.