Statistical mechanics of extensive-width Bayesian neural networks near interpolation

2026-07-27 · Source: stat.ML updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Theoretical Machine Learning · Depth: Expert, extended

Summary

A statistical mechanics analysis of two-layer fully connected Bayesian neural networks (BNNs) with extensive width, proportional to input dimension, and generic weight/activation functions, reveals a rich phenomenology of learning transitions. Focusing on Bayes-optimal learning in a teacher-student scenario near interpolation, where data and parameters are comparable (n=Θ(d^2)=Θ(kd)), the study finds that feature learning emerges. Crucially, the more strongly features (hidden neurons of the target) contribute to observed responses, the less data is needed to learn them. When data is scarce, the model learns non-linear combinations of teacher weights, with "specialization" (aligning weights with the teacher's) occurring only with sufficient data, often proving algorithmically hard to find.

Key takeaway

For AI and research scientists developing or deploying extensive-width Bayesian neural networks, understanding the "specialization" phase is critical. Your models may struggle to align with true teacher features, especially with limited data or certain readout distributions, leading to sub-optimal performance. Be aware that achieving full specialization can be computationally hard, potentially requiring exponential training time with increasing dimensions. Consider using algorithms like GAMP-RIE, which can match universal performance even when specialization is difficult to reach.

Key insights

Extensive-width Bayesian neural networks exhibit distinct learning phases, including "specialization" where student weights align with teacher features.

Principles

• Feature learning in BNNs depends on data availability and the strength of feature contribution.
• "Specialization" transitions occur when sufficient data allows student weights to align with teacher weights.
• Higher-order Hermite decomposition terms become informative only after the specialization phase.

Method

A statistical mechanics framework combines the replica method with the Harish Chandra–Itzykson–Zuber (HCIZ) integral to derive predictive formulas for generalization error. It also adapts GAMP-RIE for generic activation functions.

In practice

• GAMP-RIE offers polynomial-time prediction matching universal theory, even in computationally hard phases.
• Algorithmic hardness for achieving specialization can increase exponentially with input dimension "d" for certain readout distributions.
• Consider the feature-dependent "signal-to-noise ratio" (readout amplitude) for data efficiency in learning specific features.

Topics

• Statistical Mechanics
• Bayesian Neural Networks
• Feature Learning
• Generalization Error
• Replica Method
• Algorithmic Hardness
• GAMP-RIE

Code references

• maillard-j/GAMP-RIE

Best for: AI Scientist, Research Scientist

Related on AIssential

• Optimal generalisation and learning transition in extensive-width shallow neural networks near interpolation — Extensive-width shallow neural networks exhibit a phase transition between universal and specialization learning, impacting generalization and algorithmic tractability.
• Statistical physics of deep learning: Optimal learning of a multi-layer perceptron near interpolation — Statistical physics reveals deep learning's feature learning limits and optimal specialization in the interpolation regime.
• Equivariance and Augmentation for Bayesian Neural Networks — Data augmentation can achieve exact equivariance in Bayesian Neural Networks under specific conditions, enhanced by novel symmetrization techniques.
• Beyond Surface Statistics: Robust Conformal Prediction for LLMs via Internal Representations — Internal LLM representations offer more robust uncertainty signals for conformal prediction, especially under distribution shifts.
• Kernel Renormalization in Bayesian Deep Neural Networks: the Equivalent Wishart Ansatz in the Proportional Regime — The paper proposes an equivalent Wishart Ansatz and kernel renormalization to predict Bayesian DNN generalization in the proportional limit.
• The Evolution and Mechanics of Neural Networks:rom Brain-Inspired Models to Modern AI — Neural networks learn complex patterns from data by adjusting interconnected artificial neurons through iterative optimization.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by stat.ML updates on arXiv.org.