Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards a Unified Analysis of Neural Networks in Nonparametric Instrumental Variable Regression: Optimization and Generalization

Zonghao Chen, Atsushi Nitanda, Arthur Gretton, Taiji Suzuki

TL;DR

This work tackles nonparametric instrumental variable regression (NPIV) with neural-network features through a global-optimization lens. By lifting the two-stage regression into mean-field Langevin dynamics and reformulating the bilevel problem as a first-order penalty-Lagrangian, the authors design F$^2$BMLD, a fully first-order algorithm that achieves global convergence guarantees despite nonconvexity and nested dependence. They establish both convergence to the Lagrangian optimum and a Gamma-convergence result ensuring consistency with the original bilevel objective as regularization vanishes, along with a generalization bound that reveals a trade-off controlled by the Lagrange multiplier $\lambda$. Empirically, F$^2$BMLD attains competitive offline policy evaluation performance with more stable training and smaller batch-size requirements than prior DFIV methods. This unified, theory-backed framework advances understanding of optimization and generalization for neural NPIV models and suggests practical routes toward scalable, robust NPIV analysis in downstream tasks such as offline RL.

Abstract

We establish the first global convergence result of neural networks for two stage least squares (2SLS) approach in nonparametric instrumental variable regression (NPIV). This is achieved by adopting a lifted perspective through mean-field Langevin dynamics (MFLD), unlike standard MFLD, however, our setting of 2SLS entails a \emph{bilevel} optimization problem in the space of probability measures. To address this challenge, we leverage the penalty gradient approach recently developed for bilevel optimization which formulates bilevel optimization as a Lagrangian problem. This leads to a novel fully first-order algorithm, termed \texttt{F$^2$BMLD}. Apart from the convergence bound, we further provide a generalization bound, revealing an inherent trade-off in the choice of the Lagrange multiplier between optimization and statistical guarantees. Finally, we empirically validate the effectiveness of the proposed method on an offline reinforcement learning benchmark.

Towards a Unified Analysis of Neural Networks in Nonparametric Instrumental Variable Regression: Optimization and Generalization

TL;DR

This work tackles nonparametric instrumental variable regression (NPIV) with neural-network features through a global-optimization lens. By lifting the two-stage regression into mean-field Langevin dynamics and reformulating the bilevel problem as a first-order penalty-Lagrangian, the authors design FBMLD, a fully first-order algorithm that achieves global convergence guarantees despite nonconvexity and nested dependence. They establish both convergence to the Lagrangian optimum and a Gamma-convergence result ensuring consistency with the original bilevel objective as regularization vanishes, along with a generalization bound that reveals a trade-off controlled by the Lagrange multiplier . Empirically, FBMLD attains competitive offline policy evaluation performance with more stable training and smaller batch-size requirements than prior DFIV methods. This unified, theory-backed framework advances understanding of optimization and generalization for neural NPIV models and suggests practical routes toward scalable, robust NPIV analysis in downstream tasks such as offline RL.

Abstract

We establish the first global convergence result of neural networks for two stage least squares (2SLS) approach in nonparametric instrumental variable regression (NPIV). This is achieved by adopting a lifted perspective through mean-field Langevin dynamics (MFLD), unlike standard MFLD, however, our setting of 2SLS entails a \emph{bilevel} optimization problem in the space of probability measures. To address this challenge, we leverage the penalty gradient approach recently developed for bilevel optimization which formulates bilevel optimization as a Lagrangian problem. This leads to a novel fully first-order algorithm, termed \texttt{FBMLD}. Apart from the convergence bound, we further provide a generalization bound, revealing an inherent trade-off in the choice of the Lagrange multiplier between optimization and statistical guarantees. Finally, we empirically validate the effectiveness of the proposed method on an offline reinforcement learning benchmark.

Paper Structure

This paper contains 35 sections, 33 theorems, 233 equations, 1 figure, 2 algorithms.

Table of Contents

  1. Introduction
  2. Structure of the paper:
  3. Related work
  4. Notations:
  5. Two-stage Least Squares Regression (2SLS)
  6. Fixed feature 2SLS regression
  7. Deep feature instrumental variable regression
  8. Mean Field Langevin Dynamics (MFLD)
  9. Bilevel Mean Field Langevin Dynamics
  10. Mean field formulation of DFIV
  11. Penalty gradient method
  12. Convergence of F$^2$BMLD
  13. Convergence of the inner loop
  14. Convergence of the outer loop
  15. Generalization of F$^2$BMLD
  16. ...and 20 more sections

Key Result

Proposition 4.1

For any fixed $\mu_x\in\mathcal{P}_2(\mathbb{R}^{d_x})$, the mappings $\mu_z\mapsto F_1(\mu_x, \mu_z)$, $\mu_z\mapsto F_2(\mu_x, \mu_z)$ are linear convex. For any fixed $\mu_z\in\mathcal{P}_2(\mathbb{R}^{d_z})$, the mappings $\mu_x\mapsto F_1(\mu_x, \mu_z)$ the mapping $\mu_x\mapsto F_2(\mu_x, \mu_

Figures (1)

  • Figure 1: Left: Comparison of DFIV and F2BMLD in terms of target policy value. Right: Comparison of DFIV and F2BMLD training trajectories.

Theorems & Definitions (86)

  • Definition 1: First variation
  • Definition 2: Logarithmic Sobolev inequality bakry2013analysis
  • Remark 3.1
  • Definition 3: Bregman divergence
  • Remark 3.2
  • Remark 3.3
  • Proposition 4.1: Partial convexity of $F_1$ and $F_2$
  • proof
  • Proposition 4.2: Continuity of mapping $\mu_x \mapsto \mu_z^\ast(\mu_x)$
  • proof
  • ...and 76 more