Compute GLM design-matrix products with multiply to avoid per-coefficient autodiff nodes in forward mode

[jachymb]

AI use disclosure: The code as well as this commentary was done with the help of claude Fable 5.1

Summary

The GLM log densities compute their two design-matrix products, x * beta and x^T * theta_derivative, with multiply instead of Eigen's operator*. With forward-mode coefficients (Hessian-vector products, the third-order term of the embedded Laplace approximation) this puts one autodiff node per result coefficient on the tape instead of two per matrix entry, which makes such sweeps several times faster.

Why it matters: the partials are computed in partials_return_t, which is var for fvar<var> inputs, and Eigen evaluates a double-by-var product coefficient by coefficient, 2 n K nodes per product on every sweep. multiply is resolved by argument-dependent lookup at instantiation (reverse-mode overload for var partials, forward-mode overload for fvar<var> partials), so prim gains no dependency on rev. For double and var inputs the prim overload returns the same lazy product as before.

categorical_logit_glm_lpmf with a matrix beta did not compile with fvar<var> at all (Eigen's GEMM kernel cannot mix scalar types), which the existing one-column test never exercised; its x-partials products are switched as well.

Measured on an embedded-Laplace hierarchical GLM (49 coefficients, 30 to 400 observations per group): 214 to 22 nodes per observation per second-order sweep, 508 to 131 per third-order sweep, laplace_marginal 3.1 to 3.8 times faster. Values and gradients match develop to 12 digits at every autodiff order.

Tests

• test/unit/math/mix/prob/glm_forward_mode_tape_test.cpp (new): one fvar<var> sweep of five GLM densities must add fewer than 40 n nodes to the tape; before the change it added about 4 n K + 20 n.
• categorical_logit_glm_lpmf_test.cpp: new three-category expect_ad case, which did not compile before.
• The existing prim and mix GLM tests are unchanged.

Side Effects

None for double and var arguments. With forward-mode arguments the reverse-mode multiply copies the design matrix into the arena once per call (n K doubles), which is what replaces the 2 n K node allocations.

Release notes

GLM log densities evaluated with forward-mode scalars (nested autodiff, for example the embedded Laplace approximation) now create O(n) instead of O(n K) autodiff nodes per sweep. categorical_logit_glm_lpmf now compiles with forward-mode coefficients for more than one category.

Checklist

• Copyright holder: Jachym Barvinek

  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

[jachymb]

Not sure why the CI is failing, I think jenkins is running out of memory or something like that.

[SteveBronder]

Yeah we just moved our jenkins over to a new set of servers recently so there are some hiccups I'm looking at it now

[SteveBronder]

I think the meta here is that we really want Matrix * Matrix etc. stuff to dispatch correctly to multiply so we can stop calling it manually everywhere. I'm fine with this pr, but let me take a crack at seeing what we would need to specialize in Eigen so that var matrices can use our multiply and add functions etc.

[jachymb]

Yes, I agree, this feels a bit like hacking and it's quite non-obvious to a new developer why this is even a semantic difference in the first place. If there was a more conceptual way to improve this that would be great, but it's probably outside my ability to address at this point.

I just noticed this is a way I could speed-up an application I'm building on top of stan that has a hierarchical GLM model, so this has a practical motivation for me and I went for the approach I could do myself.

[SteveBronder]

Yes tbc this is a good PR and once I restart the jenkins I think this is good to merge as is. We have not done the matrix multiply overloads because it requires messing with a lot of Eigens internals. I'm taking a crack at that, but for now I think what you have here is best