Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models

Abstract: Conditional generative models have emerged as powerful tools for sampling from target conditional distributions, driving substantial advances across a wide range of scientific and applied domains. As these models proliferate, practitioners often face multiple plausible generators whose performance can vary with the task, data, or input condition. We propose an optimal model averaging framework for conditional generative models, allowing candidate generators to be combined even when they are accessible only through conditional samples without tractable densities. Specifically, we use a sample-based maximum mean discrepancy between conditional distributions, which first leads to a static model averaging method, StaticMA, assigning fixed weights to different candidates. In addition, we develop MoEMA (mixture-of-experts model averaging), an input-adaptive method that parameterizes covariate-dependent weights through a softmax neural-network gate. We establish in-sample and out-of-sample asymptotic optimality for the proposed methods, together with consistency of the estimated adaptive weight function under regularity conditions. The framework applies directly to Euclidean responses and extends to unstructured data by combining our formulation with fixed representation maps. Across a broad set of simulations and real-data studies spanning tabular, image, and text modalities, MoEMA generally improves over competing baselines, demonstrating the effectiveness of our proposed methods.

• The paper introduces a novel sample-based aggregation framework that uses conditional MMD to optimally average arbitrary black-box conditional generative models.
• It presents two strategies—StaticMA and MoEMA—with MoEMA employing a softmax-gated neural network for input-adaptive weighting, achieving superior performance across modalities.
• Extensive theoretical and empirical evaluations demonstrate both in-sample and out-of-sample risk optimality, with significant improvements over established baselines.

Optimal Aggregation of Conditional Generative Models via Mixture-of-Experts Model Averaging

Motivation and Problem Formulation

Contemporary advances in conditional generative modeling have resulted in practitioners routinely facing multiple plausible generator candidates for a given task, especially in high-stakes settings where marginal improvements in generation quality are valuable. Given that each conditional generative model is typically accessible as a black-box sample generator with potentially intractable or unknown likelihoods, conventional model selection and aggregation techniques based on likelihood or explicit scoring rules are often infeasible. This paper proposes and rigorously studies an optimal model averaging framework for conditional generative models, focusing on sample-based aggregation strategies that can leverage candidate generators—regardless of their internal form—by using only conditional samples.

The aggregation paradigm is formalized through the Maximum Mean Discrepancy (MMD) as the sample-based conditional distribution discrepancy. The key object is the conditional MMD (cMMD), which enables distributional risk assessment and model combination directly from samples, thus bypassing likelihood computations.

Methodological Contributions

The paper introduces two main aggregation strategies based on the cMMD criterion:

• StaticMA: A static model averaging method which assigns fixed, global weights to each model in the pool. Estimation of the optimal weights is formulated as a quadratic programming problem over the probability simplex, minimizing the empirical cMMD between the target distribution and the mixture.
• MoEMA: An input-adaptive Mixture-of-Experts Model Averaging method, where the weighting function is a softmax-gated neural network, hence enabling locally varying mixture weights as a function of the input/covariate. This architecture draws on classical MoE frameworks but generalizes them to arbitrary conditional samplers as experts.

Both methods rely only on conditional samples from each generator, and do not require knowledge of their closed-form densities. For unstructured response spaces (e.g., images, sequences), the methodology is extended to operate in feature spaces via pretrained representation maps, maintaining the flexibility of the approach for practical, high-dimensional tasks.

Theoretical Results

A central technical contribution is the development of asymptotic theory for both in-sample and out-of-sample risk optimality:

• Both StaticMA and MoEMA are shown to achieve asymptotic in-sample and out-of-sample optimality: the empirical risk of the learned aggregation converges to that of an oracle with access to population-optimal weights.
• Consistency of the estimated weighting function itself is established for MoEMA, ensuring that the adaptive gate identifies locally optimal generator relevance.
• The analysis explicitly handles complexities absent from classical frequentist model averaging, including the infinite-dimensionality of the weight function, neural network approximation error (for the adaptive setting), and algorithmic stability for out-of-sample generalization.

These results generalize optimal frequentist model averaging theory to the setting of conditional distributions evaluated under sample-based integral probability metrics, and further, to functionally adaptive weighting.

Empirical Studies

Extensive simulation and real-data experiments are conducted across tabular, image, and text modalities, benchmarking StaticMA and MoEMA against competing methods and oracle single-generator baselines.

Following the introduction and empirical evaluation of the key simulation design—(Figure 1)—MoEMA and StaticMA consistently outperform both equal-weight and per-metric best single generator baselines in terms of cMMD and task-specific accuracy metrics.

Figure 1: Simulation results across sample sizes. Rows report $\mathrm{MMD}^2$, MeanAE, and QuantileLoss; columns show dimension settings; solid/dashed lines denote architecture/training-data heterogeneity.

Among the salient empirical findings are:

• StaticMA and MoEMA demonstrate strong numerical improvements over established baselines: In simulated Gaussian mixture designs, relative reductions in conditional MMD and accuracy metrics (MeanAE, QuantileLoss) due to MoEMA span 1–75%, with the largest gains when generator pools are diverse in local support.
• In real-data tabular regression (UCI CT Slices, Protein), simple averaging and static model averaging prove robust, but MoEMA achieves pronounced gains (10–68% metric reductions) in scenarios with substantial local generator heterogeneity.
• For image generation (MNIST, ImageNet-100) and inpainting (MNIST, CIFAR-10), adaptive weighting in MoEMA further improves upon static rules, particularly when generator specialization is present (e.g., per-class training).
• In text continuation tasks using autoregressive LLMs from multiple source domains, both static and adaptive sample-based model averaging deliver significant perplexity reductions over best single-generator selection, highlighting the approach’s extensibility to generative transformers.

Contradictory Claims and Methodological Guarantees

A strong claim, supported by both theory and experiment, is that MoEMA with input-dependent gating surpasses not only all static aggregation baselines but also the best performing single generator across conditional input regimes—this is especially pronounced in scenarios with input-conditioned generator specialization. Furthermore, the theoretical guarantees ensure that, under mild regularity and approximation conditions, risk-optimality and weight consistency persist even as the complexity of the generator pool and the input distribution increase.

Practical and Theoretical Implications

Practically, the framework enables rigorous, agnostic combination of conditional generative models across a wide variety of response types and domains, including those where explicit likelihoods are either unavailable or unreliable. In contexts such as generative design, scientific simulation, or data augmentation, this approach provides a principled means to combine and specialize pretrained or independently-trained generative models without retraining or access to their internals.

Theoretically, the paper sets a new standard for model averaging in conditional generative modeling by:

• Extending optimal FMA from scalar point estimation to full conditional distribution approximation under an IPM (here, cMMD);
• Enabling weight adaptivity via high-capacity function approximators (deep neural networks with softmax gating), while maintaining statistical consistency, and
• Providing the first comprehensive analysis of model averaging for arbitrary conditional generative samplers accessible only via samples.

Future Directions

Natural continuations of this line of work include generalizing the theoretical framework to dependent data (notably, time series), extending cMMD-based aggregation to more general IPMs and joint distribution criteria, and exploring scalable, computation-efficient implementations for high-dimensional or autoregressively structured response spaces. There is also scope for synthesizing this approach with advanced model-composition strategies in the context of large, sparse MoE architectures for deep learning.

Conclusion

This paper establishes a mathematically rigorous, empirically validated approach for optimal model averaging of conditional generative models. Through the integration of sample-based discrepancy minimization and mixture-of-experts gating, the proposed MoEMA method consistently and significantly improves upon both static model averaging and best single-generator baselines across modalities and domains. The theoretical and practical contributions considerably broaden the methodological toolkit for generative model ensemble construction and evaluation, and the sample-based nature of the proposed criteria ensures broad applicability and independence from model-specific likelihood access or architecture constraints.