Numerically stable derivatives of ordered_logistic_lpmf and ordered_logistic_glm_lpmf

stan/math/opencl/kernels/ordered_logistic_glm_lpmf.hpp

@@ -83,19 +83,22 @@ static constexpr const char* ordered_logistic_glm_kernel_code = STRINGIFY(
             logp -= log1p_exp(-cut2);
           }
           if (y != 1 && y != N_classes) {
-            logp += log1m_exp(cut1 - cut2);
+            logp += log1m_exp(cut_y2 - cut_y1);
           }
 
           if (need_location_derivative || need_cuts_derivative) {
-            double exp_cuts_diff = exp(cut_y2 - cut_y1);
             d1 = inv_logit(-cut2);
-            d1 -= exp_cuts_diff / (exp_cuts_diff - 1);
-            d2 = 1 / (1 - exp_cuts_diff);
-            d2 -= inv_logit(-cut1);
-
+            d2 = inv_logit(cut1);
             if (need_location_derivative) {
               location_derivative[gid] = d1 - d2;
             }
+            if (need_cuts_derivative) {
+              // -1 / expm1(c_y - c_{y-1}), zero for the first and last class
+              const double cuts_diff = cut_y2 - cut_y1;
+              const double q = exp(cuts_diff) / expm1(cuts_diff);
+              d1 -= q;
+              d2 -= q;
+            }
           }
         }
       }

stan/math/opencl/kernels/ordered_logistic_lpmf.hpp

@@ -79,19 +79,22 @@ static constexpr const char* ordered_logistic_kernel_code = STRINGIFY(
             logp -= log1p_exp(-cut2);
           }
           if (y != 1 && y != N_classes) {
-            logp += log1m_exp(cut1 - cut2);
+            logp += log1m_exp(cut_y2 - cut_y1);
           }
 
           if (need_lambda_derivative || need_cuts_derivative) {
-            double exp_cuts_diff = exp(cut_y2 - cut_y1);
             d1 = inv_logit(-cut2);
-            d1 -= exp_cuts_diff / (exp_cuts_diff - 1);
-            d2 = 1 / (1 - exp_cuts_diff);
-            d2 -= inv_logit(-cut1);
-
+            d2 = inv_logit(cut1);
             if (need_lambda_derivative) {
               lambda_derivative[gid] = d1 - d2;
             }
+            if (need_cuts_derivative) {
+              // -1 / expm1(c_y - c_{y-1}), zero for the first and last class
+              const double cuts_diff = cut_y2 - cut_y1;
+              const double q = exp(cuts_diff) / expm1(cuts_diff);
+              d1 -= q;
+              d2 -= q;
+            }
           }
         }
       }

stan/math/prim/prob/ordered_logistic_glm_lpmf.hpp

@@ -133,11 +133,18 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
   Array<T_partials_return, Dynamic, 1> cut1
       = as_array_or_scalar(location) - cuts_y1;
 
+  // exp(-|cut|) on the same sign partition as every select below
+  Array<T_partials_return, Dynamic, 1> exp_m_abs_cut1
+      = (cut1 > 0.0).select(-cut1, cut1).exp();
+  Array<T_partials_return, Dynamic, 1> exp_m_abs_cut2
+      = (cut2 > 0.0).select(-cut2, cut2).exp();
+  Array<T_cuts_partials, Dynamic, 1> cuts_diff = cuts_y2 - cuts_y1;
+
   // Not immediately evaluating next two expressions benefits performance
   auto m_log_1p_exp_cut1
-      = (cut1 > 0.0).select(-cut1, 0) - (-cut1.abs()).exp().log1p();
+      = (cut1 > 0.0).select(-cut1, 0) - exp_m_abs_cut1.log1p();
   auto m_log_1p_exp_m_cut2
-      = (cut2 <= 0.0).select(cut2, 0) - (-cut2.abs()).exp().log1p();
+      = (cut2 <= 0.0).select(cut2, 0) - exp_m_abs_cut2.log1p();
 
   T_partials_return logp(0);
   if constexpr (is_vector<T_y>::value) {
@@ -147,7 +154,7 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
                .select(m_log_1p_exp_cut1,
                        y_vec.cwiseEqual(N_classes).select(
                            m_log_1p_exp_m_cut2,
-                           m_log_1p_exp_m_cut2 + log1m_exp(cut1 - cut2).array()
+                           m_log_1p_exp_m_cut2 + log1m_exp(cuts_diff).array()
                                + m_log_1p_exp_cut1))
                .sum();
   } else {
@@ -156,27 +163,23 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
     } else if (y_seq[0] == N_classes) {
       logp = m_log_1p_exp_m_cut2.sum();
     } else {
-      logp = (m_log_1p_exp_m_cut2 + log1m_exp(cut1 - cut2).array()
+      logp = (m_log_1p_exp_m_cut2 + log1m_exp(cuts_diff).array()
               + m_log_1p_exp_cut1)
                  .sum();
     }
   }
 
   auto ops_partials = make_partials_propagator(x_ref, beta_ref, cuts_ref);
   if constexpr (is_any_autodiff_v<T_x, T_beta, T_cuts>) {
-    Array<T_partials_return, Dynamic, 1> exp_m_abs_cut1 = (-cut1.abs()).exp();
-    Array<T_partials_return, Dynamic, 1> exp_m_abs_cut2 = (-cut2.abs()).exp();
-    Array<T_partials_return, Dynamic, 1> exp_cuts_diff = exp(cuts_y2 - cuts_y1);
-    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut2 = (cut2 > 0).select(
-        exp_m_abs_cut2 / (1 + exp_m_abs_cut2), 1 / (1 + exp_m_abs_cut2));
-    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut1 = (cut1 > 0).select(
-        exp_m_abs_cut1 / (1 + exp_m_abs_cut1), 1 / (1 + exp_m_abs_cut1));
-    Array<T_partials_return, Dynamic, 1> d1
-        = inv_logit_neg_cut2 - exp_cuts_diff / (exp_cuts_diff - 1);
-    Array<T_partials_return, Dynamic, 1> d2
-        = 1 / (1 - exp_cuts_diff) - inv_logit_neg_cut1;
+    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut2
+        = (cut2 > 0).select(exp_m_abs_cut2, T_partials_return(1))
+          / (1 + exp_m_abs_cut2);
+    Array<T_partials_return, Dynamic, 1> inv_logit_cut1
+        = (cut1 > 0).select(T_partials_return(1), exp_m_abs_cut1)
+          / (1 + exp_m_abs_cut1);
     if constexpr (is_any_autodiff_v<T_x, T_beta>) {
-      Matrix<T_partials_return, 1, Dynamic> location_derivative = d1 - d2;
+      Matrix<T_partials_return, 1, Dynamic> location_derivative
+          = inv_logit_neg_cut2 - inv_logit_cut1;
       if constexpr (is_autodiff_v<T_x>) {
         if constexpr (T_x_rows == 1) {
           edge<0>(ops_partials).partials_
@@ -189,8 +192,7 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
       if constexpr (is_autodiff_v<T_beta>) {
         if constexpr (T_x_rows == 1) {
           edge<1>(ops_partials).partials_
-              = (location_derivative * x_val.replicate(N_instances, 1))
-                    .transpose();
+              = (location_derivative.sum() * x_val).transpose();
         } else if constexpr (T_x::ColsAtCompileTime == 1) {
           edge<1>(ops_partials).partials_[0]
               = dot_product(location_derivative, x_val);
@@ -201,6 +203,11 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
       }
     }
     if constexpr (is_autodiff_v<T_cuts>) {
+      // -1 / expm1(c_y - c_{y-1}), zero for the first and last class
+      Array<T_partials_return, Dynamic, 1> q
+          = cuts_diff.exp() / cuts_diff.expm1();
+      Array<T_partials_return, Dynamic, 1> d1 = inv_logit_neg_cut2 - q;
+      Array<T_partials_return, Dynamic, 1> d2 = inv_logit_cut1 - q;
       for (int i = 0; i < N_instances; i++) {
         int c = y_seq[i];
         if (c != N_classes) {

stan/math/prim/prob/ordered_logistic_lpmf.hpp

@@ -6,13 +6,12 @@
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/as_value_array_or_scalar.hpp>
 #include <stan/math/prim/fun/exp.hpp>
-#include <stan/math/prim/fun/inv_logit.hpp>
 #include <stan/math/prim/fun/is_integer.hpp>
-#include <stan/math/prim/fun/log1p_exp.hpp>
-#include <stan/math/prim/fun/log_inv_logit_diff.hpp>
+#include <stan/math/prim/fun/log1m_exp.hpp>
 #include <stan/math/prim/fun/scalar_seq_view.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_mvt.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <stan/math/prim/fun/vector_seq_view.hpp>
 #include <stan/math/prim/functor/partials_propagator.hpp>
@@ -30,25 +29,13 @@ namespace math {
  * will be the dot product of a vector of regression coefficients
  * and a vector of predictors for the outcome
  *
-  \f[
-    \frac{\partial }{\partial \lambda} =
-    \begin{cases}\\
-    -\mathrm{logit}^{-1}(\lambda - c_1) & \mbox{if } k = 1,\\
-    -(((1-e^{c_{k-1}-c_{k-2}})^{-1} - \mathrm{logit}^{-1}(c_{k-2}-\lambda)) +
-    ((1-e^{c_{k-2}-c_{k-1}})^{-1} - \mathrm{logit}^{-1}(c_{k-1}-\lambda)))
-    & \mathrm{if } 1 < k < K, \mathrm{and}\\
-    \mathrm{logit}^{-1}(c_{K-2}-\lambda) & \mathrm{if } k = K.
-    \end{cases}
-  \f]
-
   \f[
     \frac{\partial }{\partial \lambda} =
     \begin{cases}
     -\mathrm{logit}^{-1}(\lambda - c_1) & \text{if } k = 1,\\
-    -(((1-e^{c_{k-1}-c_{k-2}})^{-1} - \mathrm{logit}^{-1}(c_{k-2}-\lambda)) +
-    ((1-e^{c_{k-2}-c_{k-1}})^{-1} - \mathrm{logit}^{-1}(c_{k-1}-\lambda)))
+    \mathrm{logit}^{-1}(c_{k-1} - \lambda) - \mathrm{logit}^{-1}(\lambda - c_k)
     & \text{if } 1 < k < K, \text{ and}\\
-    \mathrm{logit}^{-1}(c_{K-2}-\lambda) & \text{if } k = K.
+    \mathrm{logit}^{-1}(c_{K-1} - \lambda) & \text{if } k = K.
     \end{cases}
   \f]
  *
@@ -144,16 +131,23 @@ inline return_type_t<T_loc, T_cut> ordered_logistic_lpmf(const T_y& y,
   Array<T_partials_return, Dynamic, 1> cut2 = lambda_val - cuts_y2;
   Array<T_partials_return, Dynamic, 1> cut1 = lambda_val - cuts_y1;
 
+  // exp(-|cut|) on the same sign partition as every select below
+  Array<T_partials_return, Dynamic, 1> exp_m_abs_cut1
+      = (cut1 > 0.0).select(-cut1, cut1).exp();
+  Array<T_partials_return, Dynamic, 1> exp_m_abs_cut2
+      = (cut2 > 0.0).select(-cut2, cut2).exp();
+  Array<T_cuts_val, Dynamic, 1> cuts_diff = cuts_y2 - cuts_y1;
+
   // Not immediately evaluating next two expressions benefits performance
   auto m_log_1p_exp_cut1
-      = (cut1 > 0.0).select(-cut1, 0) - (-cut1.abs()).exp().log1p();
+      = (cut1 > 0.0).select(-cut1, 0) - exp_m_abs_cut1.log1p();
   auto m_log_1p_exp_m_cut2
-      = (cut2 <= 0.0).select(cut2, 0) - (-cut2.abs()).exp().log1p();
+      = (cut2 <= 0.0).select(cut2, 0) - exp_m_abs_cut2.log1p();
 
   if constexpr (is_vector<T_y>::value) {
     Eigen::Map<const Eigen::Matrix<value_type_t<T_y>, Eigen::Dynamic, 1>> y_vec(
         y_seq.data(), y_seq.size());
-    auto log1m_exp_cuts_diff = log1m_exp(cut1 - cut2);
+    auto log1m_exp_cuts_diff = log1m_exp(cuts_diff);
     logp = y_vec.cwiseEqual(1)
                .select(m_log_1p_exp_cut1,
                        y_vec.cwiseEqual(K).select(m_log_1p_exp_m_cut2,
@@ -167,29 +161,29 @@ inline return_type_t<T_loc, T_cut> ordered_logistic_lpmf(const T_y& y,
     } else if (y_seq[0] == K) {
       logp = m_log_1p_exp_m_cut2.sum();
     } else {
-      logp = (m_log_1p_exp_m_cut2 + log1m_exp(cut1 - cut2).array()
+      logp = (m_log_1p_exp_m_cut2 + log1m_exp(cuts_diff).array()
               + m_log_1p_exp_cut1)
                  .sum();
     }
   }
 
   auto ops_partials = make_partials_propagator(lambda_ref, c_ref);
   if constexpr (is_any_autodiff_v<T_loc, T_cut>) {
-    Array<T_partials_return, Dynamic, 1> exp_m_abs_cut1 = (-cut1.abs()).exp();
-    Array<T_partials_return, Dynamic, 1> exp_m_abs_cut2 = (-cut2.abs()).exp();
-    Array<T_partials_return, Dynamic, 1> exp_cuts_diff = exp(cuts_y2 - cuts_y1);
-    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut2 = (cut2 > 0).select(
-        exp_m_abs_cut2 / (1 + exp_m_abs_cut2), 1 / (1 + exp_m_abs_cut2));
-    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut1 = (cut1 > 0).select(
-        exp_m_abs_cut1 / (1 + exp_m_abs_cut1), 1 / (1 + exp_m_abs_cut1));
-    Array<T_partials_return, Dynamic, 1> d1
-        = inv_logit_neg_cut2 - exp_cuts_diff / (exp_cuts_diff - 1);
-    Array<T_partials_return, Dynamic, 1> d2
-        = 1 / (1 - exp_cuts_diff) - inv_logit_neg_cut1;
+    Array<T_partials_return, Dynamic, 1> inv_logit_neg_cut2
+        = (cut2 > 0).select(exp_m_abs_cut2, T_partials_return(1))
+          / (1 + exp_m_abs_cut2);
+    Array<T_partials_return, Dynamic, 1> inv_logit_cut1
+        = (cut1 > 0).select(T_partials_return(1), exp_m_abs_cut1)
+          / (1 + exp_m_abs_cut1);
     if constexpr (is_autodiff_v<T_loc>) {
-      partials<0>(ops_partials) = d1 - d2;
+      partials<0>(ops_partials) = inv_logit_neg_cut2 - inv_logit_cut1;
     }
     if constexpr (is_autodiff_v<T_cut>) {
+      // -1 / expm1(c_y - c_{y-1}), zero for the first and last class
+      Array<T_partials_return, Dynamic, 1> q
+          = cuts_diff.exp() / cuts_diff.expm1();
+      Array<T_partials_return, Dynamic, 1> d1 = inv_logit_neg_cut2 - q;
+      Array<T_partials_return, Dynamic, 1> d2 = inv_logit_cut1 - q;
       for (int i = 0; i < N; i++) {
         int c = y_seq[i];
         if (c != K) {

test/unit/math/mix/prob/ordered_logistic_glm_lpmf_test.cpp

@@ -1,5 +1,6 @@
 #include <stan/math/mix.hpp>
 #include <test/unit/math/test_ad.hpp>
+#include <test/unit/math/mix/prob/ordered_logistic_test_helpers.hpp>
 
 TEST_F(AgradRev, mathMixScalFun_ordered_logistic_glm_lpmf) {
   auto f = [](const auto y) {
@@ -100,3 +101,12 @@ TEST_F(AgradRev, mathMixScalFun_ordered_logistic_glm_lpmf_x_column_vector) {
 
   stan::test::expect_ad(f, x, beta, cutpoints);
 }
+
+// Against closed forms at and near the cuts, in the tails and for close cuts.
+TEST_F(AgradRev, mathMixScalFun_ordered_logistic_glm_lpmf_derivatives) {
+  auto lpmf
+      = [](const auto& y, const auto& x, const auto& beta, const auto& cuts) {
+          return stan::math::ordered_logistic_glm_lpmf(y, x, beta, cuts);
+        };
+  ordered_logistic_test::expect_all(lpmf, true);
+}

test/unit/math/mix/prob/ordered_logistic_test.cpp

@@ -1,6 +1,7 @@
 #include <stan/math/mix.hpp>
 #include <test/unit/math/test_ad.hpp>
 #include <test/unit/math/rev/fun/util.hpp>
+#include <test/unit/math/mix/prob/ordered_logistic_test_helpers.hpp>
 #include <gtest/gtest.h>
 #include <vector>
 
@@ -722,3 +723,24 @@ TEST_F(AgradRev, ProbDistributionsOrdLog_top_category_value_and_gradient) {
   EXPECT_FLOAT_EQ(c[0].adj(), 0.0);
   EXPECT_FLOAT_EQ(c[1].adj(), -0.574442516811659);
 }
+
+// Against closed forms at and near the cuts, in the tails and for close cuts.
+TEST_F(AgradRev, ProbDistributionsOrdLog_derivatives) {
+  auto lpmf
+      = [](const auto& y, const auto& x, const auto& beta, const auto& cuts) {
+          return stan::math::ordered_logistic_lpmf(
+              y, stan::math::multiply(x, beta), cuts);
+        };
+  ordered_logistic_test::expect_all(lpmf, false);
+}
+
+// The same with one cut vector per location.
+TEST_F(AgradRev, ProbDistributionsOrdLog_derivatives_cuts_array) {
+  auto lpmf
+      = [](const auto& y, const auto& x, const auto& beta, const auto& cuts) {
+          std::vector<std::decay_t<decltype(cuts)>> cuts_arr(x.rows(), cuts);
+          return stan::math::ordered_logistic_lpmf(
+              y, stan::math::multiply(x, beta), cuts_arr);
+        };
+  ordered_logistic_test::expect_all(lpmf, false);
+}