Improved numerical accuracy for log_sum_exp, softmax, and log_softmax

[helske]

Summary

This PR concerns #2802, I already started to work on this on Stancon hackathon with @SteveBronder.

The current primitives of softmax, log_softmax and log_sum_exp are numerically stable according to this paper. However, the fwd and rev specializations are numerically less stable, and their handling of non-finite values is not consistent with the primitives.

The new fwd and rev specializations of these functions rely on the stable softmax and log_sum_exp primitives rather than independently reimplementing them. To make this possible, I relaxed the require_* constraints on the softmax and log_softmax primitives so that they support also also matrix input (operating over all elements of the matrix). If this is not a good idea, alternatively fwd and rev could use reshaping/flattening or reimplement the stable softmax (see also below).

There is still potential room for improvement:

1. fwd::log_softmax, fwd::log_sum_exp, and rev::log_sum_exp now compute both softmax(x) and log_sum_exp(x) which have overlap in computations which could be avoided by computing both in a single pass.
2. While the paper linked above concludes that the current max + std::log((v_ref.array() - max).exp().sum()); is safe, their suggested algorithm 4.1 excludes the maximum from the sum and uses log1p, which, based on some standalone tests I made could make a difference in some cases. Basically we could do something like

const auto v_flat = v_ref.reshaped(); 
Eigen::Index k; 
const double max = v_flat.maxCoeff(&k); 
if (!std::isfinite(max)) { 
 return max; 
}
Eigen::ArrayXd w = (v_flat.array() - max).exp();
w(k) = 0.0; 
return max + std::log1p(w.sum());

Tests

The new tests covers numerical stability and non-finite-value handling in the fwd and rev specializations of log_softmax and log_sum_exp. The tests cover inputs with large dynamic ranges, many small probabilities, and infinite values. These tests fail with the current implementations and pass with the proposed changes. At least on my setup, few of the tests have tolerances of 1e-12 which I'm not sure are actually portable. For writing the tests, I used help from LLM.

Side Effects

I guess this now exposes support for matrices in softmax and log_softmax?

Release notes

Improved numerical accuracy of log_sum_exp, softmax and log_softmax functions.

Checklist

• Copyright holder: Jouni Helske

  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

[WardBrian]

Hi @helske -- it looks like the actual changes to the source may be missing from your branch, only the test changes seem to have been checked in

[helske]

Yes I run out of time earlier today. Here's the actual changes.

Note: I did not change the current two-scalar versions of fwd::log_sum_exp(x1, x2) and rev::log_sum_exp(x1, x2), but I think there is still a problem here as well, or at least inconsistency with the vector valued version:

template <typename T>
inline fvar<T> log_sum_exp(const fvar<T>& x1, const fvar<T>& x2) {
  return fvar<T>(log_sum_exp(x1.val_, x2.val_),
                 x1.d_ * inv_logit(-(x2.val_ - x1.val_))
                     + x2.d_ * inv_logit(-(x1.val_ - x2.val_)));
}

If x1 is infinite and x2 is finite, this returns infinite value but potentially finite derivative, whereas log_sum_exp(T&& x) version always returns NaN? But I'm not sure if this matters in practice at all given nonfinite value in both cases?