Numerically stable derivatives of bernoulli_logit_lpmf and bernoulli_logit_glm_lpmf

stan/math/opencl/prim/bernoulli_logit_glm_lpmf.hpp

@@ -92,21 +92,18 @@ inline return_type_t<T_x_cl, T_alpha_cl, T_beta_cl> bernoulli_logit_glm_lpmf(
   auto ytheta_expr = matrix_vector_multiply(x_val, beta_val) + alpha_val;
   auto signs_expr = 2 * y_val - 1;
   auto ytheta_signs_expr = elt_multiply(ytheta_expr, signs_expr);
-  auto exp_m_ytheta_expr = exp(-ytheta_signs_expr);
-  const double cutoff = 20.0;
-  auto high_bound_expr = ytheta_signs_expr > cutoff;
-  auto low_bound_expr = ytheta_signs_expr < -cutoff;
+  auto positive_expr = ytheta_signs_expr > 0;
+  auto exp_m_abs_ytheta_expr
+      = exp(select(positive_expr, -ytheta_signs_expr, ytheta_signs_expr));
   auto err_cond_expr = y_val < 0 || y_val > 1;
   auto logp_expr
       = colwise_sum(select(err_cond_expr, NOT_A_NUMBER,
-                           select(high_bound_expr, -exp_m_ytheta_expr,
-                                  select(low_bound_expr, ytheta_signs_expr,
-                                         -log1p(exp_m_ytheta_expr)))));
-  auto theta_derivative_expr
-      = select(high_bound_expr, elt_multiply(signs_expr, exp_m_ytheta_expr),
-               select(low_bound_expr, signs_expr,
-                      elt_divide(elt_multiply(signs_expr, exp_m_ytheta_expr),
-                                 (exp_m_ytheta_expr + 1))));
+                           select(positive_expr, 0.0, ytheta_signs_expr)
+                               - log1p(exp_m_abs_ytheta_expr)));
+  auto theta_derivative_expr = elt_divide(
+      elt_multiply(signs_expr,
+                   select(positive_expr, exp_m_abs_ytheta_expr, 1.0)),
+      exp_m_abs_ytheta_expr + 1);
 
   const int wgs = logp_expr.rows();
   matrix_cl<double> logp_cl(wgs, 1);

stan/math/opencl/prim/bernoulli_logit_lpmf.hpp

@@ -55,18 +55,14 @@ inline return_type_t<T_prob_cl> bernoulli_logit_lpmf(const T_n_cl& n,
 
   auto signs_expr = 2 * n - 1.0;  // subtracting 1.0 converts int to double
   auto ntheta_expr = elt_multiply(signs_expr, theta_val);
-  auto exp_m_ntheta_expr = exp(-ntheta_expr);
-  static constexpr double cutoff = 20.0;
-  auto condition1_expr = ntheta_expr > cutoff;
-  auto condition2_expr = ntheta_expr < -cutoff;
-  auto logp_expr = colwise_sum(
-      select(condition1_expr, -exp_m_ntheta_expr,
-             select(condition2_expr, ntheta_expr, -log1p(exp_m_ntheta_expr))));
-  auto deriv_expr = select(
-      condition1_expr, elt_multiply(signs_expr, exp_m_ntheta_expr),
-      select(condition2_expr, signs_expr,
-             elt_multiply(signs_expr, elt_divide(exp_m_ntheta_expr,
-                                                 (exp_m_ntheta_expr + 1)))));
+  auto positive_expr = ntheta_expr > 0;
+  auto exp_m_abs_ntheta_expr
+      = exp(select(positive_expr, -ntheta_expr, ntheta_expr));
+  auto logp_expr = colwise_sum(select(positive_expr, 0.0, ntheta_expr)
+                               - log1p(exp_m_abs_ntheta_expr));
+  auto deriv_expr = elt_multiply(
+      signs_expr, elt_divide(select(positive_expr, exp_m_abs_ntheta_expr, 1.0),
+                             exp_m_abs_ntheta_expr + 1));
 
   matrix_cl<double> logp_cl;
   matrix_cl<double> deriv_cl;

stan/math/prim/prob/bernoulli_logit_glm_lpmf.hpp

@@ -11,8 +11,10 @@
 #include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
+#include <stan/math/prim/fun/log1p.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <stan/math/prim/functor/partials_propagator.hpp>
@@ -112,15 +114,13 @@ inline return_type_t<T_x, T_alpha, T_beta> bernoulli_logit_glm_lpmf(
     ytheta = signs * (ytheta + as_array_or_scalar(alpha_val_vec));
   }
 
-  // Compute the log-density and handle extreme values gracefully
-  // using Taylor approximations.
-  // And compute the derivatives wrt theta.
-  static constexpr double cutoff = 20.0;
-  Eigen::Array<T_partials_return, Dynamic, 1> exp_m_ytheta = exp(-ytheta);
-  T_partials_return logp = sum(
-      (ytheta > cutoff)
-          .select(-exp_m_ytheta,
-                  (ytheta < -cutoff).select(ytheta, -log1p(exp_m_ytheta))));
+  // exp(-|ytheta|) in (0, 1], so no overflow enters the autodiff tape
+  Array<T_partials_return, Dynamic, 1> exp_m_abs_ytheta
+      = exp((ytheta > 0).select(-ytheta, ytheta));
+  // two sums: no per-element node for the difference in nested autodiff
+  T_partials_return logp
+      = sum((ytheta > 0).select(T_partials_return(0), ytheta))
+        - sum(log1p(exp_m_abs_ytheta));
 
   if (!isfinite(logp)) {
     check_finite(function, "Weight vector", beta);
@@ -132,11 +132,8 @@ inline return_type_t<T_x, T_alpha, T_beta> bernoulli_logit_glm_lpmf(
   // Compute the necessary derivatives.
   if constexpr (is_any_autodiff_v<T_beta, T_x, T_alpha>) {
     Matrix<T_partials_return, Dynamic, 1> theta_derivative
-        = (ytheta > cutoff)
-              .select(signs * exp_m_ytheta,
-                      (ytheta < -cutoff)
-                          .select(signs * T_partials_return(1.0),
-                                  signs * exp_m_ytheta / (exp_m_ytheta + 1)));
+        = signs * (ytheta > 0).select(exp_m_abs_ytheta, T_partials_return(1))
+          / (1 + exp_m_abs_ytheta);
     if constexpr (is_autodiff_v<T_beta>) {
       if constexpr (T_x_rows == 1) {
         edge<2>(ops_partials).partials_ = theta_derivative.sum() * x_val;

stan/math/prim/prob/bernoulli_logit_lpmf.hpp

@@ -8,8 +8,8 @@
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/log1p.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
-#include <stan/math/prim/fun/promote_scalar.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <stan/math/prim/fun/value_of_rec.hpp>
@@ -70,23 +70,18 @@ inline return_type_t<T_prob> bernoulli_logit_lpmf(const T_n& n,
     T_partials_return ntheta_s = signs * theta_val;
     ntheta = T_partials_array::Constant(1, 1, ntheta_s);
   }
-  T_partials_array exp_m_ntheta = exp(-ntheta);
-  static constexpr double cutoff = 20.0;
-  T_partials_return logp = sum(
-      (ntheta > cutoff)
-          .select(-exp_m_ntheta,
-                  (ntheta < -cutoff).select(ntheta, -log1p(exp_m_ntheta))));
+  // exp(-|ntheta|) in (0, 1], so no overflow enters the autodiff tape
+  T_partials_array exp_m_abs_ntheta = exp((ntheta > 0).select(-ntheta, ntheta));
+  // two sums: no per-element node for the difference in nested autodiff
+  T_partials_return logp
+      = sum((ntheta > 0).select(T_partials_return(0), ntheta))
+        - sum(log1p(exp_m_abs_ntheta));
 
   auto ops_partials = make_partials_propagator(theta_ref);
   if constexpr (is_autodiff_v<T_prob>) {
     edge<0>(ops_partials).partials_
-        = (ntheta > cutoff)
-              .select(
-                  promote_scalar<T_partials_return>(signs * exp_m_ntheta),
-                  (ntheta >= -cutoff)
-                      .select(promote_scalar<T_partials_return>(
-                                  signs * exp_m_ntheta / (exp_m_ntheta + 1)),
-                              promote_scalar<T_partials_return>(signs)));
+        = signs * (ntheta > 0).select(exp_m_abs_ntheta, T_partials_return(1))
+          / (1 + exp_m_abs_ntheta);
   }
   return ops_partials.build(logp);
 }

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

@@ -1,5 +1,8 @@
 #include <stan/math/mix.hpp>
 #include <test/unit/math/test_ad.hpp>
+#include <cmath>
+#include <limits>
+#include <vector>
 
 TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_glm_lpmf) {
   auto f = [](const auto y) {
@@ -37,3 +40,110 @@ TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_glm_lpmf_x_column_vector) {
 
   stan::test::expect_ad(f, x, alpha, beta);
 }
+
+namespace bernoulli_logit_glm_test {
+// theta = alpha + x beta = J v with v = (alpha, beta), J = [I, x]. With
+// s = 2y - 1, t = s theta, p = inv_logit(t), q = 1 - p, the n-th summand is
+// log(p_n) with theta_n derivatives s q, -p q and -s p q (q - p); the
+// derivatives in v follow by the chain rule through J. Each tolerance is
+// relative to the same sum taken over absolute values.
+template <typename T_y, typename T_x>
+void expect_derivatives(const T_y& y, const T_x& x,
+                        const Eigen::VectorXd& alpha,
+                        const Eigen::VectorXd& beta) {
+  using stan::math::fvar;
+  using stan::math::var;
+  const int N = alpha.size();
+  const int K = beta.size();
+  const int M = N + K;
+  Eigen::MatrixXd J(N, M);
+  J << Eigen::MatrixXd::Identity(N, N), x.replicate(N / x.rows(), 1);
+  Eigen::VectorXd v(M);
+  v << alpha, beta;
+  Eigen::VectorXd theta = J * v;
+  Eigen::ArrayXd lp(N), d1(N), d2(N), d3(N), w(N);
+  for (int n = 0; n < N; ++n) {
+    double s = 2 * stan::get(y, n) - 1, t = s * theta(n);
+    double p = stan::math::inv_logit(t), q = stan::math::inv_logit(-t);
+    lp(n) = stan::math::log_inv_logit(t);
+    d1(n) = s * q;
+    d2(n) = -p * q;
+    d3(n) = -s * p * q * (q - p);
+    w(n) = p * q;
+  }
+  const double tiny = std::numeric_limits<double>::min();
+  Eigen::MatrixXd A = J.cwiseAbs();
+  Eigen::VectorXd grad_exact = J.transpose() * d1.matrix();
+  Eigen::VectorXd grad_tol
+      = (1e-13 * A.transpose() * d1.abs().matrix()).array() + tiny;
+  Eigen::MatrixXd H_exact = J.transpose() * d2.matrix().asDiagonal() * J;
+  Eigen::MatrixXd H_tol
+      = (1e-13 * A.transpose() * w.matrix().asDiagonal() * A).array() + tiny;
+
+  auto f = [&](const auto& u) {
+    return stan::math::bernoulli_logit_glm_lpmf(y, x, u.head(N), u.tail(K));
+  };
+  double fx;
+  Eigen::VectorXd grad, grad_fwd;
+  Eigen::MatrixXd H, H_fwd, H_ffv;
+  std::vector<Eigen::MatrixXd> grad_H;
+  stan::math::hessian(f, v, fx, grad, H);
+  stan::math::hessian<double>(f, v, fx, grad_fwd, H_fwd);  // fvar<fvar<double>>
+  stan::math::grad_hessian(f, v, fx, H_ffv, grad_H);
+  EXPECT_NEAR(fx, lp.sum(), 1e-13 * lp.abs().sum() + tiny);
+  for (int i = 0; i < M; ++i) {
+    EXPECT_NEAR(grad(i), grad_exact(i), grad_tol(i)) << i;
+    EXPECT_NEAR(grad_fwd(i), grad_exact(i), grad_tol(i)) << i;
+    for (int j = 0; j < M; ++j) {
+      EXPECT_NEAR(H(i, j), H_exact(i, j), H_tol(i, j)) << i << ", " << j;
+      EXPECT_NEAR(H_fwd(i, j), H_exact(i, j), H_tol(i, j)) << i << ", " << j;
+      EXPECT_NEAR(H_ffv(i, j), H_exact(i, j), H_tol(i, j)) << i << ", " << j;
+      Eigen::ArrayXd Jij = J.col(i).array() * J.col(j).array();
+      for (int k = 0; k < M; ++k) {
+        Eigen::ArrayXd Jijk = Jij * J.col(k).array();
+        EXPECT_NEAR(grad_H[i](j, k), (d3 * Jijk).sum(),
+                    1e-13 * (w * Jijk.abs()).sum() + tiny)
+            << i << ", " << j << ", " << k;
+      }
+    }
+  }
+
+  // derivative of the value itself
+  stan::math::nested_rev_autodiff nested;
+  Eigen::Matrix<fvar<var>, Eigen::Dynamic, 1> v_fv(M);
+  for (int k = 0; k < M; ++k) {
+    v_fv(k) = fvar<var>(v(k), 0);
+  }
+  f(v_fv).val_.grad();
+  for (int i = 0; i < M; ++i) {
+    EXPECT_NEAR(v_fv(i).val_.adj(), grad_exact(i), grad_tol(i)) << i;
+    Eigen::Matrix<fvar<fvar<double>>, Eigen::Dynamic, 1> v_ff(M);
+    for (int k = 0; k < M; ++k) {
+      v_ff(k) = fvar<fvar<double>>(fvar<double>(v(k), i == k), 0);
+    }
+    EXPECT_NEAR(f(v_ff).val_.d_, grad_exact(i), grad_tol(i)) << i;
+  }
+}
+}  // namespace bernoulli_logit_glm_test
+
+// t = (2y - 1) theta in the tails (exp(-t) overflows at t < -709, curvature
+// was 0 below -20) and at t = +-0, where the value's own derivative
+// must not switch branch. Matrix and row-vector x, vector and scalar y.
+TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_glm_lpmf_tails) {
+  using bernoulli_logit_glm_test::expect_derivatives;
+  std::vector<int> y{1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0};
+  Eigen::MatrixXd x(13, 2);
+  x << 0.5, -1, 1, 0.25, -0.5, 2, 1.5, -0.75, 0.25, 0.125, -1, -0.25, 0.75, 0.5,
+      -0.625, 1.25, 0.375, -1, 1.125, 0.625, -0.5, -1.25, 0.75, 0.5, -0.25, 1;
+  Eigen::VectorXd beta(2);  // dyadic x and beta: theta = 0 exactly
+  beta << 0.5, -0.25;
+  Eigen::VectorXd theta(13);
+  theta << -2500, 720, -30, 20.5, -0.25, 0, 0, -0.0, -0.0, 21, 21, 40, -19.5;
+  Eigen::VectorXd alpha = theta - x * beta;
+  expect_derivatives(y, x, alpha, beta);
+  expect_derivatives(1, x, alpha, beta);
+
+  Eigen::RowVectorXd x_row = x.row(0);
+  alpha = theta.array() - x_row.dot(beta);
+  expect_derivatives(y, x_row, alpha, beta);
+}

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

@@ -0,0 +1,91 @@
+#include <stan/math/mix.hpp>
+#include <test/unit/math/test_ad.hpp>
+#include <cmath>
+#include <limits>
+#include <vector>
+
+namespace bernoulli_logit_lpmf_test {
+// With s = 2y - 1, t = s theta, p = inv_logit(t), q = 1 - p, the log density
+// is log(p) and its theta derivatives are s q, -p q and -s p q (q - p).
+struct exact {
+  double lp, d1, d2, d3, scale;
+  exact(int y, double theta) {
+    double s = 2 * y - 1, t = s * theta;
+    double p = stan::math::inv_logit(t), q = stan::math::inv_logit(-t);
+    lp = stan::math::log_inv_logit(t);
+    d1 = s * q;
+    d2 = -p * q;
+    d3 = -s * p * q * (q - p);
+    scale = p * q;
+  }
+};
+
+inline double tol(double x) {
+  return 1e-13 * std::fabs(x) + std::numeric_limits<double>::min();
+}
+}  // namespace bernoulli_logit_lpmf_test
+
+// t = (2y - 1) theta in the tails (exp(-t) overflows at t < -709, curvature
+// was 0 below -20) and at t = +-0, where the value's own derivative
+// must not switch branch. Closed forms in every autodiff mode, scalar and
+// vector theta.
+TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_lpmf_tails) {
+  using bernoulli_logit_lpmf_test::exact;
+  using bernoulli_logit_lpmf_test::tol;
+  using stan::math::fvar;
+  using stan::math::var;
+  std::vector<int> y{1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0};
+  Eigen::VectorXd theta(13);
+  theta << -2500, 720, -30, 20.5, -0.25, 0, 0, -0.0, -0.0, 21, 21, 40, -19.5;
+  const int N = theta.size();
+
+  for (int n = 0; n < N; ++n) {
+    exact e(y[n], theta(n));
+    auto f = [&](const auto& v) {
+      return stan::math::bernoulli_logit_lpmf(y[n], v(0));
+    };
+    Eigen::VectorXd v = theta.segment(n, 1);
+    double fx;
+    Eigen::VectorXd grad;
+    Eigen::MatrixXd H;
+    stan::math::hessian(f, v, fx, grad, H);
+    EXPECT_NEAR(fx, e.lp, tol(e.lp)) << n;
+    EXPECT_NEAR(grad(0), e.d1, tol(e.d1)) << n;
+    EXPECT_NEAR(H(0, 0), e.d2, tol(e.d2)) << n;
+    stan::math::hessian<double>(f, v, fx, grad, H);  // fvar<fvar<double>>
+    EXPECT_NEAR(grad(0), e.d1, tol(e.d1)) << n;
+    EXPECT_NEAR(H(0, 0), e.d2, tol(e.d2)) << n;
+    std::vector<Eigen::MatrixXd> grad_H;
+    stan::math::grad_hessian(f, v, fx, H, grad_H);
+    EXPECT_NEAR(H(0, 0), e.d2, tol(e.d2)) << n;
+    EXPECT_NEAR(grad_H[0](0, 0), e.d3, tol(e.scale)) << n;
+
+    // derivative of the value itself
+    fvar<fvar<double>> theta_ff(fvar<double>(theta(n), 1), 0);
+    EXPECT_NEAR(stan::math::bernoulli_logit_lpmf(y[n], theta_ff).val_.d_, e.d1,
+                tol(e.d1))
+        << n;
+    stan::math::nested_rev_autodiff nested;
+    fvar<var> theta_fv(theta(n), 0);
+    stan::math::bernoulli_logit_lpmf(y[n], theta_fv).val_.grad();
+    EXPECT_NEAR(theta_fv.val_.adj(), e.d1, tol(e.d1)) << n;
+  }
+
+  auto f
+      = [&](const auto& v) { return stan::math::bernoulli_logit_lpmf(y, v); };
+  double fx;
+  Eigen::VectorXd grad;
+  Eigen::MatrixXd H;
+  stan::math::hessian(f, theta, fx, grad, H);
+  double lp = 0, lp_tol = 0;
+  for (int i = 0; i < N; ++i) {
+    exact e(y[i], theta(i));
+    lp += e.lp;
+    lp_tol += tol(e.lp);
+    EXPECT_NEAR(grad(i), e.d1, tol(e.d1)) << i;
+    for (int j = 0; j < N; ++j) {
+      EXPECT_NEAR(H(i, j), i == j ? e.d2 : 0, tol(e.d2)) << i << ", " << j;
+    }
+  }
+  EXPECT_NEAR(fx, lp, lp_tol);
+}