Improved numerical accuracy for log_sum_exp, softmax, and log_softmax

stan/math/fwd/fun/log_softmax.hpp

@@ -14,10 +14,11 @@ 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,31 +29,35 @@ 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));
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline auto log_softmax(Mat&& x) {
+  using mat = std::decay_t<Mat>;
+  constexpr int Rows = mat::RowsAtCompileTime;
+  constexpr int Cols = mat::ColsAtCompileTime;
+  using T = typename value_type_t<mat>::Scalar;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
   if (x_ref.size() == 0) {
     return Eigen::Matrix<fvar<T>, Rows, Cols>{};
   }
-  const auto s = softmax(value_of(x_ref));
+  const auto x_val = value_of(x_ref);
+  const auto lse = log_sum_exp(x_val);
+  const auto s = softmax(x_val);
   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();
+  const auto dot_sd = (s.array() * d_in.array()).sum();
+
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
+  result.val() = (x_val.array() - lse).matrix();
   result.d() = (d_in.array() - dot_sd).matrix();
   return result;
 }
 
 }  // namespace math
 }  // namespace stan
-#endif
+#endif

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.hpp>
 #include <cmath>
 #include <vector>
 
@@ -56,11 +57,10 @@ 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 probs = softmax(vals);
+        return fvar<T_fvar_inner>(log_sum_exp(vals),
+                                  v.d().cwiseProduct(probs).sum());
       });
 }
 

stan/math/fwd/fun/softmax.hpp

@@ -29,24 +29,25 @@ 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));
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline auto softmax(Mat&& x) {
+  using mat = std::decay_t<Mat>;
+  constexpr int Rows = mat::RowsAtCompileTime;
+  constexpr int Cols = mat::ColsAtCompileTime;
+  using T = typename value_type_t<mat>::Scalar;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
   if (x_ref.size() == 0) {
     return Eigen::Matrix<fvar<T>, Rows, Cols>{};
   }
   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());
+  const auto dot_sd = (s.array() * d_in.array()).sum();
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
   result.val() = s;
   result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
   return result;

stan/math/prim/fun/log_softmax.hpp

@@ -12,9 +12,10 @@ 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)
  * \ = \ y - \log \sum_{k=1}^K \exp(y_k)
@@ -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
- * @return log softmax of the input, preserving the container structure; an
- *   empty result if any input vector is empty
+ * @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 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) {

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.hpp>
 #include <cmath>
 #include <vector>
 
@@ -71,8 +72,7 @@ inline var log_sum_exp(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();
+    arena_v.adj() += res.adj() * softmax(arena_v_val);
   });
 
   return res;
@@ -87,7 +87,7 @@ 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();
+    x.adj() += res.adj() * softmax(x.val());
   });
 }
 

stan/math/rev/fun/softmax.hpp

@@ -15,9 +15,10 @@ namespace stan {
 namespace math {
 
 /**
- * Return the softmax of the specified vector or row vector.
+ * Return the 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 softmax of the input, or an empty result if the input is empty
  */
@@ -31,17 +32,17 @@ inline auto softmax(T&& x) {
       = return_var_matrix_t<plain_type_t<decltype(x_arena.val())>, T>;
   arena_t<return_t> res = softmax(x_arena.val());
   reverse_pass_callback([x_arena, res]() mutable {
-    x_arena.adj().array()
-        += res.val().array() * (res.adj().array() - res.val().dot(res.adj()));
+    const auto dot_sd = (res.val().array() * res.adj().array()).sum();
+    x_arena.adj().array() += res.val().array() * (res.adj().array() - dot_sd);
   });
   return res;
 }
 
 /**
- * Return the softmax of each vector in an array.
+ * Return the softmax of each vector or matrix in an array.
  *
  * @tparam T `std::vector` whose scalar type is `var`
- * @param x array of vectors to transform
+ * @param x array of vectors or matrices to transform
  * @return array of softmax results
  */
 template <typename T, require_std_vector_st<is_var, T>* = nullptr>