Beyond Point Predictions: A Practical Introduction to Bayesian

Machine Learning


Beyond Point Predictions: A Practical Introduction to Bayesian Neural Networks


More informed decision-making through uncertainty quantification


By Tom Narock | September 1, 2026 | 16 min read


Image generated with AI


The Problem with Point Estimates


When a machine learning model estimates the median value of a house in California, it typically returns a single number—say, $385,000. That figure may appear precise, but it conceals significant uncertainty. Is the prediction uncertain because the neighborhood features are inherently noisy, or because the model has never encountered a location quite like this one?


Standard neural networks that output single point values cannot answer these questions. They rely on aggregate error metrics like mean absolute error (MAE), which only reflect average performance across a test dataset. MAE offers no insight into the specific house in question. By adopting a Bayesian framework, we shift from asking "What is the predicted home value?" to "What is the range of plausible values, and how confident are we?"


Recent discussions, such as the article on Bayesian Guardrails for AI Decisions, have compellingly argued for incorporating uncertainty in automated decision-making. In this article, we provide a hands-on implementation and explore the associated engineering decisions. Using Python and the California Housing dataset, we will build a Bayesian Neural Network (BNN), parameterize weight distributions, perform variational inference, and extract actionable uncertainty bounds that reveal when the model is—and isn't—trustworthy.


Ideally, we'd apply Bayes' theorem directly to compute exact solutions for a neural network. However, this exact calculation is computationally intractable; even a network with just a few thousand parameters presents a math problem that, while theoretically solvable, requires too much time and memory to be practical. Consequently, BNN implementation involves several design choices and training strategies. Our focus here is to demystify these choices and demonstrate their implementation. The shared Python notebooks are not optimized for network size or hyperparameters; instead, they illustrate the engineering nuances of BNNs. Each notebook highlights one specific design choice, and a production system would likely combine several of these approaches.


Uncertainty Quantification


Before diving into the data and code, we must deepen our understanding of uncertainty quantification, as not all uncertainty is equal. Philosophically, uncertainty can be broadly divided into two types: aleatoric (irreducible, inherent noise in the data) and epistemic (reducible, stemming from model limitations or lack of data). In 2026, with the proliferation of AI in high-stakes domains like healthcare and autonomous systems, distinguishing between these types is not just academic—it's critical for risk management. Aleatoric uncertainty, for instance, might be the natural variability in housing prices in a bustling urban area, while epistemic uncertainty could arise from a model trained only on suburban data being asked to predict urban property values.


Bayesian Neural Networks excel at capturing both, with epistemic uncertainty being particularly valuable for out-of-distribution detection—a hot topic in AI safety and reliability. By representing model parameters as probability distributions rather than fixed values, BNNs naturally provide a measure of confidence in their predictions.


Design Choices for BNN Implementation


Implementing a BNN involves several key decisions:


  1. Prior Selection: Choosing appropriate prior distributions for weights (e.g., Gaussian, Laplace) reflects our initial beliefs before seeing data.
  2. Approximation Method: Since exact inference is intractable, we rely on techniques like variational inference (VI) or Markov Chain Monte Carlo (MCMC). VI is often preferred for scalability, but MCMC offers asymptotic exactness.
  3. Loss Function: The loss combines data likelihood and KL divergence, balancing fit with complexity.

  4. Our notebooks explore these choices, but let's first set the stage with the data and a baseline.


    Setting Up the Experiment


    To illustrate BNN concepts, we'll use the classic California Housing dataset, which contains median house values and features like income, housing age, and location (latitude/longitude). In 2026, such datasets are still standard for benchmarking, but they also serve as a proxy for more complex problems—think predicting energy consumption or financial risk.


    We'll build a simple feedforward network with a single hidden layer. In a BNN, we won't assign fixed weights; instead, each weight will be a random variable, typically initialized with a prior distribution. For simplicity, we'll use a variational inference approach, where we learn the parameters (mean and variance) of these distributions to approximate the true posterior.


    Step-by-Step Implementation


    Here's a high-level breakdown of our BNN pipeline:


    1. Data Preparation: Split the dataset into training and test sets, standardizing features.
    2. Model Definition: Define a neural network where each layer outputs a distribution over activations.
    3. Variational Inference Setup: Implement the Evidence Lower Bound (ELBO) loss, which we maximize.
    4. Training: Optimize the ELBO using gradient descent, tracking both loss and uncertainty metrics.
    5. Prediction: For a new input, sample from the weight distributions to obtain a distribution of outputs.

    6. Let's dive into the code. We'll use PyTorch and the torch.distributions module for flexibility.


      # Import necessary libraries
      import torch
      import torch.nn as nn
      import torch.nn.functional as F
      from torch.distributions import Normal, kl_divergence
      from sklearn.datasets import fetch_california_housing
      from sklearn.model_selection import train_test_split
      from sklearn.preprocessing import StandardScaler
      import numpy as np
      
      # Load data
      data = fetch_california_housing()
      X, y = data.data, data.target
      X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
      
      # Standardize features
      scaler = StandardScaler()
      X_train = scaler.fit_transform(X_train)
      X_test = scaler.transform(X_test)
      
      # Convert to tensors
      X_train_t = torch.tensor(X_train, dtype=torch.float32)
      y_train_t = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)
      X_test_t = torch.tensor(X_test, dtype=torch.float32)
      y_test_t = torch.tensor(y_test, dtype=torch.float32).view(-1, 1)
      

      Now, we define a Bayesian layer that learns the distribution of weights.


      class BayesianLinear(nn.Module):
          def __init__(self, in_features, out_features, prior_var=1.0):
              super().__init__()
              # Define mean and log variance parameters for weights and biases
              self.weight_mu = nn.Parameter(torch.randn(out_features, in_features) * 0.1)
              self.weight_log_var = nn.Parameter(torch.randn(out_features, in_features) * 0.1)
              self.bias_mu = nn.Parameter(torch.randn(out_features) * 0.1)
              self.bias_log_var = nn.Parameter(torch.randn(out_features) * 0.1)
              self.prior_var = prior_var
      
          def forward(self, x, sample=False):
              if sample:
                  # Sample weights from their distributions
                  weight_std = torch.exp(self.weight_log_var / 2)
                  bias_std = torch.exp(self.bias_log_var / 2)
                  weight = Normal(self.weight_mu, weight_std).rsample()
                  bias = Normal(self.bias_mu, bias_std).rsample()
              else:
                  # Use mean weights for deterministic forward pass
                  weight = self.weight_mu
                  bias = self.bias_mu
              return F.linear(x, weight, bias)
      
          def kl_loss(self):
              # Compute KL divergence between posterior and prior
              def kl_div(mu, log_var, prior_var):
                  var = torch.exp(log_var)
                  kl = 0.5 * (log_var - torch.log(torch.tensor(prior_var)) - 1 + (var + mu**2) / prior_var)
                  return kl.sum()
              return kl_div(self.weight_mu, self.weight_log_var, self.prior_var) + kl_div(self.bias_mu, self.bias_log_var, self.prior_var)
      

      Next, we build the BNN model with one Bayesian layer and a standard output layer.


      class BNN(nn.Module):
          def __init__(self, input_dim, hidden_dim, output_dim):
              super().__init__()
              self.bayes_layer = BayesianLinear(input_dim, hidden_dim)
              self.output_layer = nn.Linear(hidden_dim, output_dim)
      
          def forward(self, x, sample=False):
              h = torch.relu(self.bayes_layer(x, sample))
              return self.output_layer(h)
      
          def loss(self, x, y, sample=True):
              y_pred = self.forward(x, sample)
              # Gaussian likelihood with fixed variance (e.g., 0.1)
              likelihood = Normal(y_pred, torch.tensor(0.1))
              log_lik = likelihood.log_prob(y).sum()
              kl = self.bayes_layer.kl_loss()
              # ELBO: maximize log_lik - kl, so minimize negative
              return -log_lik + kl
      

      Training the BNN


      We'll train the model by minimizing the negative ELBO, which combines reconstruction loss and KL divergence.


      model = BNN(input_dim=X_train_t.shape[1], hidden_dim=64, output_dim=1)
      optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
      
      epochs = 200
      for epoch in range(epochs):
          optimizer.zero_grad()
          loss = model.loss(X_train_t, y_train_t, sample=True)
          loss.backward()
          optimizer.step()
          if (epoch+1) % 50 == 0:
              print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}')
      

      Making Predictions with Uncertainty


      After training, we can predict by sampling multiple times and computing mean and standard deviation.


      def predict_with_uncertainty(model, X, samples=100):
          model.eval()
          preds = []
          with torch.no_grad():
              for _ in range(samples):
                  pred = model.forward(X, sample=True)
                  preds.append(pred.numpy())
          preds = np.array(preds)
          mean = preds.mean(axis=0)
          std = preds.std(axis=0)
          return mean, std
      
      mean_pred, std_pred = predict_with_uncertainty(model, X_test_t)
      

      We can then compute metrics like RMSE and examine the uncertainty on specific test points. For example, we might find that predictions with higher standard deviations are less accurate, allowing us to flag them for manual review.


      Visualizing Uncertainty


      To better interpret the results, we can plot the true values against predicted means, with error bars representing ±1 standard deviation. This visualization reveals where the model is confident and where it isn't, guiding us in deploying the model appropriately.


      Key Takeaways and Engineering Considerations


      • Bayesian Neural Networks offer a principled way to quantify uncertainty, but they come with computational overhead. In 2026, enhanced hardware and optimized libraries are making BNNs more viable for real-time applications.
      • Variational inference is a practical choice for large-scale problems, but MCMC remains valuable for smaller datasets where exactness is crucial.
      • Prior selection significantly impacts performance; non-informative priors are often a safe default, but domain knowledge can improve results.
      • The ELBO loss is not a perfect proxy for model accuracy; it balances fit with complexity, so monitoring additional metrics is wise.
      • Production systems should combine multiple techniques, such as ensembling BNNs or using them in conjunction with deterministic networks, to leverage their strengths.

      Conclusion


      In this article, we've walked through a practical implementation of a Bayesian Neural Network, highlighting the key design choices and engineering decisions. By moving beyond point predictions, we enable more informed decision-making, especially in critical applications where understanding uncertainty is paramount. As we look toward the future of AI, tools like BNNs will become indispensable for building trustworthy systems.


      Further reading: Explore the provided notebooks for more detailed experiments and comparisons with alternative approaches.

      via Towards Data Science

Related