Embedded Laplace: compute only the entries of A that compute_s2() reads

[avehtari]

This PR was Claude assisted

Summary

When a likelihood argument is a var (for example eta in laplace_marginal_neg_binomial_2_log_lpmf, or the mean vector), laplace_marginal_density() builds an n × n matrix A for laplace_likelihood::compute_s2() and evaluates v = K s2 - K R K s2. Both steps cost O(n³), and most of what they compute is never used.

1. Only the diagonal blocks of A. compute_s2() reads A(k + j b, i + j b) with i, k < b, where b = hessian_block_size: the diagonal blocks of size b, or only the diagonal when b = 1. The current code builds the full matrix, as covariance_val - C.transpose() * C for solver 1, C.transpose() * C for solver 2 and covariance_val - covariance_val * W_r * LU_solve_covariance for solver 3. A new helper, internal::block_diagonal_entries(), fills only the diagonal blocks in O(n² b) and leaves the other entries at zero. The branch without var likelihood arguments needs only the diagonal of C.transpose() * C, which it now gets from C.colwise().squaredNorm().transpose(); that change made no measurable difference.

2. v from the right. Eigen evaluates covariance_val * s2 - covariance_val * R * covariance_val * s2 from the left, with two n × n × n products. Written as covariance_val * s2 - covariance_val * (R * (covariance_val * s2)), it costs O(n²).

I timed the gradient of laplace_marginal_neg_binomial_2_log_lpmf with respect to eta and the 2 hyperparameters of a squared-exponential GP covariance, with eta a var, 2 observations per group, hessian_block_size = 1 and solver 1 (g++ 14.3, -O3). Each entry is the time of this PR divided by the time of develop, for two runs:

n	10	50	100	200	500	1000
time ratio	0.98, 0.98	0.94, 0.94	0.88, 0.89	0.75, 0.76	0.64, 0.67	0.60, 0.61

On develop at n = 1000, this gradient took 1.7 times as long as the gradient with eta as data; with the PR the two take about the same time. With eta as data, the value and the gradient run as fast as on develop (ratios 0.94 to 1.03). Values and gradients agree with develop to a relative difference of 3.2e-11 or less.

Tests

The new file test/unit/math/laplace/laplace_block_diagonal_entries_test.cpp has two tests. keeps_only_the_diagonal_blocks checks that block_diagonal_entries() copies the diagonal blocks of a dense matrix and sets every other entry to zero, for block sizes 1, 2 and 3. compute_s2_reads_only_the_diagonal_blocks checks the property that the change depends on: with a Poisson likelihood, compute_s2() returns identical values for the full matrix and for its diagonal blocks.

The existing tests already run the changed code. laplace_marginal_neg_binomial_log_lpmf_test.cpp, for example, runs expect_ad over the GP hyperparameters and eta for solvers 1 to 3 and Hessian block sizes 1 and 2.

Side Effects

compute_s2() still receives an n × n matrix, with zeros outside the diagonal blocks, so memory use does not change. Results differ from develop only by rounding.

Release notes

The gradient of the embedded Laplace approximation with respect to likelihood parameters, such as the negative binomial dispersion or the mean, no longer computes up to three O(n³) matrix products. With solver 1 and n = 1000 it is about 40 % faster.

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.23	0.22	1.01	0.52% faster
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.06	0.06	1.01	1.18% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	6.31	6.31	1.0	-0.09% slower
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.02	0.02	0.98	-1.65% slower
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	8.25	8.25	1.0	0.0% slower
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	4.44	4.48	0.99	-0.91% slower
stat_comp_benchmarks/benchmarks/sir/sir.stan	217.79	217.36	1.0	0.2% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.9	0.9	1.0	-0.06% slower
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.61	0.61	1.0	0.26% faster
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	43.02	42.61	1.01	0.96% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.59	0.59	0.99	-0.92% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	20.91	20.89	1.0	0.09% faster
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.11	0.11	1.0	-0.12% slower
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.18	3.17	1.0	0.47% faster
performance.compilation	368.64	367.73	1.0	0.25% faster
Mean result: 1.0001678202169857

Jenkins Console Log
Jenkins Build Stages
Commit hash: a2236819d562837e97a9bfe97dc9ebd6d8613729

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           43 bits physical, 48 bits virtual
CPU(s):                                  256
On-line CPU(s) list:                     0-255
Thread(s) per core:                      2
Core(s) per socket:                      64
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              23
Model:                                   49
Model name:                              AMD EPYC 7742 64-Core Processor
Stepping:                                0
Frequency boost:                         enabled
CPU MHz:                                 1497.187
CPU max MHz:                             3416.0681
CPU min MHz:                             1500.0000
BogoMIPS:                                4491.85
Virtualization:                          AMD-V
L1d cache:                               4 MiB
L1i cache:                               4 MiB
L2 cache:                                64 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-63,128-191
NUMA node1 CPU(s):                       64-127,192-255
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Retpolines; IBPB conditional; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sev sev_es

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin