9  Moment matching for improved PSIS-LOO-CV

In this chapter, we demonstrate how to apply importance weighted moment matching (IWMM) to improve PSIS-LOO-CV estimates when influential observations produce high Pareto-\(k\) values. We walk through a complete applied example using the roaches dataset from Gelman and Hill (2007), first letting ArviZ handle everything automatically from the fitted model, and then building the required functions ourselves. In both cases we resolve problematic importance sampling approximations without expensive model refitting. The workflow demonstrates ArviZ’s loo_moment_match() function with a deliberately misspecified Poisson regression model where multiple observations are flagged as problematic. For theoretical background and details of how moment matching works, we recommend you read Section 7.7 and references therein.

9.1 Roaches data and Poisson regression model

Let’s now walk through a concrete example using the roaches dataset to see this algorithm in action. The roaches dataset from Gelman and Hill (2007) examines the efficacy of a pest management system at reducing cockroach infestations in urban apartments. The study followed 264 apartments over several months, recording the number of roaches caught during follow-up (y), pre-treatment roach counts (roach1, square root transformed), treatment status (treatment), whether the building is restricted to elderly residents (senior), and trap exposure time in days (exposure2).

We intentionally use a Poisson model rather than negative binomial regression to demonstrate how moment matching handles misspecified models. The exposure time varies across apartments, so we include log(exposure2) as an offset term.

9.1.1 Loading and preparing the data

We start by loading the roaches dataset and examining its basic structure.

import pandas as pd

roaches = pd.read_csv("../data/roaches.csv", index_col=0)
roaches['log_exposure2'] = np.log(roaches['exposure2'])

roaches.describe()

	y	roach1	treatment	senior	exposure2	log_exposure2
count	262.000000	262.000000	262.000000	262.000000	262.000000	262.000000
mean	25.648855	4.419861	0.603053	0.305344	1.021047	-0.012538
std	50.846539	4.769184	0.490201	0.461434	0.320757	0.249755
min	0.000000	0.000000	0.000000	0.000000	0.200000	-1.609438
25%	0.000000	1.000000	0.000000	0.000000	1.000000	0.000000
50%	3.000000	2.645751	1.000000	0.000000	1.000000	0.000000
75%	24.000000	7.106071	1.000000	1.000000	1.000000	0.000000
max	357.000000	21.213203	1.000000	1.000000	4.285714	1.455287

The dataset exhibits substantial overdispersion with mean 25.6 and standard deviation 50.8. Approximately 36% of apartments had zero roaches while the maximum reached 357. This combination of high variability, many zeros, and extreme counts makes certain observations potentially influential, which is precisely where moment matching becomes valuable for improving LOO-CV estimates.

9.1.2 Model specification and fitting

We fit a Poisson regression model with the pre-treatment roach count, treatment indicator, and senior building indicator as predictors, and log exposure as an offset.

PyMC

import pymc as pm

coords = {"obs_id": np.arange(len(roaches))}

with pm.Model(coords=coords) as model:
    roach1 = pm.Normal("roach1", mu=0, sigma=2.5)
    treatment = pm.Normal("treatment", mu=0, sigma=2.5)
    senior = pm.Normal("senior", mu=0, sigma=2.5)
    intercept = pm.Normal("Intercept", mu=0, sigma=5.0)

    eta = (
        intercept
        + roach1 * roaches["roach1"]
        + treatment * roaches["treatment"]
        + senior * roaches["senior"]
        + roaches["log_exposure2"]
    )

    pm.Poisson("y", mu=pm.math.exp(eta), observed=roaches["y"], dims="obs_id")

    idata = pm.sample(
        draws=1000,
        tune=1000,
        chains=4,
        random_seed=SEED
    )
    pm.compute_log_likelihood(idata)

Bambi

import bambi as bmb

model = bmb.Model(
    'y ~ roach1 + treatment + senior + offset(log_exposure2)',
    data=roaches,
    family='poisson',
    priors={
        'roach1': bmb.Prior('Normal', mu=0, sigma=2.5),
        'treatment': bmb.Prior('Normal', mu=0, sigma=2.5),
        'senior': bmb.Prior('Normal', mu=0, sigma=2.5),
        'Intercept': bmb.Prior('Normal', mu=0, sigma=5.0)
    }
)

idata = model.fit(
    draws=1000,
    tune=1000,
    chains=4,
    random_seed=SEED,
    idata_kwargs={'log_likelihood': True}
)

CmdStanPy

from cmdstanpy import CmdStanModel

stan_code = """
data {
  int<lower=0> N;
  matrix[N, 3] X;
  array[N] int<lower=0> y;
  vector[N] log_exposure2;
}
parameters {
  real roach1;
  real treatment;
  real senior;
  real Intercept;
}
model {
  roach1 ~ normal(0, 2.5);
  treatment ~ normal(0, 2.5);
  senior ~ normal(0, 2.5);
  Intercept ~ normal(0, 5);
  y ~ poisson_log_glm(X, Intercept + log_exposure2, [roach1, treatment, senior]');
}
generated quantities {
  vector[N] log_lik;
  {
    vector[N] eta = X * [roach1, treatment, senior]' + Intercept + log_exposure2;
    for (n in 1:N)
      log_lik[n] = poisson_log_lpmf(y[n] | eta[n]);
  }
}
"""

with open("roaches_poisson.stan", "w") as f:
    f.write(stan_code)

stan_data = {
    "N": len(roaches),
    "X": roaches[["roach1", "treatment", "senior"]].values,
    "y": roaches["y"].values,
    "log_exposure2": roaches["log_exposure2"].values,
}

stan_model = CmdStanModel(stan_file="roaches_poisson.stan")

fit = stan_model.sample(
    data=stan_data,
    chains=4,
    iter_warmup=1000,
    iter_sampling=1000,
    seed=SEED,
    show_progress=False,
)

idata = az.from_cmdstanpy(
    posterior=fit,
    log_likelihood={"y": "log_lik"},
    observed_data={"y": roaches["y"].values},
    coords={"obs_id": np.arange(len(roaches))},
    dims={"y": ["obs_id"]},
)

az.summary(idata.posterior.ds, var_names=['roach1', 'treatment', 'senior', 'Intercept'])

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
roach1	0.16061	0.00194	0.16	0.16	2569	2468	1.00	3.8e-05	2.8e-05
treatment	-0.5684	0.0251	-0.61	-0.53	3331	2815	1.00	0.00043	0.00031
senior	-0.315	0.0333	-0.37	-0.26	3444	2716	1.00	0.00057	0.00041
Intercept	2.527	0.0259	2.5	2.6	2210	2279	1.00	0.00055	0.00039

9.1.3 Initial PSIS-LOO-CV evaluation

With the posterior draws in hand, we compute the PSIS-LOO-CV estimate of the model’s predictive performance.

loo_result = az.loo(idata, pointwise=True, var_name="y")
loo_result

/opt/hostedtoolcache/Python/3.13.14/x64/lib/python3.13/site-packages/arviz_stats/loo/loo_helper.py:1146: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations.
  warnings.warn(

Computed from 4000 posterior samples and 262 observations log-likelihood matrix.

         Estimate       SE
elpd_loo -5456.03   692.65
p_loo      250.23        -

There has been a warning during the calculation. Please check the results.
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.70]   (good)      245   93.5%
   (0.70, 1]   (bad)        10    3.8%
    (1, Inf)   (very bad)    7    2.7%

The output shows that 17 observations have Pareto-\(k\) values exceeding 0.7, so the importance sampling approximation is unreliable for these cases. For these problematic observations, the most accurate (but computationally expensive) approach would be to refit the model, leaving out each observation in turn. However, this quickly becomes impractical as the number of flagged observations grows. Moment matching offers an efficient solution by improving the reliability of the PSIS-LOO-CV estimates for exactly those cases where brute-force refitting would otherwise be required, reducing computational burden without sacrificing much accuracy.

9.2 Applying moment matching

Now we apply importance weighted moment matching to improve the PSIS-LOO-CV estimates for the problematic observations. ArviZ implements this through the loo_moment_match() function, which needs the posterior draws in the unconstrained parameter space, plus functions that evaluate the log posterior and the log likelihood of a single left-out observation. We can either pass our fitted model and let ArviZ build all of this for us, or provide the pieces ourselves.

9.2.1 Automatic moment matching

The easiest option is to pass the fitted model object through the model argument, which works with both PyMC and Bambi models. ArviZ extracts everything it needs in the background, and we don’t have to write any functions ourselves.

loo_mm_result = az.loo_moment_match(
    idata,
    loo_result,
    model=model,
    var_name="y",
)

loo_mm_result

Computed from 4000 posterior samples and 262 observations log-likelihood matrix.

         Estimate       SE
elpd_loo -5479.52   700.19
p_loo      273.73        -
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.70]   (good)      262  100.0%
   (0.70, 1]   (bad)         0    0.0%
    (1, Inf)   (very bad)    0    0.0%

After moment matching, all 262 observations have Pareto-\(k\) values below 0.7. Every one of the 17 problematic cases has been resolved, and we can now trust the PSIS-LOO-CV estimates for the entire dataset. The ELPD estimate also decreases from -5456.03 to -5479.52, telling us the original estimate was too optimistic and loo() overestimated the predictive performance of the model. The whole adjustment took a few seconds, compared to the minutes or hours we would have spent refitting the model 17 times.

9.2.2 Custom moment matching functions

If we are working with a framework that does not support automatic extraction, or we want full control over the computations, we can write the functions ourselves. Unlike standard PSIS, which only evaluates the likelihood, moment matching needs the full-data posterior density \(p(\theta^{*(s)} \mid y)\) and the likelihood \(p(y_i \mid \theta^{*(s)})\) at each transformed draw \(\theta^{*(s)}\). So we need one function that computes the log posterior \(\log p(\theta \mid y)\) and another that computes the log likelihood for a single observation \(\log p(y_i \mid \theta)\), both in the unconstrained parameter space where all parameters are real-valued.

Our Poisson regression model is

\[ y_i \sim \text{Poisson}(\lambda_i), \quad \lambda_i = \exp(\eta_i), \]

where the linear predictor is

\[ \eta_i = \alpha + \sum_{j=1}^3 \beta_j x_{ij} + \log(\text{exposure2}_i), \]

with \(x_i = (\text{roach1}_i, \text{treatment}_i, \text{senior}_i)\). We use independent normal priors \(\beta_j \sim N(0, 2.5^2)\) and \(\alpha \sim N(0, 5^2)\).

The log posterior combines the log likelihood with the log prior

\[ \log p(\theta \mid y) = \sum_{i=1}^n \log p(y_i \mid \theta) + \log p(\theta), \]

where the log likelihood for observation \(i\) is

\[ \log p(y_i \mid \theta) = y_i \log(\lambda_i) - \lambda_i - \log(y_i!). \]

We start by converting our data into DataArrays with properly labeled dimensions to use xarray’s dimension-aware operations for the computations.

import xarray as xr
from scipy.special import gammaln

coef_names = ['roach1', 'treatment', 'senior']

design_matrix = xr.DataArray(
    roaches[coef_names].values,
    dims=['obs_id', 'coef'],
    coords={'coef': coef_names}
)
y_da = xr.DataArray(roaches['y'].values, dims=['obs_id'])
offset_da = xr.DataArray(roaches['log_exposure2'].values, dims=['obs_id'])
factorial_term = xr.DataArray(gammaln(roaches['y'].values + 1), dims=['obs_id'])

beta_prior_scale = 2.5
alpha_prior_scale = 5.0

First, we construct an array of unconstrained parameters from the posterior draws. All of the parameters are unconstrained in this model, so we can simply stack the posterior draws along a new uparam dimension.

upars = (
    idata.posterior.ds[['roach1', 'treatment', 'senior', 'Intercept']]
    .to_dataarray(dim='uparam')
    .transpose('chain', 'draw', 'uparam')
)

We can now define the log posterior function and the leave-one-out log likelihood function. The log posterior function computes the log posterior probability for a given set of unconstrained parameters, and the log likelihood function computes the log likelihood for a given set of unconstrained parameters and a given observation.

def log_prob_upars(upars):
    """Compute log posterior for unconstrained parameters."""
    beta = upars.sel(uparam=coef_names).rename({'uparam': 'coef'})
    intercept = upars.sel(uparam='Intercept')

    lin = xr.dot(design_matrix, beta, dims='coef') + intercept + offset_da
    log_lik = (y_da * lin - np.exp(lin) - factorial_term).sum('obs_id')

    log_prior = (
        (-0.5 * (beta / beta_prior_scale) ** 2).sum('coef')
        - 0.5 * (intercept / alpha_prior_scale) ** 2
    )
    return log_lik + log_prior


def log_lik_i_upars(upars, i):
    """Compute log likelihood for observation i."""
    beta = upars.sel(uparam=coef_names).rename({'uparam': 'coef'})
    intercept = upars.sel(uparam='Intercept')

    features_i = design_matrix.isel(obs_id=i)
    lin_i = (beta * features_i).sum('coef') + intercept + offset_da.isel(obs_id=i)
    return y_da.isel(obs_id=i) * lin_i - np.exp(lin_i) - factorial_term.isel(obs_id=i)

9.2.3 Running moment matching with custom functions

With the required functions defined, we can now run PSIS-LOO-CV with importance weighted moment matching, this time passing the unconstrained draws and functions explicitly. We will specify that we want to match the covariance structure of the posterior draws given the complicated structure of the data.

Keep in mind that the split argument, which specifies whether to use the split proposal density, is a boolean that defaults to True. This is highly recommended in most applied cases. When split=True, the split proposal density is

\[ g_{\text{split,loo}}(\theta) \propto p(\theta \mid y) + |\mathbf{J}_{T_w}|^{-1}\, p(T_w^{-1}(\theta) \mid y), \]

so each fold mixes the full-data posterior with the transformed draws from the denominator adaptation.

loo_mm_custom = az.loo_moment_match(
    idata,
    loo_result,
    log_prob_upars_fn=log_prob_upars,
    log_lik_i_upars_fn=log_lik_i_upars,
    upars=upars,
    var_name="y",
    cov=True,
)

loo_mm_custom

Computed from 4000 posterior samples and 262 observations log-likelihood matrix.

         Estimate       SE
elpd_loo -5479.52   700.19
p_loo      273.73        -
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.70]   (good)      262  100.0%
   (0.70, 1]   (bad)         0    0.0%
    (1, Inf)   (very bad)    0    0.0%

The estimates are identical to the ones we got by passing the model directly, as they should be. ArviZ builds the equivalent of these functions from the model behind the scenes, so both routes run the same algorithm on the same inputs.

9.2.4 Interpreting diagnostic quantities

After moment matching, two distinct Pareto-\(k\) diagnostics are available. The values in pareto_k reflect the accuracy of the importance sampling approximation after applying the moment matching transformations. These are the values we use to assess whether our PSIS-LOO-CV estimates are reliable. The original values are preserved in influence_pareto_k and serve a different purpose: they indicate how much each observation influences the posterior distribution. An observation with high influence_pareto_k substantially affects the model fit when included, regardless of whether moment matching successfully improved the sampling accuracy.

Moment matching also provides per-observation effective sample size estimates in n_eff_i. These values quantify how many independent samples effectively contribute to each observation’s PSIS-LOO-CV estimate after accounting for the importance weights and MCMC sampling efficiency. Lower effective sample sizes indicate observations where the importance weights are more variable, suggesting greater uncertainty in those particular estimates.

These diagnostic quantities can be accessed directly from the loo_mm_result object as attributes.

9.3 Limitations of moment matching

Moment matching has some important limitations worth keeping in mind:

• It targets only first and second moments, so improvements depend on whether these moments adequately capture differences between proposal and target distributions
• When importance weights have large variance, the computation of weighted moments can become unreliable (mitigated through weight regularization or larger sample sizes)
• The algorithm may fail to find sufficiently helpful transformations when target and proposal distributions differ substantially in tail behavior, correlation structure, or number of modes
• For extremely high-dimensional problems, the most sophisticated transformation (matching full covariance) may become numerically unstable
• The split proposal approximation introduces some inefficiency by placing unnecessary probability mass in regions where the integrand is near its expectation, though this trade-off is least problematic precisely when adaptive methods are most needed

Despite these limitations, moment matching succeeds in most practical applications when the original posterior simulation was reasonably successful, providing both computational efficiency and improved reliability for model assessment without requiring complex tuning or auxiliary assumptions.

Gelman, Andrew, and Jennifer Hill. 2007. Data Analysis Using Regression and Multilevel/Hierarchical Models. Cambridge: Cambridge University Press. https://sites.stat.columbia.edu/gelman/arm/.