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?

[andrjohns]

Thanks for the contribution! Looking pretty good, just some extra simplifications you can take advantage of and a couple of performance improvements

[andrjohns]

Suggested change

	const auto probs = softmax(vals);
	return fvar<T_fvar_inner>(log_sum_exp(vals),
	v.d().cwiseProduct(probs).sum());
	return fvar<T_fvar_inner>(log_sum_exp(vals),
	v.d().cwiseProduct(softmax(vals)).sum());

[andrjohns]

A couple of 'gotchas' to be aware of. When calling .val()/value_of()/.d() on a matrix of fvar types, the resulting expression is essentially a view over non-contiguous memory. This makes Eigen expressions less efficient and vectorisable, so if you need to use them multiple times then it's generally more performant to pay the copy-cost once so that the following expressions will be faster.

The size of expressions can be checked without triggering an evaluation (correct me if I'm wrong there @SteveBronder) so you can use if (x.size() == 0) for the early return, since to_ref() can evaluate expressions.

There are also a few helper functions in the math library for simplifying the construction which you could use. Putting it all together, a simplified implementation could look like:

template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
inline plain_type_t<Mat> log_softmax(Mat&& x) {
  if (x.size() == 0) {
    return {};
  }
  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
  const auto x_val = value_of(x_ref).eval();
  const auto d_in = x_ref.d().eval();
  return to_fvar(x_val.array() - log_sum_exp(x_val),
                 d_in.array() - softmax(x_val).cwiseProduct(d_in).sum());
}

[andrjohns]

Using the same approach from log_softmax, this could simplify down to:

template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
inline plain_type_t<Mat> softmax(Mat&& x) {
  if (x.size() == 0) {
    return {};
  }
  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
  const auto s = softmax(value_of(x_ref)).eval();
  const auto d_in = x_ref.d().eval();

  return to_fvar(s, s.array() * (d_in.array() - s.cwiseProduct(d_in).sum()));
}

[andrjohns]

Suggested change

	/*
	* log_softmax(x) = x - log_sum_exp(x).
	*
	* exp(-1000) underflows to zero in double precision. Therefore:
	*
	* log(softmax(x)[2])
	*
	* gives -inf while
	*
	* x[2] - log_sum_exp(x)
	*
	* correctly retains the log-probability.
	*/

Comment isn't really necessary, since the test title already specifies that it's testing a large range

[andrjohns]

Suggested change

	EXPECT_TRUE(std::isfinite(y.val()(0)));
	EXPECT_TRUE(std::isfinite(y.val()(1)));
	EXPECT_TRUE(std::isfinite(y.val()(2)));

Will be covered by the value tests below

[andrjohns]

You can take advantage of some of utility/helper functions for matrix equalities to simplify these tests. For example, this test could be reduced down to:

TEST(AgradFwd, log_softmax_large_dynamic_range) {
  using stan::math::fvar;
  using stan::math::vector_fd;
  using stan::math::log_softmax;

  Eigen::VectorXd val(3);
  val << 0, -100, -1000;

  Eigen::VectorXd d(3);
  d << 1, 2, 3;

  vector_fd x = stan::math::to_fvar(val, d);

  auto y = log_softmax(x);

  EXPECT_MATRIX_EQ(y.val(), val);
  EXPECT_MATRIX_EQ(y.d(), d.array() - (1.0 + 2.0 * std::exp(-100.0)));
}

This also applies to the other tests in the PR, so I won't duplicate the review comments there

[andrjohns]

You can remove all comments from this one since the test name covers it