Stochasticity and Uncertainty › Ensembles, Bayesian Models, and Distributional Heads

How should a practitioner produce these distributions? Distributional heads, empirical ensembles, and Bayesian approximations make different cost and modeling choices. They are not interchangeable: a distributional likelihood addresses within-model outcome spread, while ensembles and Bayesian approximations target uncertainty over fitted functions or parameters.

For a scalar target, let the network emit $\mu_\theta(s_t, a_t)$ and $\ell_\theta(s_t, a_t)=\log \sigma_\theta^2(s_t, a_t)$, and set $\sigma_\theta^2=\exp(\ell_\theta)>0$ before evaluating the Gaussian negative log likelihood, up to the additive constant $\tfrac{1}{2}\log(2\pi)$ per observation:

$$\mathcal{L}(\theta) = \sum_{t} \left[ \frac{\big(s_{t+1} - \mu_\theta(s_t, a_t)\big)^2}{2\,\sigma^2_\theta(s_t, a_t)} + \frac{1}{2} \log \sigma^2_\theta(s_t, a_t) \right].$$

where:

• $s_{t+1}$: the observed next state in the training data
• $\mu_\theta(s_t, a_t)$: the network's predicted mean for that transition
• $\sigma^2_\theta(s_t, a_t)$: the network's predicted variance, which is a function of the state and action
• $\sum_t$: the sum over all transitions in the training set

For a vector state, this expression applies componentwise only under a diagonal conditional covariance; a full-covariance Gaussian uses a quadratic form and a log-determinant.

The first term is a squared error scaled by the predicted variance, and the second prices large variances. Without it, a model could send $\sigma^2$ to infinity and drive the scaled residual toward zero. The fitted likelihood trades variance against observed residuals; it does not prove that a learned large variance is causally irreducible noise.

The state- and action-dependent variance represents heteroscedastic uncertainty. A single global variance is misspecified when conditional noise varies with state or action, but is appropriate for genuinely homoscedastic systems. The likelihood lets the model trade residual size against predicted variance, and the logarithmic term prevents unbounded variance from being free. This learned attenuation can improve mean prediction on some tasks, but the result depends on the model, data, and optimization; the loss does not guarantee that high residuals are irreducible noise.

Quantile regression offers an alternative for selected marginals that does not assume a named parametric likelihood: pinball loss estimates selected conditional quantiles, from which intervals can be formed. Pinball fitting alone does not provide distribution-free finite-sample coverage; that guarantee requires additional assumptions or a calibration method such as an appropriate conformal procedure. A finite quantile grid does not specify an arbitrary full distribution, and independently trained quantiles can cross unless monotonicity or rearrangement is enforced. Within those limits, quantiles can represent asymmetric intervals and heavy tails more directly than a symmetric Gaussian head.

Pinball loss makes the asymmetry behind a target quantile explicit. With residual $u=y-q_\tau(x)$, the conventional loss is $\rho_\tau(u)=u(\tau-\mathbf{1}\{u<0\})$. When $\tau=0.1$, positive residuals, outcomes above the prediction and hence underprediction, receive weight $0.1$, while negative residuals receive weight $0.9$; at $\tau=0.9$ the weights reverse. At $\tau=0.5$ they are equal, yielding absolute-error conditional-median regression.

The second family trains $M$ models with controlled diversity, such as independent initializations, bootstrap samples, or shuffled data. Their predictions are empirical ensemble members, not exact posterior samples. Lakshminarayanan, Pritzel, and Blundell (2017) showed strong uncertainty and calibration results on their benchmarks. Training can run in parallel when hardware permits, but total compute and inference cost still grow with $M$. Member spread establishes predictive disagreement; interpreting it as epistemic uncertainty requires a validated diversity construction.

Initialization, resampling, data order, and hyperparameters can lead members to different fitted solutions. Their predictions often agree more in well-covered regions, but coverage does not guarantee agreement, and disagreement need not be caused by data scarcity alone. Because members usually share architecture and representation, they can also share the same extrapolation bias. Ensemble spread is therefore a model-class-dependent proxy, not a guarantee.

With equal-weight Gaussian members, averaging their means and variances and adding the variance of member means implements the corresponding finite-mixture variance identity. A mixture-headed ensemble can represent within-member multimodality and between-member disagreement, but labeling those terms aleatoric and epistemic remains approximate and does not guarantee calibration.

The third family maintains or approximates $p(\theta \mid \mathcal{D})$. Exact inference for general modern neural networks is computationally intractable, although small, discrete, linearized, or conjugate special cases can be tractable. Practical methods trade posterior fidelity for computational cost:

• Variational inference fits a tractable distribution $q(\theta)$, often a diagonal Gaussian, by optimizing the evidence lower bound. A diagonal Gaussian cannot represent parameter correlations; omitting those correlations can materially distort function uncertainty, alongside other limitations such as missing multimodality.
• Laplace approximation fits a local Gaussian around a trained differentiable network using loss curvature. Low-overhead post-hoc variants are available, but a single local Gaussian cannot represent separated posterior modes. Ensembles may span distinct fitted solutions, although they are not exact multimodal-posterior inference.
• Monte Carlo dropout uses dropout-compatible training and repeated stochastic test-time passes under stated variational assumptions. It is simple to add when dropout is already part of the model, but repeated passes have inference cost. See Gal and Ghahramani (2016).
• SWAG and related methods fit a Gaussian to statistics of stochastic-gradient iterates and use it for approximate Bayesian model averaging. See Maddox et al. (2019).

Deep ensembles are a strong empirical baseline. Distributional heads are useful when conditional outcome spread matters, while Bayesian approximations offer alternatives when ensemble cost or posterior structure matters. Some systems, such as PETS, combine probabilistic ensembles with sampling-based propagation. The appropriate layering depends on the uncertainty sources and deployment cost.