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

Baihua He; Shijin Gong; Xinyu Zhang

REVIEW 2 major objections 5 minor 14 references

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

T0 review · 2 major / 5 minor · reviewed 2026-07-11 · grok-4.5

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Conditional generators can be optimally combined from samples alone, with fixed or input-adaptive weights, and the data-driven mixtures match the oracle risk.

desk verdict Solid FMA-for-samplers paper: sample-based cMMD averaging with fixed and neural adaptive weights, real optimality theorems, and multi-modality experiments that mostly favor MoEMA. read the letter →

arxiv 2607.04360 v1 pith:JA3YBKFB submitted 2026-07-05 stat.ML cs.LG

Shijin Gong , Baihua He , Xinyu Zhang This is my paper

classification stat.MLcs.LG MSC 62G0562G2068T07

keywords conditionalgenerativemodelfrequentistaveragingmixture-of-expertsmaximummeandiscrepancyasymptoticoptimalityStaticMAMoEMA

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Practitioners now face many conditional generators that produce samples of a response given an input, but often lack closed-form densities and no single model wins everywhere. This paper shows how to combine them optimally using only generated samples, by measuring discrepancy with a sample-based conditional maximum mean discrepancy. The result is StaticMA, which assigns one fixed weight vector, and MoEMA, which routes each input through a neural-network softmax gate so different generators dominate in different regions. Under regularity conditions, both methods are asymptotically optimal in-sample and out-of-sample, and MoEMA’s learned gate recovers the population-optimal weights. The same criterion works for Euclidean responses and, with fixed feature maps, for images and text. Empirically, MoEMA improves over the best single generator and equal-weight averaging across tabular, image, and short-text tasks.

What carries the argument

Conditional MMD (cMMD): average, over the covariate distribution, of the maximum mean discrepancy between the target and mixture conditional laws, estimated from observed responses and generated samples via kernel mean embeddings; StaticMA optimises fixed simplex weights by quadratic programming, while MoEMA parameterises covariate-dependent weights by a softmax neural-network gate.

What would settle it

On a synthetic mixture where one generator is clearly best in each region of the covariate space, check whether MoEMA’s learned gate recovers those regions and whether its test cMMD approaches the oracle mixture as sample size grows; failure of either would contradict the claimed optimality and weight consistency.

Watch

Extended reading notes

Core claim

When candidate conditional generators are accessible only through samples, weights that minimise an empirical conditional MMD yield mixtures whose in-sample and out-of-sample risks are asymptotically equivalent to those of the infeasible oracle weights; for the input-adaptive MoEMA gate the estimated weight function is also consistent for the population-optimal weight set.

Load-bearing premise

Every candidate generator must converge uniformly, in the same discrepancy, to some fixed limiting conditional law at a rate that is negligible relative to residual mixture risk; if generators stay far from any such limit, the optimality and consistency theorems no longer apply.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

This is a clean methods-plus-theory paper that does something useful: it averages conditional generators when you only have samples, not densities. StaticMA is fixed weights via quadratic programming on empirical cMMD; MoEMA puts a softmax NN gate on the same objective so weights can depend on x. That combination—FMA + MoE + sample MMD for black-box conditional generators—is the real contribution, not any single ingredient.

What they do well is the theory package. Theorems 1–6 give in-sample and out-of-sample asymptotic optimality for both methods, plus consistency of the adaptive weight function, under explicit rates, kernel boundedness, generator approximation, algorithmic stability, and Hölder gates. Parallel StaticMA results under weaker conditions. The criterion is standard and identifiable under the usual RKHS conditions; the mixture is a proper mixture of samplers, not a redefinition of the target. Empirically they cover tabular, MNIST/ImageNet/CIFAR inpainting, and multi-domain GPT-2 continuation, with MoEMA usually beating BestSingle, equal weights, and StaticMA, especially when candidates are region specialists.

Soft spots are real but proportionate. Assumption 1 (uniform generator MMD rate negligible relative to oracle risk) and the stability assumptions for out-of-sample results are load-bearing; if generators stay far from any limit, the optimality claims do not apply. Hölder smoothness of the true gate is the usual nonparametric price for MoEMA. Proofs and code are deferred; BestSingle is an oracle per-metric baseline, so the tables slightly overstate the practical gap. Free knobs (kernel, Ng, gate size, representation maps) are standard and not hidden. None of this breaks the central argument.

This is for people who already care about model averaging, conditional generation, or combining pretrained samplers without likelihoods. It is not a reorganization of the field, but it is a serious, citable toolkit paper with honest literature engagement. I would send it to peer review without hesitation.

Referee Report

2 major / 5 minor

Summary. The paper develops an optimal frequentist model-averaging framework for conditional generative models that are accessible only through samples. Using a sample-based conditional maximum mean discrepancy (cMMD) criterion, it proposes StaticMA (fixed simplex weights, solved by quadratic programming) and MoEMA (input-adaptive weights via a softmax neural-network gate). Under stated regularity conditions, Theorems 1–2 and 4–5 establish in-sample and out-of-sample asymptotic optimality of the data-driven weights relative to the corresponding oracle risks, and Theorem 6 establishes consistency of the estimated adaptive weight function for the population-optimal weight set. The framework extends to unstructured responses via fixed representation maps and is evaluated on simulations and real-data tasks spanning tabular, image, and text modalities, where MoEMA generally improves over best-single, equal-weight, and static baselines.

Significance. If the results hold under the stated conditions, the paper supplies the first optimal model-averaging theory for conditional generative models and the first such theory built around a sample-based MMD criterion for comparing conditional distributions. This is a genuine extension of classical FMA along three axes: from point prediction to conditional distributions, from likelihood-based criteria to sample-based discrepancies, and from constant weights to input-adaptive gates. The methods are architecture-agnostic (requiring only conditional samples), cover Euclidean and unstructured responses, and are accompanied by explicit asymptotic guarantees (Theorems 1–6) together with multi-modality empirical support. These contributions are of clear interest to both the statistical model-averaging literature and practitioners who face multiple pretrained or trained generators.

major comments (2)

Assumption 1 (uniform MMD convergence of every fitted generator to a limiting generator at rate a_n) together with the rate conditions in Assumptions 3 and 7 is load-bearing for Theorems 1–6. The manuscript should state more explicitly when this is expected to hold for modern generators (diffusion, flows, autoregressive LMs) that are often used as fixed pretrained black boxes, and should clarify the practical regime in which a_n is negligible relative to residual mixture risk ξ_n / ξ_Xn. Without that discussion the theorems apply cleanly only when generators are themselves consistent, which is not always the operating regime of the real-data experiments.
Out-of-sample optimality (Theorems 2 and 5) rests on uniform algorithmic stability of generator training (Assumptions 4 and 8). For pretrained generators held fixed, stability is vacuous; for generators trained on D_n the paper should either verify or cite that the concrete training procedures used in Sections 6–7 satisfy the required β_n rates, or clearly separate the theory into “fixed-pool” and “jointly-trained” cases so that the out-of-sample claims are not overstated for the pretrained setting.

minor comments (5)

Figure 1 normalises all metrics by StaticMA; the absolute levels (or an un-normalised appendix table) would make the magnitude of gains easier to interpret, especially when StaticMA itself is weak.
In Section 4.2 the Monte Carlo cost O(M² N_g²) for the Ĉ terms is noted; a short practical recommendation for choosing N_g (or a default used in the experiments) would help reproducibility.
Notation for the empirical criterion L_n(w,Q) and the in-/out-of-sample risks IcMMD² / OcMMD² is dense; a short notational glossary or a display of the population objective before the empirical version would improve readability.
Table 1 Panel A (Protein) shows essentially no gain for MoEMA over BestSingle/StaticMA; a brief remark on when complementary generators are absent would temper the “generally improves” claim.
Cross-references to Supplementary Algorithms 1–2 and Sections C–G are frequent; ensuring that the main text is self-contained for the core StaticMA/MoEMA definitions would help readers who do not immediately open the supplement.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: asymptotic optimality is standard FMA risk-equivalence relative to an independently defined cMMD target, not a fit renamed as prediction.

full rationale

The derivation chain is self-contained and non-circular. Weights for StaticMA and MoEMA minimize the empirical sample-based cMMD criterion Ln(w, Q̂) between the mixture of black-box generators and observed responses; Theorems 1–2 and 4–5 then show that the resulting data-driven weights are asymptotically equivalent, in in-sample and out-of-sample cMMD risk, to the infeasible oracle that minimizes the corresponding population risk over the same weight class, and Theorem 6 adds consistency of the adaptive gate to the population-optimal weight set. That structure is classical frequentist model averaging (risk-minimizing weights + oracle risk equivalence), not a tautology: the target conditional law P(·|x) is external to the candidate pool, generators may be pretrained or trained separately and need only produce samples, and cMMD is an independently defined IPM/RKHS discrepancy (not defined in terms of the learned weights). Empirical comparisons use held-out metrics (MMD², MeanAE, QuantileLoss, FID/KID, energy, PPL) against BestSingle, SimpleAvg, and StaticMA; BestSingle is an oracle per-metric selection baseline, not an input to the weight fit. Self-citations to prior FMA work by overlapping authors supply background on model averaging, not a load-bearing uniqueness theorem that forces the present claims. Load-bearing assumptions (uniform generator approximation rate a_n, algorithmic stability, Hölder gate smoothness) are stated openly and are rate conditions, not definitional closures. No step reduces a claimed prediction or first-principles result to its own fitted inputs by construction.

Assumptions & free parameters 4 free parameters · 6 assumptions · 3 invented entities

The central optimality claims rest on standard IPM/MMD and RKHS facts, classical uniform stability, neural approximation of Hölder gates, and several rate conditions that couple sample size, number of generators, Monte Carlo size, and oracle risk. Free knobs (kernel, Ng, gate architecture) affect finite-sample behavior; invented objects are the named averaging procedures and the cMMD risk used for weight choice, not new physical entities.

free parameters (4)

kernel bandwidth / kernel choice on Y (or feature space Z)
MMD depends on the kernel; bandwidth and kernel family are tuning parameters selected by cross-validation on the fitting criterion, not derived from first principles.
Monte Carlo generator sample size Ng
Empirical B and C terms use Ng samples per condition; larger Ng reduces variance at computational cost and is chosen as computationally feasible.
MoEMA gate network depth D and width W
Theory specifies asymptotic schedules in terms of Hölder index and manifold dimension; practice uses finite networks selected among hyperparameters, which determine approximation capacity of Wn.
representation maps ψ (response) and φ (condition) for unstructured data
Feature-space aggregation composes fixed pretrained maps (e.g., CLIP, GPT-2 features) with the kernel; choice of map changes the discrepancy being optimized.

assumptions (6)

standard math cMMD is well-defined and identifies conditional distributions P_X-a.e. under kernel/characteristic conditions (Supp. A.2).
Standard MMD/IPM theory specialized to conditional averaging; invoked to justify the risk criterion in Section 3.
domain assumption Assumption 1: each fitted generator converges uniformly in MMD to a limiting generator at rate a_n.
Load-bearing for all optimality theorems; generators are treated as approximately fixed limits after training.
standard math Kernel k is measurable and uniformly bounded (Assumption 2).
Ensures MMD/cMMD are well-defined; satisfied by Gaussian/Laplace kernels.
domain assumption Uniform stability of generator training algorithms with rate β_n (Definition 2; Assumptions 4 and 8).
Required for out-of-sample optimality when generators are fit on Dn; cites Bousquet & Elisseeff (2002).
domain assumption Population score/gate coordinates lie in a Hölder class G_κ with uniform bound (Assumption 5); covariates on a compact low-dimensional manifold (Assumption 6).
Needed for neural approximation rates and MoEMA Theorems 4–6.
ad hoc to paper Rate conditions linking n, M, Ng, N, an, network size S, and oracle risks ξ_n / ξ_Xn (Assumptions 3, 7).
Paper-specific asymptotic regimes under which empirical cMMD is uniformly consistent and data-driven weights are optimal.

invented entities (3)

StaticMA independent evidence
purpose: Fixed-weight averaging of conditional generators by minimizing empirical cMMD over the simplex.
Named method introduced in Section 4.2; evaluable by re-running the quadratic program on samples.
MoEMA independent evidence
purpose: Input-adaptive mixture-of-experts averaging with softmax neural gate trained on the same sample-based cMMD objective.
Core proposed method; falsifiable via held-out distributional metrics as in Sections 6–7.
conditional MMD (cMMD) risk for generator mixtures independent evidence
purpose: Scalar risk obtained by averaging MMD² between target and mixture conditionals under P_X, used as the FMA criterion without densities.
Defined in Eq. (1) and empirical Ln; related to averaged MMD of Huang et al. (2022) but specialized here to weight selection.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models." pith.science (2026). https://pith.science/paper/JA3YBKFB

@misc{pith2026260704360,
  author       = {Pith},
  title        = {Pith review of: Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/JA3YBKFB}},
  note         = {Machine review of arXiv:2607.04360}
}

read the original 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.

Figures

Figures reproduced from arXiv: 2607.04360 by the authors.

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

14 extracted references · 12 linked inside Pith

[1]

Balaji, S

Y. Balaji, S. Nah, X. Huang, A. Vahdat, J. Song, Q. Zhang, K. Kreis, M. Aittala, T. Aila, S. Laine, et al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324,

arXiv
[2]

W. Cai, J. Jiang, F. Wang, J. Tang, S. Kim, and J. Huang. A survey on mixture of experts. arXiv preprint arXiv:2407.06204,

arXiv
[3]

Chang, Z

J. Chang, Z. Ding, Y. Jiao, R. Li, and J. Z. Yang. Deep conditional distribution learning via conditional f¨ ollmer flow.arXiv preprint arXiv:2402.01460,

arXiv
[4]

Z. Gao, B. He, and Y. Yang. Combining pre-trained models via localized model averaging. arXiv preprint arXiv:2605.13421,

arXiv
[5]

Y. Gu, L. Kong, and D. Xia. Prediction-powered inference by mixture of experts.arXiv preprint arXiv:2604.27892,

arXiv
[6]

Han and M

X. Han and M. Zhou. Diffusion boosted trees.arXiv preprint arXiv:2406.01813,

arXiv
[7]

28 X. Han, H. Zheng, and M. Zhou. Card: Classification and regression diffusion models.arXiv preprint arXiv:2206.07275,

arXiv
[8]

Huang, H

Z. Huang, H. Lam, and H. Zhang. Evaluating aleatoric uncertainty via conditional generative models.arXiv preprint arXiv:2206.04287,

arXiv

Show all 14 references

[9]

Kotelnikov, D

A. Kotelnikov, D. Baranchuk, I. Rubachev, and A. Babenko. Tabddpm: Modelling tabular data with diffusion models.arXiv preprint arXiv:2209.15421,

arXiv
[10]

Madhusudhanan, V

K. Madhusudhanan, V. K. Yalavarthi, J. Sonntag, M. Stubbemann, and L. Schmidt-Thieme. Tabresflow: A normalizing spline flow model for probabilistic univariate tabular regression. arXiv preprint arXiv:2508.17056,

arXiv
[11]

Mirza and S

M. Mirza and S. Osindero. Conditional generative adversarial nets.arXiv preprint arXiv:1411.1784,

arXiv
[12]

Gpt-4 technical report.arXiv preprint arXiv:2303.08774,

OpenAI. Gpt-4 technical report.arXiv preprint arXiv:2303.08774,

arXiv
[13]

Rezaei, F

P. Rezaei, F. Farnia, and C. T. Li. Be more diverse than the most diverse: Optimal mixtures of generative models via mixture-ucb bandit algorithms. InInternational Conference on Learning Representations, volume 2025, pages 17446–17474,

2025
[14]

Shazeer, A

N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. V. Le, G. E. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.arXiv preprint arXiv:1701.06538,

arXiv

Pith tools

Reviewed July 11, 2026 · model on record in the stance chip above.

[1] [1]

Balaji, S

Y. Balaji, S. Nah, X. Huang, A. Vahdat, J. Song, Q. Zhang, K. Kreis, M. Aittala, T. Aila, S. Laine, et al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324,

arXiv

[2] [2]

W. Cai, J. Jiang, F. Wang, J. Tang, S. Kim, and J. Huang. A survey on mixture of experts. arXiv preprint arXiv:2407.06204,

arXiv

[3] [3]

Chang, Z

J. Chang, Z. Ding, Y. Jiao, R. Li, and J. Z. Yang. Deep conditional distribution learning via conditional f¨ ollmer flow.arXiv preprint arXiv:2402.01460,

arXiv

[4] [4]

Z. Gao, B. He, and Y. Yang. Combining pre-trained models via localized model averaging. arXiv preprint arXiv:2605.13421,

arXiv

[5] [5]

Y. Gu, L. Kong, and D. Xia. Prediction-powered inference by mixture of experts.arXiv preprint arXiv:2604.27892,

arXiv

[6] [6]

Han and M

X. Han and M. Zhou. Diffusion boosted trees.arXiv preprint arXiv:2406.01813,

arXiv

[7] [7]

28 X. Han, H. Zheng, and M. Zhou. Card: Classification and regression diffusion models.arXiv preprint arXiv:2206.07275,

arXiv

[8] [8]

Huang, H

Z. Huang, H. Lam, and H. Zhang. Evaluating aleatoric uncertainty via conditional generative models.arXiv preprint arXiv:2206.04287,

arXiv