Accurate derivatives of inv_logit and log_inv_logit_diff at large arguments

stan/math/fwd/fun/inv_logit.hpp

@@ -18,8 +18,9 @@ namespace math {
  */
 template <typename T, require_fvar_t<T>* = nullptr>
 inline auto inv_logit(T&& x) {
-  return std::decay_t<T>(inv_logit(x.val_),
-                         x.d_ * inv_logit(x.val_) * (1 - inv_logit(x.val_)));
+  const auto val = inv_logit(x.val_);
+  // inv_logit(-x) rather than 1 - inv_logit(x), which cancels for x >> 0
+  return std::decay_t<T>(val, x.d_ * val * inv_logit(-x.val_));
 }
 
 }  // namespace math

stan/math/fwd/fun/log_inv_logit_diff.hpp

@@ -21,11 +21,11 @@ namespace math {
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^x}{e^y-e^x}-\frac{e^x}{e^x+1}
+    \frac{\partial}{\partial x} = \mbox{logit}^{-1}(-x) + \frac{1}{e^{x-y}-1}
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^y}{e^x-e^y}-\frac{e^y}{e^y+1}
+    \frac{\partial}{\partial y} = -\mbox{logit}^{-1}(y) - \frac{1}{e^{x-y}-1}
    \f]
  *
  * @tparam T inner type of the fvar
@@ -36,16 +36,17 @@ namespace math {
  */
 template <typename T>
 inline fvar<T> log_inv_logit_diff(const fvar<T>& x, const fvar<T>& y) {
+  // both partials as sums of positive terms, no cancellation for x >> 0
+  const auto c = inv(expm1(x.val_ - y.val_));
   return fvar<T>(
       log_inv_logit_diff(x.val_, y.val_),
-      -x.d_ * (inv(expm1(y.val_ - x.val_)) + inv_logit(x.val_))
-          - y.d_ * (inv(expm1(x.val_ - y.val_)) + inv_logit(y.val_)));
+      x.d_ * (c + inv_logit(-x.val_)) - y.d_ * (c + inv_logit(y.val_)));
 }
 
 template <typename T>
 inline fvar<T> log_inv_logit_diff(const fvar<T>& x, double y) {
   return fvar<T>(log_inv_logit_diff(x.val_, y),
-                 -x.d_ * (inv(expm1(y - x.val_)) + inv_logit(x.val_)));
+                 x.d_ * (inv(expm1(x.val_ - y)) + inv_logit(-x.val_)));
 }
 
 template <typename T>

stan/math/opencl/rev/inv_logit.hpp

@@ -21,7 +21,7 @@ inline var_value<matrix_cl<double>> inv_logit(const var_value<T>& A) {
   return make_callback_var(
       inv_logit(A.val()), [A](vari_value<matrix_cl<double>>& res) mutable {
         A.adj() += elt_multiply(res.adj(),
-                                elt_multiply(res.val(), 1.0 - res.val()));
+                                elt_multiply(res.val(), inv_logit(-A.val())));
       });
 }
 

stan/math/opencl/rev/log_inv_logit_diff.hpp

@@ -38,13 +38,11 @@ inline var_value<matrix_cl<double>> log_inv_logit_diff(T_x&& x, T_y&& y) {
   return make_callback_var(
       log_inv_logit_diff(value_of(x_arena), value_of(y_arena)),
       [x_arena, y_arena](const vari_value<matrix_cl<double>>& res) mutable {
+        // both partials as sums of positive terms, no cancellation for x >> 0
+        auto c = inv(expm1(value_of(x_arena) - value_of(y_arena)));
         adjoint_results(x_arena, y_arena) += expressions(
-            -elt_multiply(res.adj(),
-                          inv(expm1(value_of(y_arena) - value_of(x_arena)))
-                              + inv_logit(value_of(x_arena))),
-            -elt_multiply(res.adj(),
-                          inv(expm1(value_of(x_arena) - value_of(y_arena)))
-                              + inv_logit(value_of(y_arena))));
+            elt_multiply(res.adj(), c + inv_logit(-value_of(x_arena))),
+            -elt_multiply(res.adj(), c + inv_logit(value_of(y_arena))));
       });
 }
 

stan/math/prim/fun/log_inv_logit_diff.hpp

@@ -20,11 +20,11 @@ namespace math {
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^x}{e^y-e^x}-\frac{e^x}{e^x+1}
+    \frac{\partial}{\partial x} = \mbox{logit}^{-1}(-x) + \frac{1}{e^{x-y}-1}
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^y}{e^x-e^y}-\frac{e^y}{e^y+1}
+    \frac{\partial}{\partial y} = -\mbox{logit}^{-1}(y) - \frac{1}{e^{x-y}-1}
    \f]
  *
  * @tparam T1 type of x argument

stan/math/rev/fun/inv_logit.hpp

@@ -3,7 +3,12 @@
 
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core.hpp>
+#include <stan/math/prim/fun/abs.hpp>
+#include <stan/math/prim/fun/as_array_or_scalar.hpp>
+#include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/inv_logit.hpp>
+#include <stan/math/prim/functor/apply_scalar_binary.hpp>
+#include <cmath>
 
 namespace stan {
 namespace math {
@@ -15,20 +20,27 @@ namespace math {
  *
  * The derivative of inverse logit is
  *
- * \f$\frac{d}{dx} \mbox{logit}^{-1}(x) = \mbox{logit}^{-1}(x) (1 -
- * \mbox{logit}^{-1}(x))\f$.
+ * \f$\frac{d}{dx} \mbox{logit}^{-1}(x) = \frac{e}{(1 + e)^2}\f$
+ * with \f$e = \exp(-|x|)\f$.
  *
- * @tparam T Arithmetic or a type inheriting from `EigenBase`.
+ * @tparam T a `var` or a `var_value` of a matrix type
  * @param a Argument variable.
- * @return Inverse logit of argument.
+ * @return Inverse logit of argument, elementwise for a matrix.
  */
-template <typename T, require_var_t<T>* = nullptr>
+template <
+    typename T, require_var_t<T>* = nullptr,
+    require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto inv_logit(T&& a) {
-  return make_callback_var(inv_logit(a.val()), [a](auto& vi) mutable {
-    as_array_or_scalar(a).adj() += as_array_or_scalar(vi.adj())
-                                   * as_array_or_scalar(vi.val())
-                                   * (1.0 - as_array_or_scalar(vi.val()));
-  });
+  auto e = to_arena(exp(-abs(a.val())));
+  return make_callback_var(
+      apply_scalar_binary(
+          [](double x, double e) { return (x >= 0 ? 1.0 : e) / (1.0 + e); },
+          a.val(), e),
+      [a, e](auto& vi) mutable {
+        const auto& e_a = as_array_or_scalar(e);
+        as_array_or_scalar(a.adj())
+            += as_array_or_scalar(vi.adj()) * e_a / ((1.0 + e_a) * (1.0 + e_a));
+      });
 }
 
 /**
@@ -38,20 +50,28 @@ inline auto inv_logit(T&& a) {
  *
  * The derivative of inverse logit is
  *
- * \f$\frac{d}{dx} \mbox{logit}^{-1}(x) = \mbox{logit}^{-1}(x) (1 -
- * \mbox{logit}^{-1}(x))\f$.
+ * \f$\frac{d}{dx} \mbox{logit}^{-1}(x) = \frac{e}{(1 + e)^2}\f$
+ * with \f$e = \exp(-|x|)\f$. It is computed in the same pass as the value
+ * and stored for the reverse pass.
  *
  * @tparam T type of Eigen expression
  * @param x Eigen expression
- * @return Inverse logit of argument.
+ * @return Inverse logit of each value in x.
  */
 template <typename T, require_eigen_vt<is_var, T>* = nullptr>
 inline auto inv_logit(T&& x) {
   auto x_arena = to_arena(std::forward<T>(x));
-  arena_t<T> ret = inv_logit(x_arena.val());
-  reverse_pass_callback([x_arena, ret]() mutable {
-    x_arena.adj().array()
-        += ret.adj().array() * ret.val().array() * (1.0 - ret.val().array());
+  arena_t<T> ret(x_arena.rows(), x_arena.cols());
+  arena_t<plain_type_t<decltype(x_arena.val())>> d(x_arena.rows(),
+                                                   x_arena.cols());
+  for (Eigen::Index i = 0; i < x_arena.size(); ++i) {
+    const double xi = x_arena.coeff(i).val();
+    const double e = std::exp(-std::fabs(xi));
+    ret.coeffRef(i) = (xi >= 0 ? 1.0 : e) / (1.0 + e);
+    d.coeffRef(i) = e / ((1.0 + e) * (1.0 + e));
+  }
+  reverse_pass_callback([x_arena, ret, d]() mutable {
+    x_arena.adj().array() += ret.adj().array() * d.array();
   });
   return ret;
 }

stan/math/rev/fun/log_inv_logit_diff.hpp

@@ -21,11 +21,11 @@ namespace math {
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^x}{e^y-e^x}-\frac{e^x}{e^x+1}
+    \frac{\partial}{\partial x} = \mbox{logit}^{-1}(-x) + \frac{1}{e^{x-y}-1}
    \f]
 
    \f[
-    \frac{\partial }{\partial x} = -\frac{e^y}{e^x-e^y}-\frac{e^y}{e^y+1}
+    \frac{\partial}{\partial y} = -\mbox{logit}^{-1}(y) - \frac{1}{e^{x-y}-1}
    \f]
  *
  * @tparam T1 Type of x argument
@@ -41,11 +41,10 @@ class log_inv_logit_diff_vv_vari : public op_vv_vari {
   log_inv_logit_diff_vv_vari(vari* avi, vari* bvi)
       : op_vv_vari(log_inv_logit_diff(avi->val_, bvi->val_), avi, bvi) {}
   void chain() {
-    avi_->adj_
-        -= adj_ * (inv(expm1(bvi_->val_ - avi_->val_)) + inv_logit(avi_->val_));
-
-    bvi_->adj_
-        -= adj_ * (inv(expm1(avi_->val_ - bvi_->val_)) + inv_logit(bvi_->val_));
+    // both partials as sums of positive terms, no cancellation for x >> 0
+    double c = inv(expm1(avi_->val_ - bvi_->val_));
+    avi_->adj_ += adj_ * (c + inv_logit(-avi_->val_));
+    bvi_->adj_ -= adj_ * (c + inv_logit(bvi_->val_));
   }
 };
 
@@ -54,7 +53,8 @@ class log_inv_logit_diff_vd_vari : public op_vd_vari {
   log_inv_logit_diff_vd_vari(vari* avi, double b)
       : op_vd_vari(log_inv_logit_diff(avi->val_, b), avi, b) {}
   void chain() {
-    avi_->adj_ -= adj_ * (inv(expm1(bd_ - avi_->val_)) + inv_logit(avi_->val_));
+    avi_->adj_
+        += adj_ * (inv(expm1(avi_->val_ - bd_)) + inv_logit(-avi_->val_));
   }
 };
 

test/unit/math/mix/constraint/lub_constrain_1_test.cpp

@@ -1,5 +1,6 @@
 #include <test/unit/math/test_ad.hpp>
 #include <test/unit/math/mix/constraint/lub_constrain_helpers.hpp>
+#include <test/unit/math/mix/fun/inv_logit_tail_refs.hpp>
 
 TEST(mathMixMatFun, lub_constrain_scalars) {
   double x1 = 0.7;
@@ -24,3 +25,12 @@ TEST(mathMixMatFun, lub_constrain_scalars) {
   lub_constrain_tests::expect(x1, lb_inf, ub_inf);
   lub_constrain_tests::expect(x2, lb_inf, ub_inf);
 }
+
+TEST(mathMixMatFun, lub_constrain_fwd_derivative_tails) {
+  using inv_logit_tail_refs::ref;
+  // derivatives of -2 + 5 inv_logit(x)
+  inv_logit_tail_refs::expect_fwd_derivatives(
+      [](const auto& x) { return stan::math::lub_constrain(x, -2.0, 3.0); },
+      [](const ref& r) { return 5 * r.d1; },
+      [](const ref& r) { return 5 * r.d2; });
+}

test/unit/math/mix/fun/inv_logit_tail_refs.hpp

@@ -0,0 +1,86 @@
+#ifndef TEST_UNIT_MATH_MIX_FUN_INV_LOGIT_TAIL_REFS_HPP
+#define TEST_UNIT_MATH_MIX_FUN_INV_LOGIT_TAIL_REFS_HPP
+
+#include <stan/math/mix.hpp>
+#include <gtest/gtest.h>
+#include <cmath>
+#include <limits>
+#include <vector>
+
+namespace inv_logit_tail_refs {
+
+// inv_logit(x), inv_logit(-x), inv_logit'(x) and inv_logit''(x), computed
+// with mpmath at 60 digits and rounded to double (0 where they underflow).
+struct ref {
+  double x;
+  double s;
+  double s_neg;
+  double d1;
+  double d2;
+};
+
+inline std::vector<ref> refs() {
+  return {{-800, 0, 1.0, 0, 0},
+          {-40, 4.248354255291589e-18, 1.0, 4.248354255291589e-18,
+           4.248354255291589e-18},
+          {-37, 8.533047625744065e-17, 0.9999999999999999,
+           8.533047625744065e-17, 8.533047625744063e-17},
+          {-30, 9.357622968839299e-14, 0.9999999999999064,
+           9.357622968838423e-14, 9.357622968836672e-14},
+          {-20, 2.0611536181902037e-09, 0.9999999979388464,
+           2.061153613941849e-09, 2.061153605445141e-09},
+          {-1, 0.2689414213699951, 0.7310585786300049, 0.19661193324148185,
+           0.09085774767294841},
+          {0, 0.5, 0.5, 0.25, 0},
+          {1, 0.7310585786300049, 0.2689414213699951, 0.19661193324148185,
+           -0.09085774767294841},
+          {20, 0.9999999979388464, 2.0611536181902037e-09,
+           2.061153613941849e-09, -2.061153605445141e-09},
+          {30, 0.9999999999999064, 9.357622968839299e-14, 9.357622968838423e-14,
+           -9.357622968836672e-14},
+          {37, 0.9999999999999999, 8.533047625744065e-17, 8.533047625744065e-17,
+           -8.533047625744063e-17},
+          {40, 1.0, 4.248354255291589e-18, 4.248354255291589e-18,
+           -4.248354255291589e-18},
+          {100, 1.0, 3.720075976020836e-44, 3.720075976020836e-44,
+           -3.720075976020836e-44},
+          {700, 1.0, 9.85967654375977e-305, 9.85967654375977e-305,
+           -9.85967654375977e-305},
+          {800, 1.0, 0, 0, 0}};
+}
+
+// Relative check, absolute below the smallest normal double.
+inline void expect_rel(double expected, double actual, double x,
+                       double tol = 1e-13) {
+  EXPECT_NEAR(
+      expected, actual,
+      tol * std::fmax(std::fabs(expected), std::numeric_limits<double>::min()))
+      << "x = " << x;
+}
+
+// First and second derivatives of f at every reference point through
+// fvar<double>, fvar<fvar<double>> and fvar<var>, against d1(r) and d2(r).
+template <typename F, typename D1, typename D2>
+void expect_fwd_derivatives(const F& f, const D1& d1, const D2& d2) {
+  using stan::math::fvar;
+  using stan::math::var;
+  for (const auto& r : refs()) {
+    expect_rel(d1(r), f(fvar<double>(r.x, 1)).d_, r.x);
+
+    fvar<fvar<double>> ff(fvar<double>(r.x, 1), fvar<double>(1, 0));
+    fvar<fvar<double>> y_ff = f(ff);
+    expect_rel(d1(r), y_ff.d_.val_, r.x);
+    expect_rel(d2(r), y_ff.d_.d_, r.x);
+
+    fvar<var> fv(r.x, 1);
+    fvar<var> y_fv = f(fv);
+    expect_rel(d1(r), y_fv.d_.val(), r.x);
+    y_fv.d_.grad();
+    expect_rel(d2(r), fv.val_.adj(), r.x);
+    stan::math::recover_memory();
+  }
+}
+
+}  // namespace inv_logit_tail_refs
+
+#endif