General Bayesian Policy Learning
| Authors | Masahiro Kato |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2602.23672 |
| Download | |
| Categories | stat.ML, cs.LG, stat.ME |
Abstract
This study proposes the General Bayes framework for policy learning. We consider decision problems in which a decision-maker chooses an action from an action set to maximize its expected welfare. Typical examples include treatment choice and portfolio selection. In such problems, the statistical target is a decision rule, and the prediction of each outcome is not necessarily of primary interest. We formulate this policy learning problem by loss-based Bayesian updating. Our main technical device is a squared-loss surrogate for welfare maximization. We show that maximizing empirical welfare over a policy class is equivalent to minimizing a scaled squared error in the outcome difference, up to a quadratic regularization controlled by a tuning parameter ζ>0. This rewriting yields a General Bayes posterior over decision rules that admits a Gaussian pseudo-likelihood interpretation. We clarify two Bayesian interpretations of the resulting generalized posterior, a working Gaussian view and a decision-theoretic loss-based view. As one implementation example, we introduce neural networks with tanh-squashed outputs. Finally, we provide theoretical guarantees in a PAC-Bayes style.
Engineering Breakdown
Plain English
This paper proposes a General Bayes framework for learning optimal decision policies in settings like treatment selection and portfolio allocation, where the goal is to maximize expected welfare rather than predict outcomes accurately. The key insight is reformulating policy learning as a loss-based Bayesian updating problem, where maximizing empirical welfare over a policy class becomes equivalent to minimizing a scaled squared error in outcome differences, controlled by a tuning parameter ζ. This mathematical rewriting bridges Bayesian inference and policy optimization, enabling practitioners to leverage standard regression techniques for decision-making problems. The framework is general enough to handle diverse applications while maintaining theoretical grounding.
Core Technical Contribution
The paper's main technical novelty is the squared-loss surrogate for welfare maximization, which transforms a discrete, non-convex policy optimization problem into a continuous, tractable regression problem. By showing that empirical welfare maximization is equivalent to minimizing scaled squared error up to quadratic regularization, the authors create a bridge between Bayesian model estimation and decision rule learning. This equivalence is non-obvious and allows practitioners to apply standard statistical machinery (regularized regression, cross-validation) to policy learning without developing domain-specific optimization algorithms. The General Bayes formulation unifies several previously disparate approaches under one coherent framework.
How It Works
The framework operates as follows: given a decision problem where an agent must choose action a from action set A to maximize expected welfare, the surrogate loss function measures squared differences in predicted outcomes Y(a) between candidate actions. During Bayesian updating, the posterior distribution is computed using this squared-loss surrogate rather than a traditional likelihood, creating a 'loss-based' posterior. To learn an optimal policy, one then solves an empirical risk minimization problem: minimize the squared error over outcome predictions across the training data, subject to a quadratic penalty term weighted by ζ. The regularization parameter ζ controls the trade-off between fit and complexity, and can be tuned via cross-validation. This transforms the problem into standard supervised learning territory, where gradient-based optimization and existing statistical tools apply directly.
Production Impact
For production systems, this framework dramatically simplifies policy learning pipelines. Instead of building custom optimization routines for each decision problem (treatment allocation, recommendation ranking, portfolio construction), engineers can reuse existing regression libraries, hyperparameter tuning infrastructure, and cross-validation tools. The squared-loss surrogate makes the problem convex and differentiable, enabling efficient computation at scale on modern ML platforms. However, practitioners must carefully select the regularization parameter ζ—too high and the policy becomes overly conservative; too low and it overfits to training data. Additionally, the approach assumes outcome predictions drive welfare differences; in settings with complex interactions or non-linear reward structures, the squared-loss proxy may miss important dynamics, requiring careful validation before deployment.
Limitations and When Not to Use This
The paper does not address online or adaptive policy learning where the decision environment changes over time, limiting applicability to bandit or reinforcement learning scenarios. The squared-loss surrogate is a proxy for true welfare maximization, and no finite-sample guarantees on policy regret or welfare loss are provided in the abstract—only equivalence to a regression problem, which holds in expectation but may degrade with limited data. The framework assumes access to complete outcome data Y(a) for each action a, which is unrealistic in observational settings where only taken actions are observed (the counterfactual outcomes are missing). The approach also assumes the outcome differences between actions are what drive welfare, potentially failing when reward depends on higher-order statistics or distributional properties rather than point predictions.
Research Context
This work builds on a long tradition of Bayesian decision theory and causal inference for policy learning, extending frameworks like marginal structural models and inverse probability weighting with a modern optimization perspective. The loss-based Bayesian updating idea extends prior work on loss-based posteriors (e.g., in robust Bayesian analysis) to the policy learning domain. The paper's emphasis on reframing policy optimization as regression connects to recent trends in causal ML and targeted learning, which also use surrogate losses to improve statistical and computational efficiency. By unifying Bayesian inference with policy optimization, it opens pathways for hybrid approaches that leverage uncertainty quantification (from Bayesian posteriors) and computational efficiency (from regression) simultaneously.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::