[Paper Review] Gaussian Process Surrogate Models for Neural Networks

Michael Y. Li, Erin Grant+1 more|arXiv (Cornell University)|Aug 11, 2022

Gaussian Processes and Bayesian Inference4 citations

TL;DR

This paper proposes a data-driven approach to constructing Gaussian process (GP) surrogate models for finite-width neural networks by learning kernel functions from empirical predictions, rather than deriving them analytically from infinite-width limits. The method captures key phenomena like spectral bias and enables practical applications such as identifying influential training points and predicting generalization performance, demonstrating that learned GPs serve as interpretable, predictive surrogates for neural network behavior.

ABSTRACT

Not being able to understand and predict the behavior of deep learning systems makes it hard to decide what architecture and algorithm to use for a given problem. In science and engineering, modeling is a methodology used to understand complex systems whose internal processes are opaque. Modeling replaces a complex system with a simpler, more interpretable surrogate. Drawing inspiration from this, we construct a class of surrogate models for neural networks using Gaussian processes. Rather than deriving kernels for infinite neural networks, we learn kernels empirically from the naturalistic behavior of finite neural networks. We demonstrate our approach captures existing phenomena related to the spectral bias of neural networks, and then show that our surrogate models can be used to solve practical problems such as identifying which points most influence the behavior of specific neural networks and predicting which architectures and algorithms will generalize well for specific datasets.

Motivation & Objective

• To address the lack of interpretability and predictability in deep learning systems, especially when deployed in real-world applications.
• To develop a modeling framework that treats neural networks as black boxes and replaces them with simpler, interpretable surrogate models.
• To empirically learn kernel functions from finite neural network predictions, rather than relying on theoretical limits or analytic derivations.
• To enable practical applications such as identifying influential training points and predicting generalization performance across architectures and datasets.
• To demonstrate that GP surrogates can capture known phenomena like spectral bias and provide a cohesive, interpretable framework for understanding neural network behavior.

Proposed method

• Learn kernel hyperparameters of a GP surrogate by optimizing the marginal log-likelihood (mLL) of the GP using predictions from ensembles of finite-width neural networks.
• Use a spectral mixture kernel (SMK) with Q=10 mixture components to model the GP covariance function, allowing flexible representation of complex patterns in neural network predictions.
• Train the GP surrogate via batch gradient descent with Adam optimizer on predictions from 50 randomly initialized fully connected networks with ReLU or sinusoidal activations and varying depth (16 or 32 layers).
• Infer kernel hyperparameters by selecting the configuration with the highest mLL across three random initializations to ensure robustness.
• Leverage the analytic properties of GPs, such as closed-form marginal likelihood and leave-one-out predictive distributions, to reduce computational cost in influence and generalization gap analysis.
• Compare learned GP priors with neural network priors by visualizing samples and assessing their ability to reflect structural properties like periodicity and depth-dependent behavior.

Experimental results

Research questions

• RQ1 Can data-driven GP surrogates capture known phenomena such as spectral bias in neural networks?
• RQ2 How well can GP surrogates model the prior and inductive biases of different neural network families (e.g., ReLU vs. sinusoidal networks)?
• RQ3 Can GP surrogates identify which training points most influence the behavior of a specific neural network?
• RQ4 Can GP surrogates predict which neural network architectures and optimization algorithms will generalize well on a given dataset?
• RQ5 How do GP surrogate models compare to theoretical infinite-width limits in capturing real-world neural network behavior?

Key findings

• The learned GP surrogates successfully capture the spectral bias of neural networks, showing that the range of spectral frequencies in the kernel increases with training iterations, consistent with prior findings.
• The GP surrogates reflect architectural differences: 32-layer networks exhibit faster-varying behavior with shorter lengthscales, while ReLU networks show increasing-decreasing patterns near cusps.
• The surrogates accurately model the prior behavior of different neural network families, including periodicity in sinusoidal networks and depth-dependent variation in activation patterns.
• The method enables efficient identification of influential training points through leave-one-out predictive distributions, reducing the need for costly retraining.
• The framework predicts generalization performance across architectures and datasets by analyzing the learned GP kernel structure and hyperparameters.
• The data-driven approach outperforms theoretical infinite-width derivations in capturing practical, finite-network behavior, particularly in capturing non-asymptotic phenomena like pathological initialization effects.