[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]

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]

@helske apologies for the delay in wrapping this up! Can you fix up the cpplint issues and land any other changes you'd like, then I think this should be about good to go

[helske]

No worries, I got occupied with other stuff but I'm aiming to finish this by the end of the week.

[avehtari]

@helske @andrjohns

The current PR is 1.1 to 1.7 times slower than develop.

Time can be improved to 0.5 to 0.9 compared to develop with two changes:

1. Use a helper that returns log_sum_exp and softmax with one exp per element . It covers your item 1 and the duplicate computation that @andrjohns asked about on 2026-09-15.
2. Evaluate theta once in prim softmax: const auto theta = (v.array() - v.maxCoeff()).exp().eval();. Prim softmax then takes 0.63 to 0.85 of its present time, and the accuracy stays the same.

Also, it would be good to fix OpenCL part at the same time, as the OpenCL reverse mode has the same develop defect.

I can also make commits for these if needed.

[helske]

I made further changes, adding the helper function for softmax and log-sum-exp etc. While trying to commit the changes I ended up having some conflicts and diverging branches which I tried to fix with LMM, but it ended up messing things further, so I'm not sure if something got broken in the process.

Please check especially the softmax_log_sum_exp.hpp, I wasn't quite sure what is the correct style for these. I took an existing helper function file as a template and used Claude here.

I didn't touch the OpenCL files, not sure about the best way to modify those as I'm familiar with the syntax used on those.

[avehtari]

@helske , after you have CPU side finished, you can ask Claude (e.g. Opus 5.5) to implement the same in OpenCL as then it's just porting the algorithms and can re-use values from tests compare to CPU resutls, so it will get it correct

[avehtari]

@helske, the new version is better. The history rewrite did not break anything that I could find. On the new head the tests pass and accuracy good. I have two small suggestions for the CPU side speed

1. Rev log_softmax still recomputes softmax in the reverse pass, so it stays 1.07 to 1.34 times slower than develop. It can use your helper in the forward pass and keep the softmax in the arena, as your rev log_sum_exp does:

using val_t = plain_type_t<decltype(x_arena.val())>;
using return_t = return_var_matrix_t<val_t, T>;
const auto& x_val = to_ref(x_arena.val());
arena_t<val_t> arena_softmax;
const double lse = internal::softmax_log_sum_exp(x_val, arena_softmax);
arena_t<return_t> res = (x_val.array() - lse).matrix();
reverse_pass_callback([x_arena, res, arena_softmax]() mutable {
  const auto& res_adj = to_ref(res.adj());
  x_arena.adj().array()
      += res_adj.array() - res_adj.sum() * arena_softmax.array();
});
return res;

2. The snippet above uses a second overload of the helper that writes the softmax into a given output. In reverse mode that output is an arena matrix, so the result does not need a second copy with to_arena. Your pair version calls it, so the fwd code stays as it is:

template <typename T, typename Out,
          require_eigen_vt<std::is_arithmetic, T>* = nullptr>
inline double softmax_log_sum_exp(const T& x, Out& softmax_x) {
  if (x.size() == 0) {
    softmax_x = x;
    return NEGATIVE_INFTY;
  }
  const auto& x_ref = to_ref(x);
  const double max = x_ref.maxCoeff();
  softmax_x = (x_ref.array() - max).exp().matrix();
  const double theta_sum = softmax_x.sum();
  softmax_x /= theta_sum;
  return std::isfinite(max) ? max + std::log(theta_sum) : max;
}

template <typename T, require_eigen_vt<std::is_arithmetic, T>* = nullptr>
inline std::pair<plain_type_t<T>, double> softmax_log_sum_exp(const T& x) {
  plain_type_t<T> softmax_x;
  const double lse = softmax_log_sum_exp(x, softmax_x);
  return {std::move(softmax_x), lse};
}

Rev log_sum_exp then passes arena_t<plain_type_t<decltype(arena_v.val())>> arena_softmax; instead of calling to_arena(softmax_lse.first). The var_value overload does the same with x.val().

With both additions, every reverse-mode path is faster than develop for n = 10 to 1000. Rev log_softmax changes from 1.07 to 1.23 times develop to 0.73 to 0.93 times develop, which is 0.61 to 0.82 of the time of your current version. Rev log_sum_exp takes 0.81 to 0.96 of your current time. I timed one core of an AMD EPYC 9965 with g++ 11.5 and -O3, median of 5 rounds. The values are bitwise the same as before; a few var_value adjoints differ in the last bits.

3. OpenCL has the same reverse-mode defect as develop: opencl/rev/log_sum_exp.hpp uses exp(A.val() - res.val()) and opencl/rev/log_softmax.hpp uses exp(res.val()). Now that the CPU side is fixed, CPU and OpenCL give different adjoints for large inputs. I have a port that computes theta = exp(A - max) and its partial sums in one kernel and keeps theta in an arena_matrix_cl (as opencl/rev/variance.hpp does). It passed the OpenCL tests on a V100 against the earlier head of this PR. I can add it to this PR, or you can ask Claude to make it for you.

A small one: the .gitignore change (.positai) looks unrelated to this PR.