Skip to content

Embedded Laplace: closed-form derivatives for the built-in likelihoods #3471

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
avehtari wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

Embedded Laplace: closed-form derivatives for the built-in likelihoods #3471

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1f33f4a
closed-form derivatives for the built-in embedded Laplace likelihoods
avehtari Oct 19, 2026
d7d7502
fix cpplint whitespace/braces in the analytic-derivatives test
avehtari Oct 19, 2026
31cd426
share group_sums, test tail values against references, remove unused …
avehtari Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 141 additions & 84 deletions stan/math/mix/functor/laplace_likelihood.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@
#include <stan/math/mix/functor/hessian_times_vector.hpp>

#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

namespace stan {
namespace math {
Expand All @@ -21,7 +23,68 @@ namespace math {
*/
namespace laplace_likelihood {

/**
* Specialize this trait with `value = true` for a likelihood functor `F` if
* the Hessian of its log likelihood with respect to `theta` is diagonal and
* its derivatives with respect to `theta` are known in closed form. The
* specialization must have three static member functions, which take
* `theta` and the arguments of `F::operator()` without the message stream,
* all with arithmetic scalar types, and return an `Eigen::VectorXd`:
* - `theta_grad(theta, args...)`: the gradient;
* - `neg_hessian_diag(theta, args...)`: the diagonal of the negative
* Hessian;
* - `third_diff(theta, args...)`: the diagonal of the third derivative.
* When all arguments are arithmetic, `theta_grad()`, `diagonal_hessian()`,
* `block_hessian()` and `third_diff()` below then use these functions
* instead of nested autodiff.
* @tparam F Type of the likelihood functor
*/
template <typename F>
struct analytic_derivatives : std::false_type {};

namespace internal {

/**
* Sum of the observations and number of observations of each group, for the
* closed-form derivatives of the built-in likelihoods.
* @tparam YVec Type of the observations
* @param[in] theta_size number of groups
* @param[in] y observations
* @param[in] y_index group of each observation (1-based)
* @param[out] counts sum of `y` in each group
* @param[out] n number of observations in each group
*/
template <typename YVec>
inline void group_sums(Eigen::Index theta_size, const YVec& y,
const std::vector<int>& y_index, Eigen::VectorXd& counts,
Eigen::VectorXd& n) {
counts = Eigen::VectorXd::Zero(theta_size);
n = Eigen::VectorXd::Zero(theta_size);
for (size_t i = 0; i < y_index.size(); i++) {
counts(y_index[i] - 1) += y[i];
n(y_index[i] - 1) += 1;
}
}

/**
* True if the closed-form derivatives of `F` (see `analytic_derivatives`)
* are used for these argument types: `F` has them, and `theta` and all
* arguments have arithmetic scalar types.
*/
template <typename F, typename Theta, typename... Args>
inline constexpr bool use_analytic_derivatives_v
= analytic_derivatives<std::decay_t<F>>::value&&
std::is_arithmetic<scalar_type_t<std::decay_t<Theta>>>::value&&
std::conjunction_v<
std::is_arithmetic<scalar_type_t<std::decay_t<Args>>>...>;

template <typename F, typename Theta, typename... Args>
using require_analytic_derivatives_t = require_t<
std::bool_constant<use_analytic_derivatives_v<F, Theta, Args...>>>;

template <typename F, typename Theta, typename... Args>
using require_not_analytic_derivatives_t = require_t<
std::bool_constant<!use_analytic_derivatives_v<F, Theta, Args...>>>;
/**
* @tparam F A functor with `opertor()(Args&&...)` returning a scalar
* @tparam Theta A class assignable to an Eigen vector type
Expand Down Expand Up @@ -53,7 +116,8 @@ inline auto log_likelihood(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
* @param args Variadic arguments for the likelihood function.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_not_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline auto theta_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
using Eigen::Dynamic;
using Eigen::Matrix;
Expand Down Expand Up @@ -104,7 +168,8 @@ inline void ll_arg_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
* @param args Variadic arguments for the likelihood function.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_not_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline auto diagonal_hessian(F&& f, Theta&& theta, Stream* msgs,
Args&&... args) {
using Eigen::Dynamic;
Expand Down Expand Up @@ -134,7 +199,8 @@ inline auto diagonal_hessian(F&& f, Theta&& theta, Stream* msgs,
* @param args Variadic arguments for the likelihood function.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_not_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline auto block_hessian(F&& f, Theta&& theta,
const Eigen::Index hessian_block_size, Stream* msgs,
Args&&... args) {
Expand All @@ -159,57 +225,6 @@ inline auto block_hessian(F&& f, Theta&& theta,
}
}

/**
* Computes theta gradient and negative block diagonal Hessian of `f` wrt
* `theta` and `args...`
* @note If `Args` contains \ref var types then their adjoints will be
* calculated as a side effect.
* @tparam F A functor with `opertor()(Args&&...)` returning a scalar
* @tparam Theta A class assignable to an Eigen vector type
* @tparam Stream Type of stream for messages.
* @tparam Args Type of variadic arguments.
* @param f Log likelihood function.
* @param theta Latent Gaussian model.
* @param hessian_block_size If the Hessian of the log likelihood function w.r.t
* the latent Gaussian variable is block-diagonal,
* size of each block.
* @param msgs Stream for messages.
* @param args Variadic arguments for the likelihood function.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline auto diff(F&& f, Theta&& theta, const Eigen::Index hessian_block_size,
Stream* msgs, Args&&... args) {
using Eigen::Dynamic;
using Eigen::Matrix;
const Eigen::Index theta_size = theta.size();
auto theta_gradient = [&theta, &f, &msgs](auto&&... args) {
nested_rev_autodiff nested;
Matrix<var, Dynamic, 1> theta_var = theta;
var f_var = f(theta_var, args..., msgs);
grad(f_var.vi_);
return theta_var.adj().eval();
}(args...);
if (hessian_block_size == 1) {
auto v = Eigen::VectorXd::Ones(theta_size);
Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
value_of(args)..., msgs);
Eigen::SparseMatrix<double> hessian_theta(theta_size, theta_size);
hessian_theta.reserve(Eigen::VectorXi::Constant(theta_size, 1));
for (Eigen::Index i = 0; i < theta_size; i++) {
hessian_theta.insert(i, i) = hessian_v(i);
}
return std::make_pair(std::move(theta_gradient), (-hessian_theta).eval());
} else {
return std::make_pair(
std::move(theta_gradient),
(-hessian_block_diag(f, std::forward<Theta>(theta), hessian_block_size,
value_of(args)..., msgs))
.eval());
}
}

/**
* Compute third order derivative of `f` wrt `theta` and `args...`
* @note If `Args` contains \ref var types then their adjoints will be
Expand All @@ -224,7 +239,8 @@ inline auto diff(F&& f, Theta&& theta, const Eigen::Index hessian_block_size,
* @param args Variadic arguments for likelihood function.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_t<Theta>* = nullptr>
require_eigen_vector_t<Theta>* = nullptr,
require_not_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline Eigen::VectorXd third_diff(F&& f, Theta&& theta, Stream&& msgs,
Args&&... args) {
nested_rev_autodiff nested;
Expand Down Expand Up @@ -355,6 +371,76 @@ inline auto diff_eta_implicit(F&& f, V_t&& v, Theta&& theta, Stream* msgs,
grad(f_sum.d_.vi_);
}

/**
* Gradient of `f` with respect to `theta`, from the closed-form derivatives
* of `F` (see `analytic_derivatives`).
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline Eigen::VectorXd theta_grad(F&& f, Theta&& theta, Stream* msgs,
Args&&... args) {
return analytic_derivatives<std::decay_t<F>>::theta_grad(theta, args...);
}

/**
* Diagonal of the negative Hessian of `f` with respect to `theta`, from the
* closed-form derivatives of `F` (see `analytic_derivatives`).
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline Eigen::VectorXd diagonal_hessian(F&& f, Theta&& theta, Stream* msgs,
Args&&... args) {
return analytic_derivatives<std::decay_t<F>>::neg_hessian_diag(theta,
args...);
}

/**
* Negative Hessian of `f` with respect to `theta` as a sparse block-diagonal
* matrix, from the closed-form derivatives of `F` (see
* `analytic_derivatives`). The Hessian of such an `F` is diagonal; the
* matrix has the sparsity pattern of the autodiff version
* (`hessian_block_diag()`): full blocks of size `hessian_block_size`, with
* zero off-diagonal entries.
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr,
require_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline Eigen::SparseMatrix<double> block_hessian(
F&& f, Theta&& theta, const Eigen::Index hessian_block_size, Stream* msgs,
Args&&... args) {
const Eigen::VectorXd neg_hessian_diag
= analytic_derivatives<std::decay_t<F>>::neg_hessian_diag(theta, args...);
const Eigen::Index theta_size = theta.size();
Eigen::SparseMatrix<double> neg_hessian(theta_size, theta_size);
neg_hessian.reserve(
Eigen::VectorXi::Constant(theta_size, hessian_block_size));
const Eigen::Index n_blocks = theta_size / hessian_block_size;
for (Eigen::Index j = 0; j < n_blocks; ++j) {
const Eigen::Index start = j * hessian_block_size;
for (Eigen::Index c = start; c < start + hessian_block_size; ++c) {
for (Eigen::Index r = start; r < start + hessian_block_size; ++r) {
neg_hessian.insert(r, c) = r == c ? neg_hessian_diag(r) : 0.0;
}
}
}
neg_hessian.makeCompressed();
return neg_hessian;
}

/**
* Diagonal of the third derivative of `f` with respect to `theta`, from the
* closed-form derivatives of `F` (see `analytic_derivatives`).
*/
template <typename F, typename Theta, typename Stream, typename... Args,
require_eigen_vector_t<Theta>* = nullptr,
require_analytic_derivatives_t<F, Theta, Args...>* = nullptr>
inline Eigen::VectorXd third_diff(F&& f, Theta&& theta, Stream&& msgs,
Args&&... args) {
return analytic_derivatives<std::decay_t<F>>::third_diff(theta, args...);
}

} // namespace internal

/**
Expand Down Expand Up @@ -456,35 +542,6 @@ inline auto log_likelihood(F&& f, Theta&& theta, TupleArgs&& ll_tup,
std::forward<Theta>(theta), msgs);
}

/**
* A wrapper that accepts a tuple as arguments.
* @tparam F A functor with `opertor()(Args&&...)` returning a scalar
* @tparam Theta A class assignable to an Eigen vector type
* @tparam TupleArgs Type of arguments for covariance function.
* @tparam Stream Type of stream for messages.
* @param f Log likelihood function.
* @param theta Latent Gaussian model.
* @param hessian_block_size If Hessian of log likelihood w.r.t theta is
* block diagonal, size of block.
* @param ll_tuple Arguments for likelihood function
* @param msgs Stream messages.
*/
template <typename F, typename Theta, typename TupleArgs, typename Stream,
require_eigen_vector_t<Theta>* = nullptr,
require_tuple_t<TupleArgs>* = nullptr>
inline auto diff(F&& f, Theta&& theta, const Eigen::Index hessian_block_size,
TupleArgs&& ll_tuple, Stream* msgs) {
return stan::math::apply(
[](auto&& f, auto&& theta, auto hessian_block_size, auto* msgs,
auto&&... args) {
return internal::diff(
std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
hessian_block_size, msgs, std::forward<decltype(args)>(args)...);
},
std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
std::forward<Theta>(theta), hessian_block_size, msgs);
}

/**
* A wrapper that accepts a tuple as arguments.
* @tparam F Type of log likelhood function.
Expand Down
54 changes: 54 additions & 0 deletions stan/math/mix/prob/laplace_marginal_bernoulli_logit_lpmf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

#include <stan/math/mix/meta.hpp>
#include <stan/math/mix/functor/laplace_likelihood.hpp>
#include <stan/math/prim/fun/inv_logit.hpp>
#include <stan/math/mix/functor/laplace_marginal_density.hpp>
#include <stan/math/rev/fun/elt_multiply.hpp>
#include <stan/math/rev/fun/exp.hpp>
Expand All @@ -20,6 +21,8 @@
#include <stan/math/prim/fun/binomial_coefficient_log.hpp>
#include <stan/math/rev/fun/to_vector.hpp>
#include <stan/math/prim/fun/to_vector.hpp>
#include <type_traits>
#include <vector>

namespace stan {
namespace math {
Expand Down Expand Up @@ -58,6 +61,57 @@ struct bernoulli_logit_likelihood {
}
};

namespace laplace_likelihood {
/**
* Closed-form derivatives of `bernoulli_logit_likelihood` with respect to
* `theta` (see `analytic_derivatives`). With `p_g = inv_logit(theta_g +
* mean_g)` and `q_g = 1 - p_g = inv_logit(-theta_g - mean_g)` (computed
* without cancellation), the gradient is `counts_g - n_g p_g`, the negative
* Hessian is `n_g p_g q_g`, and the third derivative is
* `-n_g p_g q_g (q_g - p_g)`.
*/
template <>
struct analytic_derivatives<bernoulli_logit_likelihood> : std::true_type {
template <typename Theta, typename YVec, typename Mean>
static void probabilities(const Theta& theta, const YVec& y,
const std::vector<int>& y_index, const Mean& mean,
Eigen::VectorXd& counts, Eigen::VectorXd& n,
Eigen::VectorXd& p, Eigen::VectorXd& q) {
internal::group_sums(theta.size(), y, y_index, counts, n);
Eigen::VectorXd theta_offset = add(theta, mean);
p = inv_logit(theta_offset);
q = inv_logit(-theta_offset);
}

template <typename Theta, typename YVec, typename Mean>
static Eigen::VectorXd theta_grad(const Theta& theta, const YVec& y,
const std::vector<int>& y_index,
const Mean& mean) {
Eigen::VectorXd counts, n, p, q;
probabilities(theta, y, y_index, mean, counts, n, p, q);
return counts - n.cwiseProduct(p);
}

template <typename Theta, typename YVec, typename Mean>
static Eigen::VectorXd neg_hessian_diag(const Theta& theta, const YVec& y,
const std::vector<int>& y_index,
const Mean& mean) {
Eigen::VectorXd counts, n, p, q;
probabilities(theta, y, y_index, mean, counts, n, p, q);
return n.cwiseProduct(p).cwiseProduct(q);
}

template <typename Theta, typename YVec, typename Mean>
static Eigen::VectorXd third_diff(const Theta& theta, const YVec& y,
const std::vector<int>& y_index,
const Mean& mean) {
Eigen::VectorXd counts, n, p, q;
probabilities(theta, y, y_index, mean, counts, n, p, q);
return -n.cwiseProduct(p).cwiseProduct(q).cwiseProduct(q - p);
}
};
} // namespace laplace_likelihood

/**
* Wrapper function around the laplace_marginal function for
* a logistic Bernoulli likelihood. Returns the marginal density
Expand Down
Loading
Loading