Improved numerical accuracy for log_sum_exp, softmax, and log_softmax

.gitignore

@@ -106,3 +106,4 @@ lib/tbb/**
 
 # gdb debug history
 .gdb_history
+.positai

stan/math/fwd/fun/log_softmax.hpp

@@ -4,20 +4,21 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/meta.hpp>
-#include <stan/math/fwd/fun/softmax.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/log_softmax.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of each vector in a container of `fvar` values.
+ * Return the log softmax of each vector or matrix in a container of `fvar`
+ * values.
  *
  * @tparam T `std::vector` whose scalar type is `fvar`
- * @param x container of vectors to transform
+ * @param x container of vectors or matrices to transform
  * @return container of log softmax results
  */
 template <typename T, require_std_vector_st<is_fvar, T>* = nullptr>
@@ -28,29 +29,24 @@ inline auto log_softmax(T&& x) {
 }
 
 /**
- * Return the log softmax of the specified vector of `fvar` values.
+ * Return the log softmax of the specified vector or matrix of `fvar` values.
  *
- * @tparam Vec Eigen vector with `fvar` scalar
- * @param x vector to transform
- * @return log softmax of the vector, or an empty result if the input is empty
+ * @tparam Mat Eigen vector or matrix with `fvar` scalar
+ * @param x vector or matrix to transform
+ * @return log softmax of the vector or matrix, or an empty result if the
+ * input is empty
  */
-template <typename Vec, require_eigen_vector_vt<is_fvar, Vec>* = nullptr>
-inline auto log_softmax(Vec&& x) {
-  using vec = std::decay_t<Vec>;
-  constexpr int Rows = vec::RowsAtCompileTime;
-  constexpr int Cols = vec::ColsAtCompileTime;
-  using T = typename value_type_t<vec>::Scalar;
-  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
-  if (x_ref.size() == 0) {
-    return Eigen::Matrix<fvar<T>, Rows, Cols>{};
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline plain_type_t<Mat> log_softmax(Mat&& x) {
+  if (x.size() == 0) {
+    return {};
   }
-  const auto s = softmax(value_of(x_ref));
-  const auto d_in = x_ref.d();
-  const auto dot_sd = s.dot(d_in);
-  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.size());
-  result.val() = s.array().log().matrix();
-  result.d() = (d_in.array() - dot_sd).matrix();
-  return result;
+  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();
+  const auto softmax_lse = internal::softmax_log_sum_exp(x_val);
+  return to_fvar(x_val.array() - softmax_lse.second,
+                 d_in.array() - softmax_lse.first.cwiseProduct(d_in).sum());
 }
 
 }  // namespace math

stan/math/fwd/fun/log_sum_exp.hpp

@@ -8,6 +8,7 @@
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <cmath>
 #include <vector>
 
@@ -56,11 +57,9 @@ inline auto log_sum_exp(T&& x) {
         using T_fvar_inner = typename value_type_t<decltype(v)>::Scalar;
         using mat_type = Eigen::Matrix<T_fvar_inner, -1, -1>;
         mat_type vals = v.val();
-        mat_type exp_vals = vals.array().exp();
-
-        return fvar<T_fvar_inner>(
-            log_sum_exp(vals),
-            v.d().cwiseProduct(exp_vals).sum() / exp_vals.sum());
+        const auto softmax_lse = internal::softmax_log_sum_exp(vals);
+        return fvar<T_fvar_inner>(softmax_lse.second,
+                                  v.d().cwiseProduct(softmax_lse.first).sum());
       });
 }
 

stan/math/fwd/fun/softmax.hpp

@@ -29,27 +29,21 @@ inline auto softmax(T&& x) {
 /**
  * Return the softmax of the specified vector of `fvar` values.
  *
- * @tparam Vec Eigen vector with `fvar` scalar
- * @param x vector to transform
- * @return softmax of the vector, or an empty result if the input is empty
+ * @tparam Mat Eigen vector or matrix with `fvar` scalar
+ * @param x vector or matrix to transform
+ * @return softmax of the vector, matrix, or an empty result if the input is
+ *   empty
  */
-template <typename Vec, require_eigen_vector_vt<is_fvar, Vec>* = nullptr>
-inline auto softmax(Vec&& x) {
-  using vec = std::decay_t<Vec>;
-  constexpr int Rows = vec::RowsAtCompileTime;
-  constexpr int Cols = vec::ColsAtCompileTime;
-  using T = typename value_type_t<vec>::Scalar;
-  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
-  if (x_ref.size() == 0) {
-    return Eigen::Matrix<fvar<T>, Rows, Cols>{};
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline plain_type_t<Mat> softmax(Mat&& x) {
+  if (x.size() == 0) {
+    return {};
   }
-  const auto s = softmax(value_of(x_ref));
-  const auto d_in = x_ref.d();
-  const auto dot_sd = s.dot(d_in);
-  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.size());
-  result.val() = s;
-  result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
-  return result;
+  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()));
 }
 
 }  // namespace math

stan/math/prim/fun.hpp

@@ -296,6 +296,7 @@
 #include <stan/math/prim/fun/size_mvt.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <stan/math/prim/fun/sort_asc.hpp>
 #include <stan/math/prim/fun/sort_desc.hpp>
 #include <stan/math/prim/fun/sort_indices.hpp>

stan/math/prim/fun/log_softmax.hpp

@@ -12,8 +12,9 @@ namespace stan {
 namespace math {
 
 /**
- * Return the natural logarithm of the softmax of the specified
- * vector, or of each vector in a container.
+ * Return the natural logarithm of the softmax of the specified vector or
+ * matrix, or of each vector or matrix in a container. For a matrix, the
+ * log-softmax is taken over all elements.
  *
  * \f$
  * \log \mbox{softmax}(y)
@@ -35,17 +36,15 @@ namespace math {
  * \right.
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, `std::vector` of doubles,
- *   or nested container whose scalar type is arithmetic
- * @param x vector or container of vectors to transform
+ * @tparam Container type of input: an Eigen vector, Eigen matrix,
+ *   `std::vector` of vectors or matrices, or nested container whose scalar
+ *   type is arithmetic
+ * @param x vector, matrix, or container to transform.
  * @return log softmax of the input, preserving the container structure; an
- *   empty result if any input vector is empty
+ *   empty result if any input vector or matrix is empty.
  */
 template <typename Container, require_st_arithmetic<Container>* = nullptr,
-          require_container_t<Container>* = nullptr,
-          require_not_t<bool_constant<
-              is_eigen<std::decay_t<Container>>::value
-              && !is_eigen_vector<std::decay_t<Container>>::value>>* = nullptr>
+          require_container_t<Container>* = nullptr>
 inline auto log_softmax(Container&& x) {
   return make_holder(
       [](auto&& a) {

stan/math/prim/fun/softmax.hpp

@@ -11,7 +11,9 @@ namespace stan {
 namespace math {
 
 /**
- * Return the softmax of the specified vector, or of each vector in a container.
+ * Return the softmax of the specified vector or matrix, or of each
+ * vector or matrix in a container. For a matrix, the softmax is
+ * taken over all elements.
  *
  * \f$
  * \mbox{softmax}(y)
@@ -38,17 +40,15 @@ namespace math {
  * \end{array}
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, `std::vector` of doubles,
- *   or nested container whose scalar type is arithmetic
- * @param x vector or container of vectors to transform
+ * @tparam Container type of input: an Eigen vector, Eigen matrix,
+ *   `std::vector` of vectors or matrices, or nested container whose scalar
+ *   type is arithmetic
+ * @param x vector, matrix, or container to transform.
  * @return softmax of the input, preserving the container structure; an empty
- *   result if any input vector is empty
+ *   result if any input vector or matrix is empty.
  */
 template <typename Container, require_st_arithmetic<Container>* = nullptr,
-          require_container_t<Container>* = nullptr,
-          require_not_t<bool_constant<
-              is_eigen<std::decay_t<Container>>::value
-              && !is_eigen_vector<std::decay_t<Container>>::value>>* = nullptr>
+          require_container_t<Container>* = nullptr>
 inline auto softmax(Container&& x) {
   return make_holder(
       [](auto&& a) {
@@ -58,7 +58,7 @@ inline auto softmax(Container&& x) {
               if (v.size() == 0) {
                 return v;
               }
-              const auto theta = (v.array() - v.maxCoeff()).exp();
+              const auto theta = (v.array() - v.maxCoeff()).exp().eval();
               return (theta / theta.sum()).matrix();
             });
       },

stan/math/prim/fun/softmax_log_sum_exp.hpp

@@ -0,0 +1,69 @@
+#ifndef STAN_MATH_PRIM_FUN_SOFTMAX_LOG_SUM_EXP_HPP
+#define STAN_MATH_PRIM_FUN_SOFTMAX_LOG_SUM_EXP_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
+#include <cmath>
+#include <utility>
+
+namespace stan {
+namespace math {
+namespace internal {
+
+/**
+ * Return the softmax and the log sum of exponentials of the specified
+ * vector or matrix, using a single evaluation of the exponential per
+ * element. For a matrix, both are taken over all elements.
+ *
+ * With \f$m = \max(x)\f$ and \f$\theta = \exp(x - m)\f$,
+ *
+ * \f$\mbox{softmax}(x) = \theta / \sum_{n=1}^N \theta_n\f$ and
+ * \f$\mbox{log\_sum\_exp}(x) = m + \log \sum_{n=1}^N \theta_n\f$.
+ *
+ * The results match `softmax(x)` and `log_sum_exp(x)`, including for
+ * non-finite input: if \f$m\f$ is not finite, the log sum of exponentials
+ * is \f$m\f$ and the softmax is NaN.
+ *
+ * @tparam T type of Eigen vector or matrix
+ * @param x Eigen vector or matrix
+ * @return pair of the softmax of `x` and the log sum of exponentials of
+ *   `x`. For empty input, returns an empty softmax and negative infinity
+ */
+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) {
+  if (x.size() == 0) {
+    return {x, NEGATIVE_INFTY};
+  }
+  const auto& x_ref = to_ref(x);
+  const double max = x_ref.maxCoeff();
+  const auto theta = (x_ref.array() - max).exp().eval();
+  const double theta_sum = theta.sum();
+  const double lse = std::isfinite(max) ? max + std::log(theta_sum) : max;
+  return {plain_type_t<T>((theta / theta_sum).matrix()), lse};
+}
+
+/**
+ * Return the softmax and the log sum of exponentials of the specified
+ * vector or matrix of autodiff values. This happens with nested autodiff,
+ * e.g. the values of an `fvar<var>` vector, and falls back to the
+ * autodiff specializations of `softmax()` and `log_sum_exp()`.
+ *
+ * @tparam T type of Eigen vector or matrix with autodiff scalar
+ * @param x vector or matrix
+ * @return pair of the softmax of `x` and the log sum of exponentials of `x`
+ */
+template <typename T, require_eigen_vt<is_autodiff, T>* = nullptr>
+inline std::pair<plain_type_t<T>, value_type_t<T>> softmax_log_sum_exp(
+    const T& x) {
+  return {plain_type_t<T>(softmax(x)), log_sum_exp(x)};
+}
+
+}  // namespace internal
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/fun/log_softmax.hpp

@@ -8,15 +8,17 @@
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/log_softmax.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of the specified vector or row vector.
+ * Return the log softmax of the specified vector, row vector, or matrix.
  *
- * @tparam T a `var_value` or Eigen vector/row_vector with `var` scalar
+ * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with
+ *   `var` scalar
  * @param x input
  * @return log softmax of the input, or an empty result if the input is empty
  */
@@ -31,8 +33,8 @@ inline auto log_softmax(T&& x) {
   arena_t<return_t> res = log_softmax(x_arena.val());
   reverse_pass_callback([x_arena, res]() mutable {
     const auto& res_adj = to_ref(res.adj());
-    x_arena.adj().array()
-        += res_adj.array() - res_adj.sum() * res.val().array().exp();
+    const auto s = softmax(x_arena.val());
+    x_arena.adj().array() += res_adj.array() - res_adj.sum() * s.array();
   });
   return res;
 }

stan/math/rev/fun/log_sum_exp.hpp

@@ -9,6 +9,7 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <cmath>
 #include <vector>
 
@@ -68,11 +69,11 @@ template <typename T, require_eigen_st<is_var, T>* = nullptr,
           require_not_var_matrix_t<T>* = nullptr>
 inline var log_sum_exp(T&& v) {
   arena_t<T> arena_v = std::forward<T>(v);
-  auto arena_v_val = to_arena(arena_v.val());
-  var res = log_sum_exp(arena_v_val);
-  reverse_pass_callback([arena_v, arena_v_val, res]() mutable {
-    arena_v.adj()
-        += res.adj() * (arena_v_val.array().val() - res.val()).exp().matrix();
+  const auto softmax_lse = internal::softmax_log_sum_exp(arena_v.val());
+  auto arena_softmax = to_arena(softmax_lse.first);
+  var res = softmax_lse.second;
+  reverse_pass_callback([arena_v, arena_softmax, res]() mutable {
+    arena_v.adj() += res.adj() * arena_softmax;
   });
 
   return res;
@@ -86,9 +87,12 @@ inline var log_sum_exp(T&& v) {
  */
 template <typename T, require_var_matrix_t<T>* = nullptr>
 inline var log_sum_exp(const T& x) {
-  return make_callback_vari(log_sum_exp(x.val()), [x](const auto& res) mutable {
-    x.adj() += res.adj() * (x.val().array().val() - res.val()).exp().matrix();
-  });
+  const auto softmax_lse = internal::softmax_log_sum_exp(x.val());
+  auto arena_softmax = to_arena(softmax_lse.first);
+  return make_callback_vari(softmax_lse.second,
+                            [x, arena_softmax](const auto& res) mutable {
+                              x.adj() += res.adj() * arena_softmax;
+                            });
 }
 
 /**