RefineStat: Efficient Exploration for Probabilistic Program Synthesis

TL;DR

RefineStat addresses the challenge of automatically synthesizing reliable probabilistic programs with small language models by enforcing semantic correctness through a constrained decoding process and by applying diagnostic-aware refinement within the Bayesian workflow. It integrates grammar-based semantic validation with resampling of priors or likelihoods guided by standard MCMC diagnostics to prune invalid candidates and improve sampler health. Empirical results across five datasets and multiple open-weight LLMs show substantial improvements in run-rate and reliability, with competitive predictive performance compared to BoxLM and OpenAI o3. The approach generalizes across probabilistic backends (e.g., PyMC and NumPyro) and demonstrates that constrained, diagnostic-guided generation can yield high-quality probabilistic programs with smaller models. This has practical implications for scalable, reliable probabilistic-model discovery in settings where access to large LLMs is limited or costly.

Abstract

Probabilistic programming offers a powerful framework for modeling uncertainty, yet statistical model discovery in this domain entails navigating an immense search space under strict domain-specific constraints. When small language models are tasked with generating probabilistic programs, they frequently produce outputs that suffer from both syntactic and semantic errors, such as flawed inference constructs. Motivated by probabilistic programmers' domain expertise and debugging strategies, we introduce RefineStat, a language model--driven framework that enforces semantic constraints ensuring synthesized programs contain valid distributions and well-formed parameters, and then applies diagnostic-aware refinement by resampling prior or likelihood components whenever reliability checks fail. We evaluate RefineStat on multiple probabilistic-programming code-generation tasks using smaller language models (SLMs) and find that it produces programs that are both syntactically sound and statistically reliable, often matching or surpassing those from closed-source large language models (e.g., OpenAI o3).

Figures (3)

• Figure 1: RefineStat workflow: (1) A user provides data and prompt to the language model, which generates a probabilistic program. (2) Constrained semantic decoding enforces syntactic and semantic validity of the generated program. (3) A Bayesian reliability check diagnoses convergence, divergences, and predictive validity. If failures are detected, the model is refined by backtracking and resampling priors or likelihoods. (4) Upon passing checks, we get final reliable probabilistic program.
• Figure 2: Constrained decoding in RefineStat fixing a TypeError from using sd instead of sigma, as illustrated in Section \ref{['intro']}.
• Figure : RefineStat Synthesis via $\mathcal{D}\Vert \mathcal{P}\Vert\mathcal{L}$

Theorems & Definitions (10)

• Definition 1: PSIS-LOO vehtari2017practical
• Definition 2: Bayesian Workflow Reliability Score
• Definition 3: Valid model space
• Definition 4: RefineStat Objective
• Definition 5: $\widehat{R}$ Statistic
• Definition 6: Effective Sample Size
• Definition 7: Divergent Transitions
• Definition 8: Bayesian Fraction of Missing Information
• Definition 9: Pareto Shape Parameter
• Definition 10: Expected Log Pointwise Predictive Density