Exploring the Uncertainty Properties of Neural Networks' Implicit Priors in the Infinite-Width Limit (2010.07355v1)

Abstract: Modern deep learning models have achieved great success in predictive accuracy for many data modalities. However, their application to many real-world tasks is restricted by poor uncertainty estimates, such as overconfidence on out-of-distribution (OOD) data and ungraceful failing under distributional shift. Previous benchmarks have found that ensembles of neural networks (NNs) are typically the best calibrated models on OOD data. Inspired by this, we leverage recent theoretical advances that characterize the function-space prior of an ensemble of infinitely-wide NNs as a Gaussian process, termed the neural network Gaussian process (NNGP). We use the NNGP with a softmax link function to build a probabilistic model for multi-class classification and marginalize over the latent Gaussian outputs to sample from the posterior. This gives us a better understanding of the implicit prior NNs place on function space and allows a direct comparison of the calibration of the NNGP and its finite-width analogue. We also examine the calibration of previous approaches to classification with the NNGP, which treat classification problems as regression to the one-hot labels. In this case the Bayesian posterior is exact, and we compare several heuristics to generate a categorical distribution over classes. We find these methods are well calibrated under distributional shift. Finally, we consider an infinite-width final layer in conjunction with a pre-trained embedding. This replicates the important practical use case of transfer learning and allows scaling to significantly larger datasets. As well as achieving competitive predictive accuracy, this approach is better calibrated than its finite width analogue.

• The paper introduces an NNGP-based model for both classification and regression, enabling direct sampling from the Bayesian posterior in infinite-width neural networks.
• Empirical evaluations show that NNGP models achieve lower expected calibration errors and negative log-likelihoods on out-of-distribution and shifted data compared to finite-width networks.
• The study extends the approach to practical transfer learning by integrating an infinite-width final layer, bridging theoretical insights with real-world applications.

Essay on "Exploring the Uncertainty Properties of Neural Networks’ Implicit Priors in the Infinite-Width Limit"

The paper "Exploring the Uncertainty Properties of Neural Networks’ Implicit Priors in the Infinite-Width Limit" provides an in-depth analysis of the calibration and uncertainty quantification capabilities of neural networks in the infinite-width regime. By leveraging the theoretical framework of Neural Network Gaussian Processes (NNGP), the authors aim to understand better the implicit priors that infinitely-wide neural networks impose on function space. The paper offers a comparison between these infinite-width models and their finite-width counterparts, exploring their calibration under various distributional conditions.

Key Contributions

1. NNGP for Classification and Regression: The authors construct a probabilistic model for multi-class classification using the NNGP in conjunction with a softmax link function. This allows direct sampling from the Bayesian posterior, offering insights into function space priors and providing a direct comparison with finite neural network analogs. For regression, they employ NNGP to exactly solve for the posterior, using Gaussian noise to model output uncertainties.
2. Empirical Evaluations: A thorough empirical analysis demonstrates that NNGP-based models are better calibrated than finite-width neural networks, particularly on out-of-distribution (OOD) and distributionally-shifted data. The paper highlights the strength of NNGP in capturing uncertainty, reflected in lower expected calibration errors (ECE) and negative log-likelihoods (NLL).
3. Transfer Learning with NNGP Final Layer: By using an infinite-width final layer with a pre-trained model, termed NNGP-LL, the paper extends its method to more practical domains like transfer learning. This approach shows improved calibration and competitive accuracy in larger datasets by harnessing pre-trained embeddings, showcasing its potential for real-world applications.
4. Metrics and Evaluation: The paper uses extensive benchmarks and compares methods against standard baselines on well-known datasets including CIFAR-10 and UCI regression tasks. Strong performance across various scenarios is noted, especially with CNN-GP kernels, which outperform others at higher corruption levels.

Methodology

The authors employ recent advances in understanding neural networks as Gaussian processes in the infinite-width limit. They utilize elliptical slice sampling for posterior inference in classification and exact Bayesian updates in regression, allowing exact posterior calculations. These methodological choices underscore the benefits of bridging Bayesian inference techniques with deep learning's empirical successes.

Implications and Future Directions

The implications of this research are twofold. Practically, the introduction of NNGP and its variants for modeling uncertainties can lead to more reliable model predictions in domains demanding high robustness. Theoretically, understanding and utilizing the implicit priors of neural networks may provide pathways to improved training paradigms that inherently incorporate uncertainty estimates.

Future work might explore more diverse architectures, such as attention-based models or graph neural networks, to generalize the applicability of the infinite-width framework. Additionally, integrating these findings into scalable and computationally efficient algorithms remains a crucial challenge for deploying such models in real-world systems.

Conclusion

Overall, the research explores the undervalued aspect of uncertainty in neural networks using an ambitious approach that merges Gaussian processes with neural architecture. Despite computational challenges, the demonstrated benefits suggest a promising avenue for advancing both the theoretical understanding and practical deployment of uncertainty-aware AI models.