Search every L-BFGS step with a zoom line search by default

docs/source/api/optim.rst

@@ -28,6 +28,15 @@ Update rules
     adadelta
     lbfgs
 
+Line searches
+-------------
+
+.. autosummary::
+    :toctree: generated/
+
+    zoom_line_search
+    LineSearch
+
 Transforms
 ----------
 

pytensor_ml/optim/__init__.py

@@ -35,6 +35,7 @@
     nonfinite,
     skip_if,
 )
+from pytensor_ml.optim.line_search import LineSearch, zoom_line_search
 from pytensor_ml.optim.policy import reduce_on_plateau
 from pytensor_ml.optim.rules import (
     adadelta_updates,
@@ -70,6 +71,7 @@
     "Decision",
     "Gradients",
     "LearningRate",
+    "LineSearch",
     "Rate",
     "Schedule",
     "SkipCondition",
@@ -123,4 +125,5 @@
     "to_floatx",
     "to_updates",
     "trace",
+    "zoom_line_search",
 ]

pytensor_ml/optim/alias.py

@@ -7,6 +7,7 @@
     Transform,
     Updates,
 )
+from pytensor_ml.optim.line_search import LineSearch, zoom_line_search
 from pytensor_ml.optim.rules import (
     _require_numeric_learning_rate,
     adadelta_updates,
@@ -363,6 +364,7 @@ def lbfgs(
     memory_size: int = 10,
     scale_init_precond: bool = True,
     *,
+    line_search: LineSearch | None = zoom_line_search(),
     namespace: str = "lbfgs",
 ) -> Transform:
     """
@@ -373,10 +375,15 @@ def lbfgs(
 
     Examples
     --------
-    A quasi-Newton direction from a memory of recent parameter and gradient differences, taken at a
-    fixed fraction with no line search. It reads the change between consecutive gradients as curvature,
-    so the loss has to be the same function from one step to the next: full batch, no dropout. For the
-    same reason it takes the loss's own gradients: put a clip after it in a chain, never ahead of it.
+    A quasi-Newton direction from a memory of recent parameter and gradient differences, searched along
+    by a strong Wolfe line search with ``learning_rate`` as the first trial. ``line_search=None`` takes
+    every step at ``learning_rate`` instead, which is also what
+    :func:`~pytensor_ml.optim.rules.lbfgs_updates` does unless given a search. On the mlx backend every
+    search runs all ``max_steps`` of its trials, so each step costs that many loss and gradient
+    evaluations, and a smaller ``zoom_line_search(max_steps=...)`` budget makes it cheaper. It reads the
+    change between consecutive gradients as curvature, so the loss has to be the same function from one
+    step to the next: full batch, no dropout. For the same reason it takes the loss's own gradients: put
+    a clip after it in a chain, never ahead of it.
 
     .. code-block:: python
 
@@ -402,6 +409,7 @@ def rule(
             learning_rate=learning_rate,
             memory_size=memory_size,
             scale_init_precond=scale_init_precond,
+            line_search=line_search,
             namespace=namespace,
         )
 

pytensor_ml/optim/line_search.py

@@ -23,7 +23,10 @@
 class LineSearch(Protocol):
     """A step-size search written as one trial per :meth:`step`, which :func:`search_along` repeats."""
 
-    max_steps: int
+    @property
+    def max_steps(self) -> int:
+        """Return the most trials one search may take."""
+        ...
 
     def init(self, value: TensorVariable, slope: TensorVariable) -> State:
         """Return the state before the first trial, from the loss and its slope at a step of zero."""

pytensor_ml/optim/rules.py

@@ -24,6 +24,7 @@
     to_floatx,
 )
 from pytensor_ml.optim.lbfgs import LBFGSDirection, flat_dot
+from pytensor_ml.optim.line_search import LineSearch, search_along
 from pytensor_ml.params import step_counter
 
 
@@ -913,6 +914,7 @@ def lbfgs_updates(
     learning_rate: LearningRate = 1.0,
     memory_size: int = 10,
     scale_init_precond: bool = True,
+    line_search: LineSearch | None = None,
     namespace: str = "lbfgs",
 ) -> Updates:
     r"""
@@ -928,18 +930,20 @@ def lbfgs_updates(
     step is :math:`p \leftarrow p - \eta H g`.
 
     The direction is well scaled once the memory holds a pair, so :math:`\eta = 1` is the natural rate.
-    The rule takes every step at that rate: there is no line search yet (pymc-devs/pytensor-ml#58), so
-    the rate is the only safeguard against a bad direction. Consecutive gradients have to be measured on
-    the same objective for their difference to be curvature, so the rule assumes a deterministic,
-    full-batch loss.
+    Given a ``line_search``, the rule searches along :math:`-\eta H g` for a step size :math:`t` and
+    moves to :math:`p - t \eta H g`, so ``learning_rate`` is the first trial step. Without one it takes
+    every step at the rate, which is then the only safeguard against a bad direction. Consecutive
+    gradients have to be measured on the same objective for their difference to be curvature, and a
+    search compares the loss at several points, so the rule assumes a deterministic, full-batch loss.
 
     Three uses break that assumption. A gradient transform ahead of the rule in a chain, such as
     :func:`~pytensor_ml.optim.clipping.clip_by_global_norm`, hands it gradients whose differences are not
-    curvature, and the rule can diverge. Clip after the rule instead: the parameter differences are read
-    off the parameters, so a clipped step still forms a valid pair. Wrapping the rule in
-    :func:`~pytensor_ml.optim.guards.skip_if` does not rescue a bad step either, because the loss is
-    deterministic and a skipped step is recomputed unchanged on the next call until the guard raises. A
-    step that a guard would skip calls for a smaller rate. Finally, a parameter written between steps,
+    curvature, and the rule can diverge. With a line search the rule refuses that input outright. Clip
+    after the rule instead: the parameter differences are read off the parameters, so a clipped step
+    still forms a valid pair. Wrapping the rule in :func:`~pytensor_ml.optim.guards.skip_if` does not
+    rescue a bad step either, because the loss is deterministic and a skipped step is recomputed
+    unchanged on the next call until the guard raises. A step that a guard would skip calls for a line
+    search, which shortens it, or a smaller rate. Finally, a parameter written between steps,
     with ``set_value`` for instance, forms a pair from a move the rule did not make, and that pair stays
     in the memory for up to ``memory_size`` steps.
 
@@ -957,19 +961,29 @@ def lbfgs_updates(
     scale_init_precond : bool
         Start the recursion from :math:`\gamma I` as above. When False it starts from the identity, and
         the first step is the raw gradient. Default True.
+    line_search : LineSearch, optional
+        Search for a step size along each step, with ``learning_rate`` as the first trial, for example
+        :func:`~pytensor_ml.optim.zoom_line_search`. The search evaluates the loss, so the first argument
+        has to be the loss itself. Its step size multiplies the step already scaled by ``learning_rate``,
+        so a bound such as ``max_learning_rate`` caps that multiplier. The rule records the last
+        multiplier in ``{namespace}/line_search/step_size``, the last search's trial count in
+        ``{namespace}/line_search/evaluations``, and the number of searches that have failed so far in
+        ``{namespace}/line_search/failures``. Default None, which takes every step at
+        ``learning_rate``. The :func:`~pytensor_ml.optim.lbfgs` alias searches by default.
     namespace : str
         Prefix for every state slot this rule allocates, so two rules in one graph keep separate state
         rather than reusing each other's. Default is the rule's own name.
 
     Returns
     -------
     updates : Updates
-        Mapping from each parameter and its memory buffers to their next values.
+        Mapping from each parameter, its memory buffers and, with a search, the search's records to
+        their next values.
 
     Examples
     --------
     Compile the step yourself rather than going through :func:`~pytensor_ml.optim.train.compile_train`.
-    The rule returns the updates dict directly, with no line search:
+    The rule returns the updates dict directly, and takes the fixed step unless given a search:
 
     .. code-block:: python
 
@@ -989,6 +1003,15 @@ def lbfgs_updates(
     """
     if memory_size < 1:
         raise ValueError(f"memory_size must be at least 1, got {memory_size}.")
+    loss = (
+        loss_gradients_or_updates if isinstance(loss_gradients_or_updates, TensorVariable) else None
+    )
+    if line_search is not None and loss is None:
+        raise ValueError(
+            f"{namespace} searches along its step by evaluating the loss, so with a line search it needs "
+            f"the loss graph rather than gradients or an updates dict. It therefore has to come first in "
+            f"a chain, where the loss still reaches it."
+        )
 
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
     step_count = step_counter(f"{namespace}/step_count")
@@ -1086,6 +1109,25 @@ def lbfgs_updates(
         updates[previous_gradients[index]] = gradients[index]
         updates[value_memory[index]] = new_value_memory[index]
         updates[gradient_memory[index]] = new_gradient_memory[index]
-        updates[parameter] = parameter - learning_rate * directions[index]
+
+    steps = [
+        (-learning_rate * direction).astype(parameter.dtype)
+        for parameter, direction in zip(parameters, directions)
+    ]
+    if line_search is None or loss is None:
+        for parameter, step in zip(parameters, steps):
+            updates[parameter] = parameter + step
+        return updates
+
+    search = search_along(loss, parameters, gradients, steps, line_search)
+    for parameter, step in zip(parameters, steps):
+        updates[parameter] = parameter + search.step_size.astype(parameter.dtype) * step
+
+    step_size = scalar_state(f"{namespace}/line_search/step_size", fill_value=1.0)
+    evaluations = scalar_state(f"{namespace}/line_search/evaluations", dtype="int64")
+    failures = scalar_state(f"{namespace}/line_search/failures", dtype="int64")
+    updates[step_size] = search.step_size.astype(step_size.dtype)
+    updates[evaluations] = search.evaluations.astype(evaluations.dtype)
+    updates[failures] = failures + search.failed.astype(failures.dtype)
 
     return updates

tests/dispatch/jax/test_line_search.py

@@ -5,9 +5,14 @@
 
 pytest.importorskip("jax")
 
+from pytensor.compile.mode import get_mode
+
+from pytensor_ml.optim import compile_train, lbfgs
 from pytensor_ml.optim.line_search import search_along, zoom_line_search
+from pytensor_ml.params import trainable
 from tests.dispatch.jax.test_basic import jax_mode
 from tests.optim.test_line_search import OBJECTIVES, PINNED, ROSENBROCK_START, D, X
+from tests.optim.test_rules import OPTAX_LBFGS_ROSENBROCK
 
 
 @pytest.mark.parametrize("case, expected", PINNED.values(), ids=PINNED.keys())
@@ -77,3 +82,15 @@ def test_the_search_on_jax_matches_the_default_backend(
 
     np.testing.assert_allclose(float(on_jax[0]), float(on_default[0]), rtol=1e-9)
     assert (bool(on_jax[1]), int(on_jax[2])) == (bool(on_default[1]), int(on_default[2]))
+
+
+def test_lbfgs_with_a_line_search_follows_optax_on_jax():
+    """The whole training step on JAX, under the backend's own compile mode: the rule's step needs
+    the fast_run rewrites that the bare JAX mode of these tests leaves out."""
+    point = trainable(ROSENBROCK_START.copy(), name="x")
+    loss = OBJECTIVES["rosenbrock"](point)
+    step = compile_train(loss, lbfgs(), inputs=[], compile_kwargs={"mode": get_mode("JAX")})
+
+    losses = [float(step()) for _ in range(12)]
+
+    np.testing.assert_allclose(losses, OPTAX_LBFGS_ROSENBROCK, rtol=1e-10)

tests/dispatch/mlx/test_line_search.py

@@ -8,13 +8,15 @@
 from pytensor.compile.mode import Mode
 from pytensor.link.mlx.linker import MLXLinker
 
+from pytensor_ml.optim import compile_train, lbfgs
 from pytensor_ml.optim.line_search import (
     STEP,
     LineSearchOp,
     find_piece,
     search_along,
     zoom_line_search,
 )
+from pytensor_ml.params import trainable
 from tests.dispatch.mlx.test_basic import mx
 from tests.optim.test_line_search import OBJECTIVES, PINNED
 
@@ -109,3 +111,25 @@ def step_nodes_under(mode):
         return len(find_piece(search, STEP).fgraph.apply_nodes)
 
     assert step_nodes_under(mode) < step_nodes_under(mode.excluding("fusion"))
+
+
+def test_lbfgs_with_a_line_search_trains_as_on_the_default_backend():
+    """The whole training step on mlx, at single precision, with a short trial budget: mlx runs every
+    trial the budget allows, so the budget sets what each step costs here."""
+    A = np.array([[3.0, 0.5], [0.5, 1.0]], dtype="float32")
+    start = np.array([5.0, -3.0], dtype="float32")
+    point = trainable(start.copy(), name="w")
+    loss = 0.5 * point @ pt.constant(A) @ point
+
+    def losses_on(mode):
+        point.set_value(start.copy())
+        # The rule's own state and rate follow floatX, which has to be single precision on the GPU too
+        with pytensor.config.change_flags(floatX="float32"):
+            rule = lbfgs(line_search=zoom_line_search(max_steps=5))
+            step = compile_train(loss, rule, inputs=[], compile_kwargs={"mode": mode})
+        return [float(step()) for _ in range(8)]
+
+    on_mlx = losses_on(Mode(linker=MLXLinker(), optimizer="fast_run"))
+    on_default = losses_on(None)
+
+    np.testing.assert_allclose(on_mlx, on_default, rtol=1e-4, atol=1e-12)

tests/optim/test_composition.py

@@ -23,6 +23,7 @@
     scale_by_schedule,
     sgd,
     skip_if,
+    zoom_line_search,
 )
 from pytensor_ml.params import step_counter, trainable
 from pytensor_ml.pytensorf import function
@@ -55,11 +56,15 @@ def test_lbfgs_converges_with_its_step_clipped_after_it():
     np.testing.assert_allclose(p.get_value(), 0.0, atol=1e-6)
 
 
-def test_skip_if_holds_back_all_of_the_lbfgs_state():
+@pytest.mark.parametrize(
+    "line_search", [None, zoom_line_search()], ids=["fixed_step", "line_search"]
+)
+def test_skip_if_holds_back_all_of_the_lbfgs_state(line_search):
     """A skipped step must leave the memory exactly as it was, or the next applied step pairs a stale
-    previous gradient with a fresh parameter and stores a pair that is not a secant."""
+    previous gradient with a fresh parameter and stores a pair that is not a secant. The line search's
+    counters are held back with it."""
     _, loss = quadratic_problem()
-    step = compile_train(loss, skip_if(lbfgs(), max_consecutive_skips=None))
+    step = compile_train(loss, skip_if(lbfgs(line_search=line_search), max_consecutive_skips=None))
     for _ in range(2):
         step(GOOD)
     lbfgs_state = [variable for variable in step.get_shared() if "lbfgs/" in str(variable.name)]
@@ -74,6 +79,20 @@ def test_skip_if_holds_back_all_of_the_lbfgs_state():
     assert int(pairs_written.get_value()) == int(before["lbfgs/pairs_written"]) + 1
 
 
+def test_lbfgs_line_search_reaches_the_minimum_under_a_plateau_policy():
+    """The plateau policy hands the rule the loss, which the search needs, and its scale becomes the
+    search's first trial."""
+    p, loss = quadratic_problem()
+    plateau_scale = scalar_state("plateau/scale", fill_value=1.0)
+    rule = lbfgs(learning_rate=plateau_scale, line_search=zoom_line_search())
+    step = compile_train(loss, reduce_on_plateau(rule, plateau_scale, patience=2))
+
+    for _ in range(10):
+        step(GOOD)
+
+    np.testing.assert_allclose(p.get_value(), 0.0, atol=1e-6)
+
+
 def test_clipping_bounds_a_rules_step_end_to_end():
     """The clipping transform is otherwise only exercised on a hand-built updates dict; here it has to
     survive a real rule, a real gradient, and compile_train's assembly."""