ARR2
Description
The ARR2 prior (Kohns et al. 2024) is a global-local shrinkage prior for autoregressive models. It adapts the R2-D2 prior (Zhang et al. 2022) to the time-series setting by placing a joint prior on the autoregressive coefficients and their induced \(R^2\). Compared to other priors for time-series models, the ARR2 prior allows for flexible and interpretable shrinkage: the user can encode prior information about the overall explanatory power of the lagged outcomes (via \(R^2\)) and about the relative importance of each lag (via the simplex proportions \(\psi\)).
Definition
For an AR(\(p\)) model:
\[ \begin{align} y_t &\sim \text{normal}(\mu_t, \sigma^2), \quad t = p+1, \dotsc, T \\ \mu_t &= \sum_{i=1}^{p} \phi_i \, y_{t-i} \end{align} \]
the ARR2 prior is:
\[ \begin{align} \phi_i &\sim \text{normal}\!\left(0,\; \frac{\sigma^2}{\sigma_y^2}\,\tau^2\,\psi_i\right) \\[6pt] \tau^2 &= \frac{R^2}{1 - R^2} \\[6pt] R^2 &\sim \text{beta}(\mu_{R^2},\, \varphi_{R^2}) \\[6pt] \psi &\sim \text{Dirichlet}(\xi_1, \dotsc, \xi_p) \\[6pt] \sigma &\sim p(\sigma) \end{align} \]
where \(\sigma_y^2\) is the (unconditional) sample variance of \(y\).
The signal-to-noise ratio \(\tau^2 = R^2/(1-R^2)\) reparameterizes the prior in terms of \(R^2\), the proportion of variance in \(y\) explained by the autoregressive component. The proportions \(\psi_i\) lie on the simplex (\(\psi_i \geq 0\), \(\sum_i \psi_i = 1\)) and determine how the total prior variance \(\tau^2\) is allocated among the \(p\) lag coefficients. A symmetric Dirichlet concentration \(\xi_i = a_\pi\) with \(a_\pi \leq 1/2\) favours sparsity, concentrating the explanatory power on a few lags.
The scaling by \(\sigma^2 / \sigma_y^2\) ensures that the prior on the AR coefficients is invariant to the units of \(y\): the coefficients \(\phi_i\) are dimensionless (they relate \(y_t\) to past values of \(y\) in the same units), and the factor \(\sigma_y^2\) normalizes the prior variance relative to the unconditional variability of the series.
Things to specify
Prior on \(\sigma\). Hyperparameters \(\mu_{R^2}\) and \(\varphi_{R^2}\) (mean and precision of the \(R^2\) prior), and \(\xi_1, \dotsc, \xi_p\) (Dirichlet concentration parameters for the lag-coefficient proportions).
Implementation
data {
int<lower=1> T; // number of time points
vector[T] Y; // observations
int<lower=0> p; // AR order
// concentration vector of the Dirichlet prior
vector<lower=0>[p] cons;
// data for the R2D2 prior
real<lower=0> mean_R2; // mean of the R2 prior
real<lower=0> prec_R2; // precision of the R2 prior
real<lower=0> sigma_sd; // sd of sigma prior
// variance estimates of y
real<lower=0> var_y;
}
parameters {
vector[p] phi; // AR coefficients
simplex[p] psi; // decomposition simplex
real<lower=0, upper=1> R2; // coefficient of determination
real<lower=0> sigma; // observation model sd
}
transformed parameters {
real<lower=0> tau2 = R2 / (1 - R2); // Equation 18
vector[T] mu = rep_vector(0.0, T);
for (t in (p+1):T) {
for (i in 1:p) {
mu[t] += phi[i] * Y[t-i]; // Equation 16
}
}
}
model {
// priors
phi ~ normal(0, sqrt(sigma^2/var_y * tau2 * psi)); // Equation 17
R2 ~ beta(mean_R2 * prec_R2, (1 - mean_R2) * prec_R2); // Equation 19
sigma ~ normal(0, sigma_sd); // Equation 20
psi ~ dirichlet(cons); // Equation 21
// likelihood
Y ~ normal_lpdf(mu, sigma); // Equation 15
}import numpy as np
import pymc as pm
import pytensor.tensor as pt
# Specify the order of the autoregressive model and the observed data
p = 2
y_data = np.array([...]) # Replace with your actual time series data
# Hyperparameters for the ARR2 prior
cons = 0.1 * np.ones(p)
mean_R2 = 1 / 3
prec_R2 = 3
# Scale to make HalfNormal have unit variance
HALFNORMAL_SCALE = 1 / np.sqrt(1 - 2 / np.pi)
with pm.Model() as model:
sigma = pm.HalfNormal("sigma", HALFNORMAL_SCALE)
zb = pm.Normal("zb", mu=0.0, sigma=1.0, shape=p)
psi = pm.Dirichlet("psi", a=cons)
R2 = pm.Beta("R2", mu=mean_R2, nu=prec_R2)
tau2 = R2 / (1 - R2)
phi = pm.Deterministic("phi", zb * (sigma / y_data.std()) * pt.sqrt(tau2 * psi))
pm.AR(
"obs",
phi,
np.sqrt(2) * sigma,
init_dist=pm.Normal.dist(0, 1, shape=p),
observed=y_data,
)