Improve floating point accuracy

[avehtari]

Herbie https://herbie.uwplse.org/ checks floating point computations, detects inaccurate expressions and finds more accurate replacements. I used Herbie with Sol to find these fixes. I have checked that all of them make sense.

Summary

• loo.R: Compute pointwise ELPD Monte Carlo standard errors in log-relative
  coordinates, preventing underflow and overflow for extreme log likelihoods.

• E_loo.R, loo_moment_matching.R: Calculate weighted variances from centered observations instead of
  subtracting squared moments.

• loo_subsample.R: Factor differences of squares in subsampling estimators to avoid
  cancellation for nearby large values.

• psis.R: Use expm1() when calculating PSIS tail exceedances from nearby log weights.

• loo_model_weights.R: Use a sign-aware expm1() formulation for stacking gradients when model
  predictions are nearly identical.

• Added focused regression tests for all five numerical issues.

• Ran the complete test suite: 1168 passed, 0 failed.

There are three new helper functions which are used once or twice, and they could be also inlined, but I think using them improves the readability.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 85.56701% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.52%. Comparing base (bb0760c) to head (32fad02).

Files with missing lines	Patch %	Lines
R/loo.R	57.57%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #395      +/-   ##
==========================================
- Coverage   92.86%   92.52%   -0.34%     
==========================================
  Files          31       31              
  Lines        3041     3105      +64     
==========================================
+ Hits         2824     2873      +49     
- Misses        217      232      +15     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 32fad02 is merged into master:

• 🚀loo_function: 1.95s -> 1.91s [-2.37%, -1.7%]
• 🚀loo_matrix: 1.74s -> 1.63s [-6.81%, -5.98%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[jgabry]

Cool, I'll take a look soon

[avehtari]

Sol generated summary of the excpected benefits

1. ELPD MCSE calculation — high benefit

Affected code: mcse_elpd() in loo/R/loo.R.

Herbie improved the scalar log-normal variance transformation from 35% to 100%.

The original implementation converted likelihoods and expectations to linear scale. For sufficiently
negative log likelihoods, both became zero and produced 0 / 0:

log likelihoods: -1000, -1001
original result: NaN
new result:      finite

The new calculation operates on log-relative likelihoods and applies stable log1p()/expm1()
transformations.

Expected benefits:

• Prevents MCSE values from becoming NaN for valid finite log likelihoods.
• Avoids overflow for large positive log likelihoods.
• Improves accuracy when the variance-to-mean-squared ratio is small.
• Preserves ordinary results to floating-point tolerance.

2. Weighted variances — high benefit for large offsets

Affected code:

• .wvar() in loo/R/E_loo.R
• shift_and_scale() in loo/R/loo_moment_matching.R

Herbie could not improve the abstract scalar expression because it saw the first and second moments
as independent inputs. This is a limitation of scalar expression analysis.

The empirical test exposed complete cancellation:

x:                     1e8 + c(-1, 0, 1)
original variance:     0
centered variance:     1

The new implementation calculates weighted squared deviations from the weighted mean.

Expected benefits:

• Prevents zero or negative variances for nearly constant values with large offsets.
• Prevents moment matching from producing zero, NaN, or inaccurate scale transformations.
• Has little effect on well-centered, ordinary-sized parameters.

3. Subsampling difference of squares — high local benefit

Affected code: srs_diff_est() in loo/R/loo_subsample.R.

The broad Herbie sample scored both expressions at 100%, because random inputs rarely contain nearby
large values. A focused nearby-input analysis selected the factored/FMA form.

For adjacent doubles near 1e8:

high-precision result: 2.98023223876953147
original x² - y²:      2.0
factored result:       2.9802322387695312

original relative error: 32.9%
factored relative error: ~7.5e-17

Expected benefits:

• Retains small differences between large approximation and observation values.
• Improves the finite-population variance calculation used by subsampling.
• Most relevant when the approximation is accurate and both values have large magnitude.

4. PSIS tail exceedances — moderate to high benefit

Affected code: psis_smooth_tail() in loo/R/psis.R.

Herbie improved the analyzed exponential-difference kernel from 11% to 100%.

For nearby log weights around -30:

original relative error: 2.29e-4
stable relative error:   4.60e-17

The stable helper computes

[
\exp(x)-\exp(y)
]

as

[
-\exp(x)\operatorname{expm1}(y-x),
]

with x >= y.

Expected benefits:

• Preserves the smallest positive GPD exceedances.
• Reduces distortion in the fitted PSIS tail when values cluster near the cutoff.
• Handles cases where exp(cutoff) underflows but a larger tail value remains representable.
• May slightly change Pareto-(k) estimates in difficult, tightly clustered tails.

5. Stacking gradients — moderate benefit

Affected code: stacking_weights() in loo/R/loo_model_weights.R.

Herbie strongly favored an expm1() representation for nearby model predictions. Under the analyzed
domain, the score improved from 11% to 100%.

For nearly identical model log predictive densities:

original relative error: 5.7e-3
stable relative error:   ~5.8e-17

The stable helper computes

[
\frac{\exp(a)-\exp(b)}{\exp(c)}
]

without directly subtracting nearby exponentials.

Expected benefits:

• Produces more accurate gradients for models with nearly identical predictions.
• Reduces optimizer noise near flat regions of the stacking objective.
• Can improve convergence and reproducibility when models receive similar weights.
• The benefit is small when model predictions are clearly separated.

[jgabry]

@avehtari The benchmark results look a bit concerning. That's a pretty large slowdown:

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 572ae71 is merged into master:

• ❗🐌loo_function: 1.93s -> 2.31s [+19.03%, +20.07%]
• ❗🐌loo_matrix: 1.79s -> 2.02s [+11.98%, +12.76%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

Also, I was wondering if there were downsides to including these changes, so I asked Codex to review the PR looking for unintended consequences like performance regressions or other implications for existing functions and it found several things. I looked into its findings and verified most of them myself too. They're not all equally important, but some seem quite important. There are a few things that we should probably have tests for that would have caught some of these things beyond what the performance benchmarks already caught. Codex also found a pre-existing moment matching bug (see end of review) that we should probably fix as part of this PR. Here's what Codex found:

---

The core numerical reformulations are mathematically sound for ordinary finite inputs, but the new helpers introduce two regressions for supported -Inf values. The most concerning failure is stacking silently returning incorrect uniform weights.

Blocking issues

1. PSIS now errors for valid -Inf log ratios — R/helpers.R:37, called from R/psis.R:257.

   x <- c(rep(-Inf, 90), seq(-9, 0, length.out = 10))
   psis(x)

   The PR errors inside posterior::gpdfit(); the base branch returns a PSIS object. When x == y == -Inf, y - x is NaN, so -exp(x) * expm1(y - x) is NaN.

   Equal inputs, including equal infinities, need to return zero. This should have a regression test through psis(), not only the helper.

2. Stacking can silently return incorrect uniform weights — R/helpers.R:57, called from R/loo_model_weights.R:279.

   lpd <- matrix(c(
     -Inf,  0, -Inf,
       -1, -1,   -1
   ), nrow = 2, byrow = TRUE)
   
   stacking_weights(lpd)

   PR result: approximately (1/3, 1/3, 1/3)
   Base result: approximately (0, 1, 0)

   exp_diff_over_exp(-Inf, -Inf, c) returns NaN. The optimizer then effectively stays at its initial weights without warning. Equal a and b should contribute zero to the gradient.

Required changes

• In R/loo.R:506, the positive branch is still unstable for very small positive z:

  z + log1p(-exp(-z))

  When exp(-z) rounds to one, this loses the positive contribution. A symmetric example with z = 5e-17 underestimates MCSE by about 29%, although the absolute MCSE is tiny. Use:

  z + log(-expm1(-z))

• Add a NEWS.md entry. These are user-visible correctness changes and should be recorded.

Computational cost

Yes, there is a measurable cost, primarily from mcse_elpd():

• The PR’s Touchstone benchmark on a 4,000 × 1,359 matrix reports:

  • loo.matrix(): 1.79 s → 2.02 s, roughly 12–13% slower
  • function-method loo(): 1.93 s → 2.31 s, roughly 19–20% slower

• Isolated mcse_elpd() benchmarks were around 1.8–2.6× slower. It does, however, reduce peak memory by processing columns individually.

• shift_and_scale() was roughly 20–30% slower for large parameter matrices, but only affects moment matching.

• Actual PSIS runtime was essentially unchanged; Pareto fitting dominates the added expm1() cost.

• Subsampling and .wvar() showed no meaningful regression.

• Stacking becomes substantially faster with many models because the mixture log density is no longer recomputed inside the loop. With 10 models, the gradient was about 2.2× faster; with 50 models, about 8× faster. Two-model stacking was approximately 20% slower.

So the end-to-end slowdown is moderate but real—not prohibitive, though worth explicitly accepting or optimizing.

Adjacent pre-existing issue

R/loo_moment_matching.R:411 passes exp(elpd_loo_i) to mcse_elpd(), whose contract expects the log-scale elpd_loo_i. For a simple example, current output is 0.466511 versus the correct 0.4039146.

This bug exists on both base and PR branches, so I did not count it as a regression, but this PR is a natural place to fix it and add coverage.

All focused tests and current CI pass; the failures above require boundary cases not presently tested.

[avehtari]

Based on your Codex analysis the slowdown is not really a practical issue, but the blocking ones are clear. It's clear that our tests have been weak since these were not caught. I'll fix

[jgabry]

Based on your Codex analysis the slowdown is not really a practical issue

To me loo() >10% slower seems like it could be a practical issue in cases when it's already slow. Or perhaps I misunderstood what the timing referred to?

It's clear that our tests have been weak since these were not caught.

Yeah this was useful also because it pointed out a few holes in our tests that we should add coverage for

[avehtari]

All Codex issues fixed and some more. Speed might be faster now, but let's see Touchstone results

• Fix loo_moment_match() reporting mcse_elpd_loo on the wrong scale:
  loo_moment_match_i() passed exp(elpd_loo_i) to the internal mcse_elpd(),
  which already expects the log scale.
• waic() now errors on -Inf log-likelihood values through the matrix, array
  and function interfaces, because the pointwise variance is undefined. These
  previously returned NaN
• psis(), tis() and sis() accept -Inf log ratios, giving those draws zero
  weight, and reject only columns with no finite value at all.
• loo() now reports -Inf log likelihoods as "-Inf log-likelihood values are not
  allowed." instead of the internal log-ratio message "All input values must be
  finite or -Inf.".
• stacking_weights() errors when an observation has no finite log predictive
  density under any model, and pseudobma_weights() errors when no model has a
  finite total. pseudobma_weights() gives a model with -Inf total ELPD zero
  weight instead of NaN.
• psis_approximate_posterior() and ap_psis() reject undefined log density
  ratios, +Inf log ratios, and all--Inf log-ratio columns.
• mcse_elpd() is computed relative to the loo predictive density, so the added
  numerical stability costs no run time: loo() is slightly faster than 2.10.1 and
  uses about 40% less memory.

[jgabry]

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 32fad02 is merged into master:

• 🚀loo_function: 1.95s -> 1.91s [-2.37%, -1.7%]
• 🚀loo_matrix: 1.74s -> 1.63s [-6.81%, -5.98%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

Yeah speed looks much better! I'll take a look at the other changes soon.