[Paper Review] Scalable PAC-Bayesian Meta-Learning via the PAC-Optimal Hyper-Posterior: From Theory to Practice

Jonas Rothfuss, Martin Josifoski+2 more|arXiv (Cornell University)|Nov 14, 2022

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes PACOH, a scalable PAC-Bayesian meta-learning framework that derives the closed-form optimal hyper-posterior for principled meta-level regularization, enabling strong generalization guarantees with minimal meta-training tasks. By avoiding bi-level optimization and leveraging variational inference, PACOH achieves state-of-the-art performance in predictive accuracy and uncertainty estimation across Bayesian neural networks and Gaussian processes, especially in low-data regimes.

ABSTRACT

Meta-Learning aims to speed up the learning process on new tasks by acquiring useful inductive biases from datasets of related learning tasks. While, in practice, the number of related tasks available is often small, most of the existing approaches assume an abundance of tasks; making them unrealistic and prone to overfitting. A central question in the meta-learning literature is how to regularize to ensure generalization to unseen tasks. In this work, we provide a theoretical analysis using the PAC-Bayesian theory and present a generalization bound for meta-learning, which was first derived by Rothfuss et al. (2021a). Crucially, the bound allows us to derive the closed form of the optimal hyper-posterior, referred to as PACOH, which leads to the best performance guarantees. We provide a theoretical analysis and empirical case study under which conditions and to what extent these guarantees for meta-learning improve upon PAC-Bayesian per-task learning bounds. The closed-form PACOH inspires a practical meta-learning approach that avoids the reliance on bi-level optimization, giving rise to a stochastic optimization problem that is amenable to standard variational methods that scale well. Our experiments show that, when instantiating the PACOH with Gaussian processes and Bayesian Neural Networks models, the resulting methods are more scalable, and yield state-of-the-art performance, both in terms of predictive accuracy and the quality of uncertainty estimates.

Motivation & Objective

• To address overfitting in meta-learning when only a small number of related tasks are available, a common real-world limitation.
• To develop a theoretically grounded, scalable meta-learning method that ensures generalization to unseen tasks through principled regularization.
• To derive a closed-form solution for the optimal hyper-posterior (PACOH) that minimizes PAC-Bayesian generalization bounds.
• To enable practical, scalable meta-learning without relying on computationally expensive bi-level optimization.
• To empirically validate the superiority of meta-learning over per-task learning within the PAC-Bayesian framework.

Proposed method

• Derives a novel PAC-Bayesian generalization bound for meta-learning that applies to unbounded loss functions, extending applicability to regression and probabilistic inference.
• Tightens the bound by relating it directly to the generalized marginal log-likelihood instead of task-specific posteriors, simplifying optimization.
• Identifies the PAC-optimal hyper-posterior (PACOH) in closed form, which minimizes the upper bound on meta-generalization error.
• Replaces the intractable bi-level optimization of prior PAC-Bayesian meta-learning with a stochastic optimization problem amenable to standard variational inference.
• Applies the PACOH framework to Bayesian neural networks and Gaussian processes using neural network-based priors (PACOH-NN) for scalable meta-learning.
• Employs standard Bayesian inference and bandit algorithms (UCB, Thompson Sampling) to evaluate performance in sequential decision-making tasks.

Experimental results

Research questions

• RQ1 Under what conditions does meta-learning provide a significant improvement over per-task learning in terms of generalization guarantees?
• RQ2 Can a closed-form optimal hyper-posterior be derived within the PAC-Bayesian framework for meta-learning, and does it lead to better generalization?
• RQ3 How does the proposed PACOH framework scale and perform in low-data meta-learning scenarios with only a few tasks?
• RQ4 To what extent do PACOH-based priors improve performance in real-world sequential decision-making tasks such as vaccine design?
• RQ5 How do the uncertainty estimates from PACOH-based models compare to those from standard priors in complex, high-dimensional problems?

Key findings

• PACOH-UCB and PACOH-TS achieve near-optimal regret in fewer than 10 iterations on the MHC-I peptide design task, significantly outperforming standard BNN and GP priors.
• The meta-learned PACOH-NN prior reduces regret substantially compared to zero-centered Gaussian priors, demonstrating effective positive transfer across alleles with only five meta-training tasks.
• In regression and classification tasks, PACOH-based methods achieve state-of-the-art predictive accuracy and uncertainty estimation quality across both Bayesian neural networks and Gaussian processes.
• The theoretical generalization bounds derived in the paper are empirically tighter than per-task learning bounds, particularly in low-data regimes.
• The framework enables scalable meta-learning without bi-level optimization, making it practical for complex models like deep Bayesian networks.
• The case study confirms that PACOH provides strong, principled regularization, allowing reliable meta-learning even when meta-training data is scarce.