Bayesian spline model

The first Bayesian model that ChatGPT generated can be described using this notation:

Data

• \(N\): number of individuals
• \(K\): number of clusters
• \(M\): number of spline basis functions
• \(y_{ik}\): outcome for individual \(i\) in cluster \(k\), \(i \in 1,\dots,N\), \(k \in 1,\dots ,K\)
• \(\boldsymbol{X} \in \mathbb{R}^{N \times M}\): matrix of spline basis function values

Parameters

• \(\boldsymbol{\beta_{k}} \in \mathbb{R}^M\): spline coefficients for cluster \(k\) (a vector of length \(M\) for each cluster)
• \(\sigma_y\): standard deviation of the observation noise
• \(\sigma_\beta\): prior standard deviation for the spline coefficients

Model

Likelihood:

\[ y_{ik} \sim N\left( \sum_{m=1}^M X_{im} \beta_{km}, \sigma_y \right), \ i \in 1,\dots, N, \ k \in 1, \dots, K\]

Priors:

\[ \boldsymbol{\beta_{k}} \sim N(0, \sigma_{\beta} \boldsymbol{I_M}), \ \ k \in 1,...,K \\ \sigma_{y} \sim N(0, 1), \ \ \sigma_y \gt 0 \\ \sigma_{\beta} \sim N(0, 1), \ \ \sigma_{\beta} \gt 0 \]

The Stan code provided by ChatGPT aligns with this description. As part of the model, I also requested code to generate outcome predictions for each observation, which is implemented in the generated quantities block. My goal was to plot the median of those predictions for each individual \(i\) as a comparison to the GAM plot above.”

data {
  int<lower=1> N;                          // number of observations
  int<lower=1> K;                          // number of clusters
  int<lower=1> M;                          // number of basis functions
  array[N] int<lower=1, upper=K> cluster;  // cluster ids
  matrix[N, M] X_spline;                   // basis function values
  vector[N] y;                             // response variable
}

parameters {
  matrix[K, M] beta;          // cluster-specific spline coefficients
  real<lower=0> sigma_y;      // observation noise
  real<lower=0> sigma_beta;   // prior standard deviation for beta
}

model {
  sigma_y ~ normal(0, 1);
  sigma_beta ~ normal(0, 1);
  
  // Priors for beta
  
  for (k in 1:K) {
    beta[k] ~ normal(0, sigma_beta);
  }
  
  // Likelihood
  
  for (n in 1:N) {
    y[n] ~ normal(X_spline[n] * beta[cluster[n]]', sigma_y);
  }
}

generated quantities {
  
  vector[N] y_pred;                    // Vector of observations.
  
  for (n in 1:N) {
    y_pred[n] = normal_rng(X_spline[n] * beta[cluster[n]]', sigma_y);
  }
}

Spline basis functions

In the likelihood, \(y_i\) is modeled as a function of the vector \(\boldsymbol{X_i}\) rather than the single measurement \(x_i\). While I won’t delve deeply into spline estimation, I want to conceptually outline how this vector is constructed in the context of cubic splines.

We control the flexibility of the curve by specifying the number of knots. A unique curve is fitted between each pair of knots (as well as at the ends), with constraints ensuring smooth transitions between these curves. The estimation of these curves is performed using basis functions, specifically B-spline basis functions of \(x\).

The number of basis functions is determined by the number of knots. For instance, the plot below illustrates the \(M=9\) basis functions required for \(K=5\) knots. Each basis function contributes an element to the vector \(\boldsymbol{X}\) for each value of \(x\). In the case of cubic splines, at most four basis functions can be non-zero between any two knots, as indicated by the intervals on the x-axis. Consequently, the vector \(\boldsymbol{X}\) consists of the values of each basis function at a given point \(x\), with at most four non-zero entries corresponding to the active basis functions. (As an example, in the plot below there is a vertical line at a single point \(x\) that passes through four basis functions.)

This example uses \(M = 5\) knots to introduce a slight overfitting of the data, which will allow me to apply another model in the next step that will further smooth the curves. (In a real-world setting, it may have made more sense to start out with fewer knots.) The bs function (in the splines package) computes the B-spline basis function values for each observed \(x\).

n_knots <- 5
knot_dist <- 1/(n_knots + 1)
probs <- seq(knot_dist, 1 - knot_dist, by = knot_dist)
knots <- quantile(dd$x, probs = probs)
spline_basis <- bs(dd$x, knots = knots, degree = 3, intercept = TRUE)
X_spline <- as.matrix(spline_basis)

Data list for stan

To fit the model, we need to create the data set that Stan will use to estimate the parameters.

stan_data <- list(
  N = nrow(dd),           # number of observations
  K = k,                  # number of clusters
  M = ncol(X_spline),     # number of basis functions
  cluster = dd$cluster,   # vector of cluster ids
  X_spline = X_spline,    # basis function values
  y = dd$y                # response variable
)

Run stan model

ChatGPT provided code to estimate the model using the rstan package. However, I prefer using the cmdstanr package, which I find more stable and generally less finicky. From the plot, you can see that the estimation was quite good. However, the curves are a bit too wiggly, indicating the data may have been slightly overfit, particularly for clusters 1, 3, and 7.

mod <- cmdstan_model("code/spline.stan")

fit <- mod$sample(
  data = stan_data,
  chains = 4,
  iter_warmup = 500,
  iter_sampling = 2000,
  parallel_chains = 4,
  refresh = 0 # print update every 500 iters
)
## Running MCMC with 4 parallel chains...
## Chain 2 finished in 5.4 seconds.
## Chain 1 finished in 5.6 seconds.
## Chain 3 finished in 5.6 seconds.
## Chain 4 finished in 5.8 seconds.
## 
## All 4 chains finished successfully.
## Mean chain execution time: 5.6 seconds.
## Total execution time: 5.9 seconds.
draws <- as_draws_df(fit$draws())

ds <- summarize_draws(draws, .fun = median) |> data.table()
dd$np <- ds[substr(variable, 1, 3) == "y_p", 2]