lbfgs evaluates the loss and gradient twice per step when its line search accepts the first trial #171

lbfgs() evaluates the loss and gradient about twice per step, even when its line search accepts the first trial. Each step evaluates them at its starting point. The previous step's accepted trial already computed them there, but the compiled step returns without keeping that gradient. optax.lbfgs() carries the accepted value and gradient in its state (value_and_grad_from_state) and evaluates about once per step.

Measured on numba over 1000 iterations, pytensor-ml takes 2.1 gradient evaluations per iteration on a small MLP, a mean-field Bayesian network and a heat-equation PINN. scipy's L-BFGS-B and optax.lbfgs() take 1.1 to 1.3. That's 1.4x to 1.9x scipy's time per iteration, while pytensor-ml's time per evaluation is equal to or lower than scipy's.

import numpy as np
import pytensor.tensor as pt

from pytensor_ml.optim import lbfgs_updates, zoom_line_search
from pytensor_ml.params import trainable
from pytensor_ml.pytensorf import function

A = np.array([[3.0, 0.5], [0.5, 1.0]])
p = trainable(np.array([5.0, -3.0]), name="w")
loss = 0.5 * p @ pt.constant(A) @ p
updates = lbfgs_updates(loss, [p], line_search=zoom_line_search())
trials = next(key for key in updates if key.name == "lbfgs/line_search/evaluations")
step = function([], loss, updates=updates)

step()
print(int(trials.get_value()))  # 1 trial, plus the evaluation at the starting point: 2 per step