Add Wolfe line search to Laplace approximation #3250

SteveBronder · 2025-10-24T21:04:55Z

Summary

This PR makes the following changes for the laplace approximation:

Adds a wolfe line search to the Newton solver used in the laplace approximation to improve convergence.

The example code provided from Laplace Bug when passing Eigen::Map in tuple of functor arguments #3205 fails on develop. The issue arose that the initial value of 0 for theta started the model in the tail of the distribution. The quick line search we did which only tested half of a newton step was not robust enough for this model to reach convergance. This PR adds a full wolfe line search to the Newton solver used in the laplace approximation to improve convergence in such cases.
The graphic below shows the difference in estimates of the log likelihood for laplace relative to integrate_1d on the roach test data plotted along the mu and sigma estimates. There is still a bias relative to integrate_1d as mu becomes negative and sigma becomes larger, but it is much nicer than before.

The main loop for laplace_marginal_density_est is expensive as it requires calculating either a diagonal hessian or block diagonal hessian with 2nd order autodiff. The wolfe line search only requires the gradients of the likelihood with respect to theta. So with that in mind the wolfe line search tries pretty aggressively get the best step size. If our initial step size is successful, we try to keep doubling until we hit a step size where the strong wolfe conditions fail and then return the information for the step right before that failure. If our initial step size does not satisfy strong wolfe then we do a bracketed zoom with cubic interpolation until till we find a step size that satisfies the strong wolfe conditions.
Tests for the wolfe line search are added to test/unit/math/laplace/wolfe_line_search.hpp.

Fixes bugs in the laplace approximation

Fix iteration mismatch between values when line search succeeds
In the last iteration of the laplace approximation we were returning the negative block diagonal hessian and derived matrices from the previous search. This is fine if the line search in that last step failed. But if the line search succeeds then we need to go back and recalculate the negative block diagonal hessian and it's derived quantities.
Breakup diagonal and block hessian functions
Previously we had one block_hessian function that calculated both the block hessian or the diagonal hessian at runtime. But this function is only used in places where we know at compile time whether we want a block or diagonal hessian. So I split out the two functions to avoid unnecessary runtime branching.
barzilai_borwein_step_size
For an initial step size estimate before each line search we use the Barzilai-Borwein method to get an estimate.
Adjoints of ll args only calculated once
Previously we calculated them eargerly in each laplace iteration. But they are not needed within the inner loop so we wait till we finish the inner search then calculate their adjoints once afterwards.
Calculate covariance once at the start and reuse throughout.
We were calculating the covariance matrix from inside of laplace_density_est, but this required us to then return it from that function and imo looked weird. So I pulled it out and now laplace_marginal_density_est is passed the covariance matrix.

Fixes numerical stability in laplace distributions
There were a few places where we could use log_sum_exp etc. so I made those changes.
Fixes "bug" in finite difference step size calculation

Changed from cube root of epsilon to epsilon^(1/7) for 6th order
The finite difference method in Stan was previously using stepsize optimzied a 2nd order method. But the code is a 6th order method. I modified finite_diff_stepsize to use epsilon^(1/7) instead of cbrt(epsilon). With this change all of the laplace tests pass with a much higher tolerance for precision.

Tests

All the AD tests now have a tighter tolerance for the laplace approximation.
There are also tests for the wolfe line search in test/unit/math/laplace/wolfe_line_search.hpp.

./runTests.py test/unit/math/laplace

Release notes

Improve laplace approximation with wolfe line search and bug fixes.

Checklist

Copyright holder: Steve Bronder

The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
- Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
- Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)
the basic tests are passing
- unit tests pass (to run, use: ./runTests.py test/unit)
- header checks pass, (make test-headers)
- dependencies checks pass, (make test-math-dependencies)
- docs build, (make doxygen)
- code passes the built in C++ standards checks (make cpplint)
the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested

…ch method

…line-search

… the lambdas

…lues for W, B, etc. are used

…order stepsize

WardBrian

Some low-hanging fruit

WardBrian · 2026-01-08T17:08:09Z

test/unit/math/laplace/csv_reader.hpp

I think it would be better to do what is done for e.g. sigmaz.hpp and just convert the test data to hpp files with a const member, rather than include a csv reader into the tests for these functions only

WardBrian · 2026-01-08T17:10:36Z

stan/math/rev/core/set_zero_adjoints.hpp

+          } else {
+            static_assert(
+                sizeof(std::decay_t<output_i_t>*) == 0,
+                "INTERNAL ERROR:(laplace_marginal_lpdf) set_zero_adjoints was "


Since we've moved this out of the laplace code, I think this string should be updated

WardBrian · 2026-01-08T17:11:06Z

stan/math/rev/core/set_zero_adjoints.hpp

+ * @param[in, out] output The output whose adjoints will be set to zero
+ */
+template <typename Output>
+inline void set_zero_adjoint(Output&& output) {


Function name and file name should match: adjoints

Actually this code was not even used so I can delete it

WardBrian · 2026-01-08T17:14:55Z

stan/math/rev/core/collect_adjoints.hpp

+        } else {
+          static_assert(
+              sizeof(std::decay_t<output_i_t>*) == 0,
+              "INTERNAL ERROR:(laplace_marginal_lpdf) collect_adjoints was "


Similar, if these will live in core they shouldn't mention laplace functions (applies to all the functions in this file)

WardBrian · 2026-01-08T17:20:47Z

stan/math/mix/functor/conditional_copy_and_promote.hpp

+      },
+      std::forward<Args>(args)...);
+}
+


q: should the internal namespace end here? deep/shallow copy look much more like normal functions than conditional_copy_and_promote does

It's only used in other internal functions so I think it is better to have in internal

…er passes

…to fix/wolfe-zoom1

avehtari · 2026-01-14T14:40:47Z

The current version seems to be more robust than before. I have tested with integrated LOO and leave-one-group-out cross-validation

Poisson varying intercept per observation (roaches)
beta-binomial with varying intercept and slope with several observations per group. In this case, hessian_block_size = 2. This is also challenging as beta-binomial may have non-log-concave likelihood. solver=1 seems to fail, but fallback to other solvers seems to work great! I got the best behavior wtting solver=2 which then sometimes fallbacks to solver=3

In both cases, laplace_marginal_tol() is fast enough that it woul be possible to use it also in the model block, but at this point I've focused on using them in generated quantities block

I'll continue experimenting with other models

SteveBronder · 2026-01-14T15:06:34Z

Awesome thank you!

WardBrian

C++ is complicated but manageable -- but someone else (@charlesm93) is gonna need to review the actual algorithmic pieces for correctness

stan/math/mix/functor/laplace_marginal_density.hpp

WardBrian · 2026-01-15T19:46:52Z

stan/math/mix/functor/laplace_marginal_density.hpp

    if constexpr (is_any_var_scalar_v<scalar_type_t<CovarArgs>>) {
-      [&covar_args_refs, &covar_args_adj, &md_est, &R, &s2,
-       &covariance_function, &msgs]() mutable {
-        const nested_rev_autodiff nested;


Steve note to self: double check that not having a nested here is fine

WardBrian · 2026-01-15T19:50:07Z

stan/math/mix/functor/laplace_marginal_density_estimator.hpp

+namespace internal {
+
+template <std::size_t N, typename Tuple, typename CheckType>
+inline constexpr bool is_tuple_type_v


Move to prim/meta / re-use existing programs there

WardBrian · 2026-01-15T19:53:12Z

stan/math/mix/functor/laplace_marginal_density_estimator.hpp

+template <typename Ops>
+inline constexpr auto tuple_to_laplace_options(Ops&& ops) {


Please clean up the manual typechecking here -- I'm not even sure if it is necessary, since the std::gets will raise a compiler error if they're wrong in the happy path

So I think I kind of like leaving these even though it is a bit bulky, std::get would throw a compiler error or the parameters may not match, but these compiler errors will be a little bit nicer for whoever sees them. Though the stanc3 compiler will also block users from seeing.

WardBrian · 2026-01-15T19:54:25Z

stan/math/mix/functor/laplace_marginal_density_estimator.hpp

+ * @note This helper is currently unused in the Laplace solvers in this file.
+ */
+template <typename WRootMat>
+inline void block_matrix_chol_L(WRootMat& W_root,


Unused function

WardBrian · 2026-01-15T20:42:12Z

stan/math/mix/functor/barzilai_borwein_step_size.hpp

+ * @warning The vectors must have identical size. Non-finite inputs yield the
+ *          safe fallback.
+ */
+inline double barzilai_borwein_step_size(const Eigen::VectorXd& s,


@charlesm93 the C++ looks fine to me for this function but I'd appreciate someone else putting eyes on the math

WardBrian · 2026-01-15T20:45:59Z

stan/math/mix/functor/wolfe_line_search.hpp

+ * The routine assumes a 1-D bracket [x_left, x_right], together with function
+ * values and directional derivatives at both endpoints. Internally it:
+ *
+ *   1. Normalizes the interval to s ∈ [0, 1] via


q: Is doxygen happy with the unicode etc?

I'm going to look for any more of these and switch them to math mode.

WardBrian · 2026-01-15T20:48:50Z

stan/math/mix/functor/wolfe_line_search.hpp

+  struct Candidate {
+    Scalar s_;
+    Scalar value_;
+  };
+  Candidate best{0.5, eval(0.5)};  // Start from bisection.


I think this can just be two variables, s_best and value_best

WardBrian · 2026-01-15T20:52:39Z

stan/math/mix/functor/wolfe_line_search.hpp

+  auto assign_step
+      = [](WolfeData& out, WolfeData& buf, auto&& e) { out.update(buf, e); };


WardBrian · 2026-01-15T20:54:29Z

stan/math/mix/functor/wolfe_line_search.hpp

+  auto armijo_ok = [&prev, &opt](const Eval& eval) -> bool {
+    return check_armijo(eval, prev, opt);
+  };
+  auto wolfe_ok = [&prev, &opt](const Eval& eval) -> bool {
+    return check_wolfe(eval, prev, opt);
+  };


SteveBronder · 2026-01-22T19:52:34Z

EDIT: Wrong PR

SteveBronder · 2026-01-22T19:53:59Z

@WardBrian for that static init failure on jenkins I need to bring back the csv reader so the test is not trying to put all of that data into static memory. I'm pretty sure that is the cause of the error

wrong branch for eigen update

WardBrian · 2026-01-26T15:04:17Z

@SteveBronder the latest couple commits here look like mistakes meant for #3271

SteveBronder added 30 commits September 29, 2025 18:03

update the laplace line search to use a more advanced wolfe line sear…

43c3eef

…ch method

Merge remote-tracking branch 'origin/develop' into fix/laplace-wolfe-…

74a92bc

…line-search

add data for roach data test

2de7a4f

update tests

547288e

move wolfe to its own file

cb5e282

Merge remote-tracking branch 'origin' into fix/laplace-wolfe-line-search

cdaf700

update wolfe

6b22c85

Merge remote-tracking branch 'origin' into fix/laplace-wolfe-line-search

a1d0906

update to use barzilai borwein step size as initial step size estimate

5b6ffff

seperate moto from other lpdf tests

8eff766

update

f542cc5

add WolfeInfo

c845944

use WolfeInfo for extra data

40f1243

put everything for iterations in laplace into structs

6e528d2

update poisson test

d89eeb5

add swap functions

40d889f

cleanup laplace_density_est to reduce repeated code

59b7a2f

update to search for a good initial alpha on a space

c73f5aa

fix code for wolfe line search

773d417

update tests for zoom

b557dad

all tests pass for laplace with new wolfe

b18bf87

use log sum of diagonal of U matrix for solver 3 determinant

98df588

move update_step to be a user passed function

929dd47

cleanup the laplace code to remove some passed by reference values to…

2ebb01a

… the lambdas

cleanup the laplace code to remove some passed by reference values to…

3bbcef3

… the lambdas

update WolfeData with member accessors and use Eval within WolfeData

66ffec9

update docs for wolfe

ff5bee4

update logic in laplace_marginal_desntiy_est so that final updated va…

7a7415a

…lues for W, B, etc. are used

clang format

973144a

change stepsize of finite difference to use 6th order instead of 2nd …

cc5d49a

…order stepsize

SteveBronder added 2 commits January 8, 2026 13:07

minor cleanup

e6b2d74

minor cleanup

27d2fc9

WardBrian reviewed Jan 8, 2026

View reviewed changes

SteveBronder and others added 12 commits January 8, 2026 14:51

cleanup retry logic on NaN of Inf objective function values

f845a43

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

200fd9a

adds fallthrough option in user opts. User ops are now a tuple the us…

d6859f1

…er passes

Merge remote-tracking branch 'refs/remotes/origin/fix/wolfe-zoom1' in…

d2de83d

…to fix/wolfe-zoom1

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

7f1a995

update docs

d86fa69

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

6abd9da

remove .csv files for .hpp files

b30823d

Merge remote-tracking branch 'refs/remotes/origin/fix/wolfe-zoom1' in…

dd1e74e

…to fix/wolfe-zoom1

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

b292cae

make non-static so it's memory is allocated dynamically.

ed3ce3a

remove rev/core/set_zero_adjoints.hpp

d9319ad

WardBrian requested changes Jan 15, 2026

View reviewed changes

SteveBronder and others added 5 commits January 20, 2026 14:56

update for @WardBrian review

a14b58d

Merge commit 'bf9fc6bc00539c5ddb31d4ed622b08b19ea36c44' into HEAD

b61a3d6

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

28467a2

update

d43587f

[Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1

10d5012

SteveBronder and others added 3 commits January 23, 2026 12:02

use csv reader for roach_data

c2bc877

fix doxygen for new Eigen

5d24366

Update Eigen library version in Doxygen config

968b4bb

wrong branch for eigen update

WardBrian mentioned this pull request Jan 26, 2026

Update laplace _tol functions for allow_fallthrough argument, tuple of control params stan-dev/stanc3#1586

Draft

2 tasks

		template <typename Ops>
		inline constexpr auto tuple_to_laplace_options(Ops&& ops) {

		auto assign_step
		= [](WolfeData& out, WolfeData& buf, auto&& e) { out.update(buf, e); };

Uh oh!

Add Wolfe line search to Laplace approximation #3250

Are you sure you want to change the base?

Add Wolfe line search to Laplace approximation #3250

Conversation

SteveBronder commented Oct 24, 2025

Summary

Tests

Release notes

Checklist

Uh oh!

WardBrian left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

avehtari commented Jan 14, 2026

Uh oh!

SteveBronder commented Jan 14, 2026

Uh oh!

WardBrian left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SteveBronder commented Jan 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SteveBronder commented Jan 22, 2026

Uh oh!

WardBrian commented Jan 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

SteveBronder commented Jan 22, 2026 •

edited

Loading