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

TL;DR

This paper introduces a novel neural network approach for nonparametric instrumental variable regression, establishing global convergence and generalization bounds through a bilevel optimization framework and validating it empirically.

Contribution

It develops a first-order algorithm for neural network-based 2SLS in NPIV, addressing bilevel optimization challenges with a mean-field Langevin dynamics perspective.

Findings

Proves global convergence of neural networks in NPIV 2SLS.

Provides generalization bounds highlighting the trade-off in Lagrange multiplier choice.

Empirically demonstrates effectiveness on reinforcement learning benchmark.

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.