Neural Stein critics with staged $L^2$-regularization

TL;DR

This paper introduces a staged $L^2$-regularization method for training neural Stein critics, connecting it to NTK theory, and demonstrates improved high-dimensional distribution testing and generative model evaluation.

Contribution

It develops a novel staged regularization procedure for neural Stein critics, with theoretical guarantees based on NTK theory and lazy training analysis.

Findings

Staged regularization improves critic training in high dimensions.

Theoretical convergence rate of $O(n^{-1/2})$ for critic learning.

Effective application to high-dimensional data and image generative models.

Abstract

Learning to differentiate model distributions from observed data is a fundamental problem in statistics and machine learning, and high-dimensional data remains a challenging setting for such problems. Metrics that quantify the disparity in probability distributions, such as the Stein discrepancy, play an important role in high-dimensional statistical testing. In this paper, we investigate the role of $L^2$ regularization in training a neural network Stein critic so as to distinguish between data sampled from an unknown probability distribution and a nominal model distribution. Making a connection to the Neural Tangent Kernel (NTK) theory, we develop a novel staging procedure for the weight of regularization over training time, which leverages the advantages of highly-regularized training at early times. Theoretically, we prove the approximation of the training dynamic by the kernel optimization, namely the ``lazy training'', when the $L^2$ regularization weight is large, and training on $n$ samples converge at a rate of ${O}(n^{-1/2})$ up to a log factor. The result guarantees learning the optimal critic assuming sufficient alignment with the leading eigen-modes of the zero-time NTK. The benefit of the staged $L^2$ regularization is demonstrated on simulated high dimensional data and an application to evaluating generative models of image data.