Optimizable Model Averaging for Ensembling Heterogeneous Model Outputs

arXiv:2505.21857 · cs.LG, stat.ML · Submitted 2025-05-28 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Optimizable Model Averaging for Ensembling Heterogeneous Model Outputs".

Jane: The gist The authors revisit classical Bayesian model averaging (BMA) to ensemble pre-trained and/or lightly-finetuned foundation models by introducing trainable linear classifiers and an optimizable model averaging scheme (OMA) to…

Tom: First, who's behind it and why it matters.

Paper summary: Jane: We’ve covered the main points of "Optimizable Model Averaging for Ensembling Heterogeneous Model Outputs," focusing on how it uses trainable linear classifiers and the OMA scheme to combine different foundation models intelligently.

Tom: The authors, Mijung Park and their team, essentially propose a principled way to ensemble pre-trained and lightly fine-tuned foundation models by learning model combination weights directly through optimization.

Lu: What this means in simpler terms is that instead of guessing which model is best, the system learns the optimal way to weight them based on what each model thinks about the data.

Meng: It's a structured approach to combining complex AI systems, which could make building robust classification tools much more systematic for developers.

Lalam: For us in culture and information systems, this points toward developing AI that is more reliable because it’s explicitly managing its uncertainty across different model outputs.

Jane: The title of the paper, "Optimizable Model Averaging for Ensembling Heterogeneous Model Outputs," really captures the essence—it's about making the averaging process itself optimizable and tailored to how different models are behaving.

Tom: So, if you’re listening and you’re wondering what this actually changes for classification tasks, it suggests a more efficient way to leverage a diverse set of AI models without needing massive retraining efforts for every combination.

Jane: It moves the conversation from just 'which model is best' to 'how do we optimally combine these different models given their individual uncertainties.'

Tom: That’s the big picture here. We’ll keep an eye on how this OMA approach scales when we start talking about even more complex, heterogeneous AI architectures down the road.

Lu: And honestly, it opens up a lot of avenues for exploring model diversity in ways that were previously too computationally demanding to explore fully.

Meng: From an engineering side, the focus on optimizing weights instead of retraining everything seems like a much more viable path forward for integrating these massive models into real applications.

Conclusion: Tom: So we’re wrapping up this look at Optimizable Model Averaging for Ensembling Heterogeneous Model Outputs by Mijung Park and their team, and what this whole setup actually means for how we use these big foundation models.

Jane: It boils down to them figuring out a way to combine different AI models—models that have different designs—without having to retrain everything from scratch every time they want something better.

Lu: They introduce this specific method, OMA, where you learn these weights that tell you how much influence each model should have on the final prediction, keeping all the actual model parameters frozen.

Meng: So it’s about using a clever mathematical trick to average out these different predictions in a way that’s computationally lighter than fine-tuning fifty models at once.

Lalam: And for us, this means we can get more reliable results from these complex systems because we're not just picking the one that looked best on a test set, but learning the right balance across all of them.

Tom: Exactly, it’s about making sure that when you use a diverse set of models, you actually get a better result than just picking the single most confident one.

Jane: The authors are showing that this optimization process works well for both images and text classification tasks, which is pretty broad.

Lu: They found that this approach improves performance on certain datasets significantly over simple output averaging, even when using models trained in different ways.

Meng: I’m interested in the practical side here—how much time does setting up this optimization take compared to just running a single model?

Lalam: It seems like it trades a little more setup work for a much more robust and adaptable final AI system.

Tom: Right, so they’re not just finding a static combination; they’re finding an adjustable one that adapts based on the data you feed it.

Jane: And the main thing is that this method lets us keep using those huge pre-trained foundation models without needing to do all that heavy lifting on them again.

Lu: It opens up possibilities for building even more complex AI pipelines where you might have a hundred different specialized models, and OMA could manage them all efficiently.

Meng: That’s interesting because it moves the challenge from training massive new things to learning smarter ways to use what we already have available.

Lalam: It really shows how improving the way we combine existing models can lead to a more versatile and trustworthy AI culture for everyone using these tools.

Department of Computer Science, University of British Columbia

cs.LG, stat.ML

Submitted: 2025-05-28

Updated: 2026-10-08

Code: https://github.com/mlfoundations/open_clip

License: http://creativecommons.org/licenses/by-nc-nd/4.0/

Importance score: 73/100

The gist: The gist The authors revisit classical Bayesian model averaging (BMA) to ensemble pre-trained and/or lightly-finetuned foundation models by introducing trainable linear classifiers and an optimizable

Key concepts

Bayesian Model Averaging (BMA): BMA is a method to combine the predictions of multiple models by weighting them based on their posterior probabilities. In this work, it uses trainable linear classifiers on frozen foundation model features to determine which model heads and features are best for a specific dataset, providing a principled way to ensemble diverse models.
Optimizable Model Averaging (OMA): OMA is a black-box approach that learns scalar weights for each candidate model. It optimizes an objective function that minimizes the expected entropy of predictions under the ensemble. This allows the system to dynamically adjust model contributions based on data, improving performance beyond simple output averaging.
Predictive Entropy Minimization: This is the core objective function for OMA. Minimizing predictive entropy aims to find an ensemble where the combined prediction distribution is as concentrated as possible around the true label. The authors note that this regularization helps control adaptation and improves calibration compared to methods that only focus on raw accuracy.
Regularization ($\lambda$): Regularization is used in OMA to control how much the ensemble adapts to the specific dataset. It is added as a penalty term during optimization, often using prior beliefs from zero-shot or MAP models. Tuning this parameter helps balance model diversity with data-specific performance, leading to better calibration results.

Terminology

Foundation Model — 5× in this paper · explained in 3 episodes
Foundation Models — 5× in this paper · explained in 7 episodes

Summary

The gist The authors revisit classical Bayesian model averaging (BMA) to ensemble pre-trained and/or lightly-finetuned foundation models by introducing trainable linear classifiers and an optimizable model averaging scheme (OMA) to enhance classification performance on image and text data.

Bayesian Model Averaging (BMA)

The BMA paradigm is revisited 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, they 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 In vision tasks, where these pretrained models provide rich and diverse features for natural images, their approach outperforms simple output averaging and is comparable to averaging fine-tuned CILP models’ parameters The posterior predictive distribution on a query point x is given by p(y∗x∗, D) = X L l=1 p(y∗x∗, Ml, D)p(MlD), which is the average of the posterior distributions under each of the models considered weighted by their posterior model probability The marginal likelihood of model Ml is obtained by integrating out the model parameters wl: p(DMl) = R p(Dwl, Ml)p(wlMl)dwl In Laplace approximation, the unnormalized log-posterior given by Ψ(wl) is approximated as a quadratic function in wl using the 2nd-order Taylor expansion The approximate marginal likelihood in eq. 7 gives us an un-normalized posterior model probability if we assume a uniform distribution for p(Ml) = 1 L, and we normalize it by dividing the sum p(MlD) = p(DMl)PL l′=1 p(DMl′)

Optimizable Model Averaging (OMA)

When there is a large distribution shift between the training and validation/test data distributions, they introduce an optimizable model averaging (OMA) method OMA improves classification by directly optimizing the model ensemble weights (just like those weights based on model posterior distributions in BMA) to reduce the expected entropy of predictions under ensembled models, significantly outperforming model output averaging The objective function for OMA is L(βl) = −1/M X M m EPL l=1 βl·p(y∗mx∗m,Ml) log "X L l=1 βl· p(y∗m = cx∗m, Ml)# This objective function can be regularized by adding a regularization term, e.g., taking zeroshot models’ or MAP estimates’ performance as a form of prior belief about each model, β0 l, with a regularization constant λ ≥ 0 The constraint is that PL lβl = 1 and βl ≥ 0 for all l The predictive distribution on an unseen point by p(y∗x∗, D) is computed as X L l=1 p(y∗x∗, M˜l, D)βl

Experimental Results

In image classification experiments on Imagenet-1K and its variant datasets, BMA improves output averaging by a large margin in the in-distribution Img-1K validation data The zero-shot models using our OMA weights improve the performance of output averaging, up to approximately 20% (the column, called Cam17 and the first row) When using both MAP and zeroshot models together with our OMA paradigm improves the performance The results are not due to random chance because across three training subsets, the accuracy of our method does not change by much For text classification tasks using transformer-based models, OMA outperforms output averaging and a single best model The OMA weights learned by Algorithm 2 show that when the models with high accuracy typically get higher weights, and vice versa

Comparison to Other Methods

Model ensembling methods like their own require more inference time than a weight-averaged single model, however, their framework requires significantly less total computation In Model Soups, in addition to the required computation time for fine-tuning 58 models, the greedy soup approach requires the inference-like computation (passing a held-out validation set through all those 58 candidate models) to select models to be included in the soup The authors show that BMA improves output averaging by a large margin in the in-distribution Img-1K validation data

Limitations and Future Directions

There are limitations to the proposed methods, first it is hard to employ our BMA framework beyond the image domain where they do not know how useful the features of the pre-trained foundation models are Alternatively, one could use OMA while OMA has its limitations as stated at the end of Subsec. 3.2 The ultimate goal is to develop computationally efficient methods by utilizing pre-trained or lightly fine-tuned models, reducing the need for additional energy-intensive training This approach will enable the incorporation of future, possibly significantly better foundation models to enhance the performance of challenging classification tasks

Implementation Details

The implementation involves two main folders, Image Classification and Text Classification In Image Classification, scripts include preprocess.py, reorganized imagenet feats.py, reorganized imagenet OOD datasets feats.py, training with preprocessed data.py, hessian computation with p 1minus p.py, predict given zeroshot individually trained weights.py and test optimization.py In Text Classification, the base code used is from https://github.com/huggingface/transformers/tree/main/examples/pytorch/text-classification The text classification experiments use four models: two BERT-base cased models (large and base) and two funnel-transformers (xlarge and xlarge-base) The performance of the OMA weights learned by Algorithm 2 is visualized in Fig. 3

Hyperparameter Tuning

The ablation study for BMA shows that the optimal prior variance α for ImageNet-1K was 80 when using a grid with values of 0.01, 0.1, 10, 40, 80, and 160 For OMA on ObjectNet when using zeroshot models, the optimal λ was 1 The ablation study for OMA shows that the optimal λ value for ObjectNet when using zeroshot models was 10,000

Model Comparisons

The performance of individual models for ImageNet-1K and its variant datasets is shown in Table 4 The MAP estimates trained with ImageNet-1K outperform zeroshot models on most datasets, but for datasets such as ImageNet-A, ImageNet-R, ImageNet-Sketch, and Objectnet, zeroshot models consistently outperform trained classifiers The performance of BMA and OMA when using both MAP and zeroshot models together with OMA paradigm improves the performance on Camelyon17, Flowers102, and Sun397 datasets

Acknowledgements

The authors thank their anonymous reviewers for their constructive feedback M.

Improvements for AI systems

The system can achieve superior classification performance on image datasets by employing an optimizable model averaging scheme (OMA) that optimizes ensemble weights to reduce the expected entropy of predictions under ensembled models, significantly outperforming model output averaging.
The system can leverage pre-trained foundation models with only minimal fine-tuning, which is a practical and sustainable choice compared to fine-tuning large models for each target dataset due to its reduced high CO2 footprint.
The system can be robust against distribution shifts by using model posteriors or optimized weights; specifically, the OMA method is proposed when there is a large distribution shift between training and validation/test sets, as it directly optimizes weights to reduce prediction surprise.
For text classification, the system can utilize OMA applied to transformer-based models by fine-tuning each model for small epochs like 3, using the learned model weights which are shown to be similarly with each model’s performance in Figure 3.
The system can outperform baseline methods like output averaging and Model Soups on out-of-distribution image classification datasets (e.g., ImageNet-R), achieving performance improvements up to 2.4% over output averaging when using zeroshot models with OMA weights.

Abstract

Foundation model ensembles can improve accuracy and robustness, but full fine-tuning and parameter-space merging are often expensive or impossible when candidate models have different architectures. We propose Optimizable Model Averaging (OMA), a black-box approach that learns nonnegative model-combination weights, one scalar coefficient per model, to average the candidate models' predictive probabilities, leaving all model parameters frozen, by minimizing regularized predictive entropy. A central challenge is that minimizing predictive entropy can favour confidently incorrect models: increased confidence need not imply improved accuracy or calibration. We therefore examine regularization as a mechanism for controlling adaptation and empirically evaluate its effect on calibration, observing lower expected calibration error in selected configurations, alongside better performance-calibration trade-offs than model output averaging. As foundation models continue to advance, these approaches will enable the incorporation of future, potentially much better models to improve performance on challenging text and image classification tasks.