Skip to content

Optimizer: avoid heap allocations per evaluation and per iteration in L-BFGS/BFGS #3445

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
jachymb wants to merge 2 commits into
base: develop
Choose a base branch
from
+38 −33
Open

Optimizer: avoid heap allocations per evaluation and per iteration in L-BFGS/BFGS #3445

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
adf582c
Avoid allocations per evaluation and per iteration in the optimizers
jachymb Oct 17, 2026
23662fc
Make LBFGSUpdate::search_direction non-const now that it writes a wor…
jachymb Oct 19, 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
62 changes: 33 additions & 29 deletions src/stan/optimization/bfgs.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ class BFGSMinimizer {

protected:
FunctorType _func;
VectorT _gk, _gk_1, _xk_1, _xk, _pk, _pk_1;
VectorT _gk, _gk_1, _xk_1, _xk, _pk, _pk_1, _sk, _yk;
Scalar _fk, _fk_1, _alphak_1;
Scalar _alpha, _alpha0;
size_t _itNum;
Expand Down Expand Up @@ -164,7 +164,6 @@ class BFGSMinimizer {

int step() {
Scalar gradNorm, stepNorm;
VectorT sk, yk;
int retCode(0);
int resetB(0);

Expand Down Expand Up @@ -226,21 +225,21 @@ class BFGSMinimizer {
_gk.swap(_gk_1);
_pk.swap(_pk_1);

sk.noalias() = _xk - _xk_1;
yk.noalias() = _gk - _gk_1;
_sk.noalias() = _xk - _xk_1;
_yk.noalias() = _gk - _gk_1;

gradNorm = _gk.norm();
stepNorm = sk.norm();
stepNorm = _sk.norm();

// Update QN approximation
if (resetB) {
// If the QN approximation was reset, automatically scale it
// and update the step-size accordingly
Scalar B0fact = _qn.update(yk, sk, true);
Scalar B0fact = _qn.update(_yk, _sk, true);
_pk_1 /= B0fact;
_alphak_1 = _alpha * B0fact;
} else {
_qn.update(yk, sk);
_qn.update(_yk, _sk);
_alphak_1 = _alpha;
}
// Compute search direction for next step
Expand Down Expand Up @@ -288,7 +287,8 @@ class ModelAdaptor {
M &_model;
std::vector<int> _params_i;
std::ostream *_msgs;
std::vector<double> _x, _g;
std::vector<double> _x;
std::vector<stan::math::var> _ad;
size_t _fevals;

public:
Expand Down Expand Up @@ -327,36 +327,40 @@ class ModelAdaptor {
}
int operator()(const Eigen::Matrix<double, Eigen::Dynamic, 1> &x, double &f,
Eigen::Matrix<double, Eigen::Dynamic, 1> &g) {
using Eigen::Dynamic;
using Eigen::Matrix;
using stan::math::index_type;
using stan::model::log_prob_grad;
typedef typename index_type<Matrix<double, Dynamic, 1> >::type idx_t;

_x.resize(x.size());
for (idx_t i = 0; i < x.size(); i++)
_x[i] = x[i];

_fevals++;

// stan::model::log_prob_grad<true, jacobian> on a member vector of vars:
// no allocation per call and no copies through std::vector<double>
_ad.resize(x.size());
for (Eigen::Index i = 0; i < x.size(); ++i) {
_ad[i] = stan::math::var(x.coeff(i));
}
try {
f = -log_prob_grad<true, jacobian>(_model, _x, _params_i, _g, _msgs);
stan::math::var lp
= _model.template log_prob<true, jacobian>(_ad, _params_i, _msgs);
f = -lp.val();
lp.grad();
} catch (const std::domain_error &e) {
stan::math::recover_memory();
if (_msgs)
(*_msgs) << e.what() << std::endl;
return 1;
} catch (...) {
stan::math::recover_memory();
throw;
}
g.resize(x.size());
for (Eigen::Index i = 0; i < x.size(); ++i) {
g.coeffRef(i) = -_ad[i].adj();
}
stan::math::recover_memory();

g.resize(_g.size());
for (size_t i = 0; i < _g.size(); i++) {
if (!std::isfinite(_g[i])) {
if (_msgs)
*_msgs << "Error evaluating model log probability: "
"Non-finite gradient."
<< std::endl;
return 3;
}
g[i] = -_g[i];
if (!g.allFinite()) {
if (_msgs)
*_msgs << "Error evaluating model log probability: "
"Non-finite gradient."
<< std::endl;
return 3;
}

if (std::isfinite(f)) {
Expand Down
2 changes: 0 additions & 2 deletions src/stan/optimization/bfgs_linesearch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,6 @@ int WolfeLineSearch(FunctorType &func, Scalar &alpha, XType &x1,
Scalar alpha1(alpha);

Scalar prevF(f0);
XType prevDF(gradx0);
Scalar prevDFp(dfp);
Scalar newDFp;

Expand Down Expand Up @@ -287,7 +286,6 @@ int WolfeLineSearch(FunctorType &func, Scalar &alpha, XType &x1,

alpha0 = alpha1;
prevF = func_val;
std::swap(prevDF, gradx1);
prevDFp = newDFp;

alpha1 *= 10.0;
Expand Down
7 changes: 5 additions & 2 deletions src/stan/optimization/lbfgs_update.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,8 +70,9 @@ class LBFGSUpdate {
* direction gk.
* @param[in] gk Gradient direction.
**/
inline void search_direction(VectorT &pk, const VectorT &gk) const {
std::vector<Scalar> alphas(_buf.size());
inline void search_direction(VectorT &pk, const VectorT &gk) {
std::vector<Scalar> &alphas = _alphas;
alphas.resize(_buf.size());
typename stan::util::ring_buffer<UpdateT>::const_reverse_iterator buf_rit;
typename stan::util::ring_buffer<UpdateT>::const_iterator buf_it;
typename std::vector<Scalar>::const_iterator alpha_it;
Expand Down Expand Up @@ -105,6 +106,8 @@ class LBFGSUpdate {
protected:
stan::util::ring_buffer<UpdateT> _buf;
Scalar _gammak;
// workspace of search_direction, kept to avoid an allocation per call
std::vector<Scalar> _alphas;
};
} // namespace optimization
} // namespace stan
Expand Down
Loading