R2-D2

Description

The R2-D2 prior (Zhang et al. 2022) is a global-local shrinkage prior. This has been expanded for generalised linear models (Yanchenko et al. 2024).

Definition

For a linear model:

\[ \begin{align} y_i &\sim \text{normal}(\mu_i,\sigma^2) \\ \mu_i &= \alpha + \sum_{k=1}^{K}\beta_k x_{ik} \end{align} \]

The prior is:

\[ \begin{align} \beta_k &\sim \text{normal}\left(0, \frac{\sigma^2}{\sigma_{y_{t}\mid \theta}^2}\tau^2\psi_k\right) \\ \tau^2 &= \frac{R^2}{1 - R^2} \\ R^2 &\sim \text{beta}(\mu_{R^2},\phi_{R^2}) \\ \psi &\sim \text{Dirichlet}(\xi_1,\dotsc,\xi_p) \\ \sigma^2 &\sim p(\sigma^2) \\ \alpha &\sim p(\alpha) \end{align} \]

The proportions \(\psi_k\) determine how the total prior variance \(\tau^2\) is allocated among the regression coefficients.

Things to specify

Priors on \(\alpha\) and \(\sigma\). Hyperparameters \(\xi\), \(\mu_{R^2}\), \(\phi_{R^2}\).

Implementation

functions {
  /* Efficient computation of the R2D2 prior
   * Args:
   *   z: standardized population-level coefficients
   *   phi: local weight parameters
   *   tau2: global scale parameter
   * Returns:
   *   population-level coefficients following the R2D2 prior
   */
  vector R2D2(vector z, vector phi, real tau2) {
    return z .* sqrt(phi * tau2);
  }
}

data {
  int<lower=0> N;
  int<lower=0> K;
  matrix[N, K] X;
  vector[N] Y;
  int prior_only;  // should the likelihood be ignored?
  // concentration vector of the D2 prior
  vector<lower=0>[K-1] R2D2_cons_D2;
  // data for the R2D2 prior
  real<lower=0> R2D2_mean_R2;  // mean of the R2 prior
  real<lower=0> R2D2_prec_R2;  // precision of the R2 prior
}
transformed data {
  int Kc = K - 1;
  matrix[N, Kc] Xc; // centered version of X without an intercept
  vector[Kc] means_X; // column means of X before centering
  real sd_Y = sd(Y);
  for (i in 2 : K) {
    means_X[i - 1] = mean(X[ : , i]);
    Xc[ : , i - 1] = X[ : , i] - means_X[i - 1];
  }
}
parameters {
  // local parameters for the R2D2 prior
  vector[Kc] zb;
  simplex[Kc] R2D2_phi;
  // R2D2 shrinkage parameters
  real<lower=0,upper=1> R2D2_R2;  // R2 parameter
  real<lower=0> sigma;  // dispersion parameter
  real a_c;
}

transformed parameters {
  vector[Kc] b;  // population-level effects
  real R2D2_tau2;  // global R2D2 scale parameter
  array[Kc+4] real lprior;
  R2D2_tau2 = sigma^2 * R2D2_R2 / (1 - R2D2_R2);
  // compute actual regression coefficients
  b = R2D2(zb, R2D2_phi, R2D2_tau2);
  lprior[1] = student_t_lpdf(a_c | 4, 0, sd_Y);
  for (k in 1:Kc) {
    lprior[k+1] = std_normal_lpdf(zb[k]);
  }
  lprior[Kc+2] = beta_lpdf(R2D2_R2 | R2D2_mean_R2 * R2D2_prec_R2, (1 - R2D2_mean_R2) * R2D2_prec_R2);
  lprior[Kc+3] = dirichlet_lpdf(R2D2_phi | R2D2_cons_D2);
  lprior[Kc+4] = student_t_lpdf(sigma | 3, 0, sd_Y);
  
}

model {
  // likelihood including constants
  if (!prior_only) {
    target += normal_id_glm_lpdf(Y | Xc, a_c, b, sigma);
  }
  // priors including constants
  target += sum(lprior);
}

import numpy as np
import pymc as pm
import pytensor.tensor as pt

# center and scale predictors
# Assume X is your predictor matrix and y is your outcome vector
X_mean = X.mean(axis=0)
X_sd = X.std(axis=0, ddof=1)
Xc = (X - X_mean) / X_sd
sd_y = y.std(ddof=1)

# HyperParameters for the R2D2 prior
# Adjust as needed for your specific problem
K = X.shape[1]
cons_D2 = np.full(K, 1)
mean_R2 = 0.3
prec_R2 = 3.0

with pm.Model() as model:
    # prior for the intercept and error term
    # not part of the R2D2 prior, but needed for the model
    sigma = pm.HalfStudentT("sigma", nu=3, sigma=sd_y)
    a_c = pm.StudentT("a_c", nu=4, mu=y.mean(), sigma=sd_y)

    # R2D2 prior
    zb = pm.Normal("zb", mu=0.0, sigma=1.0, shape=K)
    phi = pm.Dirichlet("phi", a=cons_D2)
    R2 = pm.Beta("R2", mu=mean_R2, nu=prec_R2)

    # compute coefficients following the R2D2 prior
    tau2 = sigma**2 * R2 / (1 - R2)
    b = pm.Deterministic("b", zb * pt.sqrt(phi * tau2))

    mu = a_c + Xc @ b
    pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y)

References

Yanchenko, Eric, Howard D. Bondell, and Brian J. Reich. 2024. The R2D2 Prior for Generalized Linear Mixed Models. arXiv:2111.10718. arXiv. https://doi.org/10.48550/arXiv.2111.10718.

Zhang, Yan Dora, Brian P. Naughton, Howard D. Bondell, and Brian J. Reich. 2022. “Bayesian Regression Using a Prior on the Model Fit: The R2-D2 Shrinkage Prior.” Journal of the American Statistical Association 117 (538): 862–74. https://doi.org/10.1080/01621459.2020.1825449.