Stacking-Based Weighting for Large Language Bayes in M-Open Settings

AI summary

Overview

Research area: Probabilistic machine learning — specifically Bayesian model combination and automated probabilistic programming built on large language models.

Technical level: Intermediate. Readers need comfort with Bayesian model averaging, variational inference (ELBO), and leave-one-out predictive scoring, but the paper's argument is accessible without following the algebra.

Scope: The paper replaces the evidence-based model-weighting rule in Large Language Bayes with a KL-regularized stacking objective computed from ELBO surrogates, and evaluates it on three county-level disaster-loss prediction tasks.

What This Paper Is About

Large Language Bayes (LLB) asks a language model to write down candidate probabilistic programs from a plain-text problem description, runs inference inside each program, and blends their predictions with Bayesian model averaging. The problem is that the language model has no guarantee of producing the true data-generating process, so the task lives in the "M-open" setting, where likelihood-based weights tend to collapse onto a single candidate model. The authors ask whether weights chosen for predictive performance (stacking) instead of marginal likelihood can keep the ensemble spread across several plausible models without hurting accuracy.

Key Contributions

1. A stacking-based aggregation rule for LLB. The paper swaps LLB's Bayesian model averaging weights, which are proportional to each model's marginal likelihood, for weights chosen to maximize an aggregate leave-one-out log predictive score over the candidate models.

2. A tractable ELBO-based objective for latent-variable models. Because leave-one-out predictive densities are intractable when models contain latent variables, the authors prove (Proposition 4.1) that a variational ELBO lower-bounds each log leave-one-out predictive density, and substitute exp(ELBO) into the stacking objective as a surrogate.

3. A KL regularizer toward the BMA weights. The objective adds a term −λ KL(w^BMA ‖ w) that anchors the learned weights near the original evidence-based weights while still permitting redistribution when predictive performance supports it.

4. An empirical comparison on real hazard-loss tasks. The method is compared against the original LLB aggregation rule on wildfire, inland flooding, and hurricane county-level loss prediction, measuring both weight concentration (effective sample size, ESS) and held-out expected log predictive density (ELPD).

Main Findings

• BMA collapses onto one model; stacking does not. In all three tasks the BMA weights put essentially all mass on a single candidate, while stacking assigns substantial weight to at least two models. Reported top weights: Hurricane — BMA w_139 = 1.000 versus stacking w_1 = 0.759, w_65 = 0.238, w_129 = 0.003; Inland Flood — BMA w_24 = 1.000 versus stacking w_24 = 0.588, w_13 = 0.412; Wildfire — BMA w_83 = 0.998 versus stacking w_138 = 0.666, w_83 = 0.334.

• Predictive performance is essentially unchanged. Held-out ELPD: Hurricane −5.983 (BMA) versus −5.999 (stacking); Inland Flood −370.340 for both; Wildfire −177.922 (BMA) versus −177.897 (stacking). The authors read this as evidence that dispersing weight does not cost accuracy.

• The difference is in how model uncertainty is represented, not in raw predictive power. The paper explicitly frames the contribution this way: the main effect of stacking is a more faithful spread of weight under misspecification rather than better predictions.

• The variational surrogate converges quickly. On the diagnostic plots (6 sampled data points per task), ELBO values converge within 25 iterations, typically stabilizing after 10–15 iterations, supporting the use of these surrogates in the stacking objective.

• Weight concentration is measured but not tabulated numerically. ESS, defined as 1 / Σ_k w_k², is reported graphically in the paper's Figure 2; specific ESS values are not given as numbers in the text.

Methodology in Plain English

The pipeline keeps LLB's generation and inference steps untouched and only changes the final blending step.

1. Generate candidate models. A text description of the problem is fed to Qwen2.5 Coder 32B, which produces 250 candidate NumPyro programs. Programs are deduplicated by normalized code hash and filtered to keep only those that compile and pass shape validation.

2. Fit each model. Every valid candidate is fit with the No-U-Turn Sampler (NUTS), a gradient-based MCMC method, producing posterior samples.

3. Score each model on held-out-style predictions. For every observation x_i and every model, the data point is set aside, a leave-one-out dataset is formed, and a Gaussian variational approximation is fit to the resulting posterior. The converged ELBO is used as a stand-in for the (intractable) log leave-one-out predictive density of that point under that model.

4. Learn the weights. Weights are constrained to the probability simplex and chosen to maximize the summed log of the weighted mixture of exp(ELBO) terms, minus the KL penalty pulling them toward the BMA weights. This is optimized in log-space with SciPy's BFGS using a softmax parameterization and analytical gradients; λ is picked from a small grid using a held-out validation split.

5. Aggregate. The final posterior over latent variables is the weighted mixture of the per-model posteriors, using the learned weights.

Data. Three county-level expected annual loss tasks — wildfire, inland flooding, hurricane — derived from the FEMA National Risk Index dataset, with losses in millions of USD. Each task uses n = 20 training counties and 10 held-out counties for evaluation; losses are strictly positive and heavy-tailed.

Why This Matters

Impact on research. The paper makes a clean conceptual point for the LLM-driven probabilistic modeling literature: once models are generated by a language model, the "true model in the candidate set" assumption behind Bayesian model averaging is unjustified, and the aggregation rule should be chosen accordingly. It also shows a practical route for making stacking work when models have latent variables, by borrowing the same variational machinery already used inside the pipeline.

Real-world applications (potential):

• Catastrophe risk assessment — combining heterogeneous models of county-level wildfire, flood, and hurricane losses where no single model is correct and over-committing to one is risky.
• Automated statistical modeling for non-specialists — letting analysts describe a problem in text and get a blended predictive distribution rather than one arbitrarily selected program.
• Decision support under uncertainty — settings where downstream decisions benefit from a genuinely spread mixture rather than a point-like model choice.
• Model-uncertainty reporting — any workflow where the spread of weights is itself an output that stakeholders inspect.

Industry relevance. The work targets the reliability of LLM-generated quantitative pipelines: as more organizations use language models to draft statistical or simulation code, the question of how to combine several plausible drafts — rather than betting on one — becomes an engineering concern, not just a theoretical one.

Future Directions

• Broader task coverage. The authors state that future work includes evaluating the approach on more diverse problems; the current evidence rests on three hazard-loss tasks with 20 training points each.

• Alternative regularization. The KL anchor toward BMA weights with a grid-selected λ is one design choice; the paper lists exploring other regularization strategies as open work.

• Input-dependent stacking. The paper points to input-dependent, partially pooled weighting schemes (citing Bayesian hierarchical stacking) as a natural extension beyond a single global weight vector.

• Calibration of the uncertainty claim. The authors explicitly note that their experiments do not establish that the learned weights recover the true epistemic uncertainty of the data-generating process, and that the mixture still depends on the quality and diversity of the language model's candidate programs.

Target Audience

Researchers and practitioners working on Bayesian model combination, automated probabilistic programming, or LLM-assisted statistical modeling — particularly those who already know LLB or stacking and want to understand how to combine the two when models contain latent variables. It is also useful for applied scientists in risk and hazard modeling who care about representing model uncertainty rather than picking a single best model. Readers unfamiliar with ELBOs or Bayesian model averaging will need background reading first.

Authors’ abstract

Large Language Bayes (LLB) transforms a natural language problem description into a collection of candidate probabilistic programs and performs inference within each model, combining their predictions using Bayesian model averaging. While this provides a flexible framework for probabilistic inference without manual model specification, the resulting candidate model set is often misspecified in practice, and likelihood-based aggregation can concentrate on a single model, under-representing uncertainty across the sampled models. We propose a stacking-based aggregation method for LLB that selects model weights by optimizing predictive performance rather than marginal likelihood, while preserving the original language-model-driven model generation and within-model inference pipeline. Because leave-one-out predictive quantities are intractable in latent-variable models, we derive a variational lower bound that yields a tractable ELBO-based stacking objective. Across real-world probabilistic inference tasks, our method achieves comparable predictive performance to Bayesian model averaging while yielding more dispersed model weights, reducing concentration on a single candidate model under misspecification.