Be More Diverse than the Most Diverse: Optimal Mixtures of Generative Models via Mixture-UCB Bandit Algorithms

2024-12-23

AI Summary

This work addresses generative model ensembling—improving upon the best individual pre-trained model in both sample quality and diversity via weighted combination. We formulate mixture weight optimization as a kernelized convex quadratic program, with Kernel Inception Distance (KID) and Rényi Kernel Entropy (RKE) as objective metrics. To solve it efficiently, we propose Mixture-UCB, a multi-armed bandit algorithm based on Upper Confidence Bound principles, which provably converges to the global optimum with low sample complexity and a theoretically bounded regret. Empirically, our method achieves significant reductions in FID/KID scores while improving the fidelity-diversity trade-off across standard image and text generation benchmarks. The implementation is publicly available.

Technology Category

Machine Learning: Deep Generative Models & AutoencodersSearch and Optimization: Mixed Discrete/Continuous SearchNatural Language Processing: Generation

Application Category

User Modeling, Personalization and Recommendation: Fairness-aware retrieval and rankingSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

The availability of multiple training algorithms and architectures for generative models requires a selection mechanism to form a single model over a group of well-trained generation models. The selection task is commonly addressed by identifying the model that maximizes an evaluation score based on the diversity and quality of the generated data. However, such a best-model identification approach overlooks the possibility that a mixture of available models can outperform each individual model. In this work, we numerically show that a mixture of generative models on benchmark image datasets can indeed achieve a better evaluation score (based on FID and KID scores), compared to the individual models. This observation motivates the development of efficient algorithms for selecting the optimal mixture of the models. To address this, we formulate a quadratic optimization problem to find an optimal mixture model achieving the maximum of kernel-based evaluation scores including kernel inception distance (KID) and R'enyi kernel entropy (RKE). To identify the optimal mixture of the models using the fewest possible sample queries, we view the selection task as a multi-armed bandit (MAB) problem and propose the Mixture Upper Confidence Bound (Mixture-UCB) algorithm that provably converges to the optimal mixture of the involved models. More broadly, the proposed Mixture-UCB can be extended to optimize every convex quadratic function of the mixture weights in a general MAB setting. We prove a regret bound for the Mixture-UCB algorithm and perform several numerical experiments to show the success of Mixture-UCB in finding the optimal mixture of text and image generative models. The project code is available at https://github.com/Rezaei-Parham/Mixture-UCB.

Problem

• Selecting optimal mixture of generative models for better performance.
• Maximizing evaluation scores like FID and KID via mixtures.
• Efficiently finding best model mixtures using Mixture-UCB algorithm.

Innovation

• Optimal mixture of generative models via quadratic optimization
• Mixture-UCB bandit algorithm for efficient model selection
• Convergence to optimal mixture with provable regret bounds