[avehtari]

Summary

This PR computes the predictive variance of the reference model for the Gaussian and Student-t families with two matrix products instead of one R-level weighted.mean() call per observation. project() with 10 000 observations and 4 000 projected draws gets about 7 times faster. For small numbers of observations, the run time does not change. The results do not change.

This PR was assisted by Claude.

Overhead

predvar_gauss() and predvar_student_t() compute, for each observation, the weighted variance of mu across the draws with one weighted.mean() call per observation. get_refdist() calls family$predvar() once per projected draw, and get_p_clust() once per cluster, so a projection with S draws makes n × S R-level calls. In a profile of project() for a Gaussian model with n = 10 000 and 4 000 projected draws, get_refdist() took 88 % of the time, nearly all in weighted.mean().

Changes

• The new internal rowvar_wdraws() computes the same weighted variances for all observations with two matrix products.
• predvar_gauss() and predvar_student_t() use it. predvar_gamma() does not compute a variance, so it does not change.

Results

For random inputs (one draw, 400 draws, unequal weights, missing values), the largest relative difference to the old code is 1.1e-15, with the same missing values. For one draw, the variance term is exactly 0, as before. One call with n = 10 000 and one draw takes 0.05 ms instead of 34 ms.

Gaussian reference model with P = 50 and 4 000 draws; all results identical to master (largest relative difference 1.6e-16, identical rankings):

case	master	this PR
project(), n = 10 000, 4 000 projected draws	239 s	33 s
project(), n = 50 000, 400 projected draws	114 s	11 s
project(), n = 1 000, 4 000 projected draws	29 s	14 s
varsel(), n = 10 000, 400 projected draws	122 s	99 s

For small n (N = 50, P = 200), the run time does not change.

Tests

The new test_predvar.R checks that:

• rowvar_wdraws() equals the row-wise weighted variance of the old code, for unequal and equal weights; that it gives exactly 0 for one draw; and that a missing value gives a missing variance for its observation only;
• predvar() of the extended gaussian() and Student_t() families equals the weighted mean of the squared dispersion (times nu / (nu - 2) for Student-t) plus the weighted variance of mu.

The first test fails on master (function not found). The second test passes on master and on this branch; it checks that the change keeps the results. The full suite passes with NOT_CRAN=true: 255 tests, 55 813 expectations, 0 failures, the same 4 warnings as master. Tested with this change applied to master 73cf0ee (projpred 2.11.0.9000), R 4.6.1, brms 2.23.0, rstanarm 2.32.2, cmdstanr 0.9.0.9001, CmdStan 2.39.0, on CSC Roihu (Linux, AMD EPYC).

Formula

For the draws $s$ in one cluster (or one projected draw), with weights $w_s$ normalized to $\sum_s w_s = 1$, means $\mu_{is}$ and dispersions $\sigma_s$, the predictive variance of observation $i$ is

$$ v_i = c \sum_s w_s \sigma_s^2 + \sum_s w_s \left(\mu_{is} - \bar{\mu}_i\right)^2 , \qquad \bar{\mu}_i = \sum_s w_s \mu_{is} , $$

with $c = 1$ for the Gaussian family and $c = \nu / (\nu - 2)$ for the Student-t family. The first term is the same for all $i$.

Master

for i in 1..n:
  v_mu[i] <- sum(w * (mu[i, ] - weighted.mean(mu[i, ], w))^2)

get_refdist() calls predvar() once per projected draw, and get_p_clust() once per cluster. So one projection makes $n \cdot S$ R-level calls of weighted.mean(). For $n = 10000$ and $S = 4000$, these are $4 \cdot 10^7$ calls.

This PR

With $M$ the $n \times S_c$ matrix of $\mu_{is}$ for the $S_c$ draws in the cluster and $w$ the weight vector:

$$ \bar{\mu} = M w , \qquad v^{\mu} = \left( (M - \bar{\mu} \mathbf{1}^\top) \circ (M - \bar{\mu} \mathbf{1}^\top) \right) w , $$

where $\circ$ is the elementwise product. These are two matrix-vector products with $O(n S_c)$ operations, as in master, but in compiled code.

mu_mean <- drop(mu %*% w)
v_mu    <- drop(((mu - mu_mean)^2) %*% w)   # mu - mu_mean recycles down columns

For one draw ($S_c = 1$, $w = 1$), $v^{\mu} = 0$ exactly. A missing $\mu_{is}$ gives a missing $v_i$, as in master.