Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revisiting Bayesian Model Averaging in the Era of Foundation Models

Mijung Park

TL;DR

This work tackles robust ensembling of foundation models by revisiting Bayesian model averaging (BMA) with frozen pre-trained features and trainable linear heads, enabling principled weighting of diverse representations. To address practical constraints and distribution shifts, it introduces Optimized Model Averaging (OMA), directly optimizing ensemble weights to minimize predictive entropy without full fine-tuning. The approach relies on a Laplace approximation to compute approximate marginal likelihoods and a tractable (block-diagonal) Hessian to scale to large models, enabling effective fusion of both zero-shot and lightly-finetuned models for image and text tasks. Empirically, BMA improves in-distribution accuracy, while OMA offers notable gains under distribution shifts and in multilingual text settings, with substantial reductions in training effort and energy consumption compared to full fine-tuning of foundation models.

Abstract

We revisit the classical, full-fledged Bayesian model averaging (BMA) paradigm to ensemble pre-trained and/or lightly-finetuned foundation models to enhance the classification performance on image and text data. To make BMA tractable under foundation models, we introduce trainable linear classifiers that take frozen features from the pre-trained foundation models as inputs. The model posteriors over the linear classifiers tell us which linear heads and frozen features are better suited for a given dataset, resulting in a principled model ensembling method. Furthermore, we propose a computationally cheaper, optimizable model averaging scheme (OMA). In OMA, we directly optimize the model ensemble weights, just like those weights based on model posterior distributions in BMA, by reducing the amount of surprise (expected entropy of the predictions) we get from predictions of ensembled models. With the rapid development of foundation models, these approaches will enable the incorporation of future, possibly significantly better foundation models to enhance the performance of challenging classification tasks.

Revisiting Bayesian Model Averaging in the Era of Foundation Models

TL;DR

This work tackles robust ensembling of foundation models by revisiting Bayesian model averaging (BMA) with frozen pre-trained features and trainable linear heads, enabling principled weighting of diverse representations. To address practical constraints and distribution shifts, it introduces Optimized Model Averaging (OMA), directly optimizing ensemble weights to minimize predictive entropy without full fine-tuning. The approach relies on a Laplace approximation to compute approximate marginal likelihoods and a tractable (block-diagonal) Hessian to scale to large models, enabling effective fusion of both zero-shot and lightly-finetuned models for image and text tasks. Empirically, BMA improves in-distribution accuracy, while OMA offers notable gains under distribution shifts and in multilingual text settings, with substantial reductions in training effort and energy consumption compared to full fine-tuning of foundation models.

Abstract

We revisit the classical, full-fledged Bayesian model averaging (BMA) paradigm to ensemble pre-trained and/or lightly-finetuned foundation models to enhance the classification performance on image and text data. To make BMA tractable under foundation models, we introduce trainable linear classifiers that take frozen features from the pre-trained foundation models as inputs. The model posteriors over the linear classifiers tell us which linear heads and frozen features are better suited for a given dataset, resulting in a principled model ensembling method. Furthermore, we propose a computationally cheaper, optimizable model averaging scheme (OMA). In OMA, we directly optimize the model ensemble weights, just like those weights based on model posterior distributions in BMA, by reducing the amount of surprise (expected entropy of the predictions) we get from predictions of ensembled models. With the rapid development of foundation models, these approaches will enable the incorporation of future, possibly significantly better foundation models to enhance the performance of challenging classification tasks.

Paper Structure

This paper contains 48 sections, 1 theorem, 23 equations, 4 figures, 11 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. OpenClip
  4. Bayesian Model Averaging (BMA)
  5. Method
  6. Bayesian model averaging
  7. Notation
  8. A Challenge in computing the marginal likelihood of each model
  9. Likelihood of the data
  10. Prior
  11. The challenge
  12. Laplace approximation
  13. Approximate model marginal likelihood
  14. Posterior predictive distribution
  15. Optimized Model Averaging
  16. ...and 33 more sections

Key Result

Lemma 2.1

For any $j$ in $\{1, 2, \cdots, L\}$, where the expectation $\mathbb{E}$ is with respect to $\sum_{l=1}^L p(\mathbf{y}^*|\mathbf{x}^*, M_l, \mathcal{D}) p(M_l | \mathcal{D})$.

Figures (4)

  • Figure 1: The diagonal elements are orders of magnitude larger than off-diagonal elements (with the MAP estimates trained for ImageNet-1K. Performance in Table \ref{['tab:BMA']}). We show the values of a [400 × 400] subset of the full Hessian, which is [$10^6$x$10^6$], by subsampling 20 classes and 20 feature dimensions uniformly at random for traceability.
  • Figure 2: Model posterior weights given the Imagenet-1K training data. The first and the fourth feature extractors are deemed the most significant.
  • Figure 3: OMA weights learned by Algorithm \ref{['algo:OMA']} and individual model's performance. For visualization purpose, we divide each quantity by its maximum value, so that all the resulting values are between 0 and 1.
  • Figure 4: The diagonal elements are orders of magnitude larger than off-diagonal elements (with the MAP estimates trained for ImageNet-1K. Performance in Table \ref{['tab:BMA']}). Showing the magnitude of a subset of the Hessian for a [400 × 400] subset of the Hessian, which is originally roughly [1 million by 1 million], sampling 20 classes log-uniformly and 20 input dimensions uniformly.

Theorems & Definitions (1)

  • Lemma 2.1: Eq.4 in 11c48783-68bb-36b7-8053-175211b0eaa8