Add Wolfe line search to Laplace approximation #3250
Conversation
…lues for W, B, etc. are used
|
@SteveBronder you mentioned I need to check the math on this. Is there any particular file or document you want me to look at? |
|
@charlesm93 the main thing I'd like you to look at is my wolfe impl. I would just like a spot check that how I'm doing that is reasonable. Also if you lookover the @avehtari below is a script for pulling down a fresh command stan and building it with this branch. If you can try this out it would be appreciated! |
There was a problem hiding this comment.
Two questions:
- is this change actually relevant to the PR? It doesn't look like it's called anywhere in laplace directly
- the stencil in the couple places that the code does call it is always 2, so maybe it could be simpler?
There was a problem hiding this comment.
Yeah this can be reverted. Technically I think we are using the wrong finite diff stepsize for the order of finite differences we use, but now I don't need to fix this in this PR and we can do this after some discussion in a separate PR
| return -covariance * step.a() + covariance * step.theta_grad(); | ||
| }; | ||
| auto update_step = [&covariance, &obj_fun, &theta_grad_f, &grad_fun]( | ||
| auto& step_info, auto&& /* curr */, auto&& prev, |
There was a problem hiding this comment.
Should I find it odd that curr is unused? Why have the argument at all?
| try { | ||
| if (options.solver == 1) { | ||
| if (options.hessian_block_size == 1) { | ||
| // std::cout << "Solver: 1Diag" << std::endl; |
There was a problem hiding this comment.
| // std::cout << "Solver: 1Diag" << std::endl; |
There was a problem hiding this comment.
There are a few other commented cout calls as well
There was a problem hiding this comment.
There are a few test files that changed and I can't immediately tell if the actual values are different or just the code layout -- are these relevant to the PR? (this one and mdivide_left/mdivide_right)
There was a problem hiding this comment.
Thanks. Let me look this over. This is from a previous change I need to remove in this PR.
|
|
||
| stan::math::matrix_d I = Eigen::MatrixXd::Identity(2, 2); | ||
| EXPECT_MATRIX_FLOAT_EQ(I, stan::math::mdivide_left(Ad, Ad)); | ||
| EXPECT_MATRIX_NEAR(I, stan::math::mdivide_left(Ad, Ad), 1e-15); |
There was a problem hiding this comment.
This one is because of local failures on my desktop
Jenkins Console Log Machine informationNo LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focalCPU: G++: Clang: |
Jenkins Console Log Machine informationNo LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focalCPU: G++: Clang: |
… use across the stan math library. Adds docs for laplace helper functions. clean up control logic in reverse mode autodiff laplace approximation.
Jenkins Console Log Machine informationNo LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focalCPU: G++: Clang: |
Summary
This PR makes the following changes for the laplace approximation:
thetastarted 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
laplacerelative tointegrate_1don the roach test data plotted along the mu and sigma estimates. There is still a bias relative tointegrate_1das mu becomes negative and sigma becomes larger, but it is much nicer than before.laplace_marginal_density_estis 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.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.
Previously we had one
block_hessianfunction 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.For an initial step size estimate before each line search we use the Barzilai-Borwein method to get an estimate.
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.
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 nowlaplace_marginal_density_estis passed the covariance matrix.There were a few places where we could use
log_sum_expetc. so I made those changes.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_stepsizeto 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.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
./runTests.py test/unit)make test-headers)make test-math-dependencies)make doxygen)make cpplint)the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested