diff --git a/docs/source/api/optim.rst b/docs/source/api/optim.rst index b3fb21d..e169585 100644 --- a/docs/source/api/optim.rst +++ b/docs/source/api/optim.rst @@ -28,6 +28,15 @@ Update rules adadelta lbfgs +Line searches +------------- + +.. autosummary:: + :toctree: generated/ + + zoom_line_search + LineSearch + Transforms ---------- diff --git a/pytensor_ml/optim/__init__.py b/pytensor_ml/optim/__init__.py index f064e85..57bd0f6 100644 --- a/pytensor_ml/optim/__init__.py +++ b/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", ] diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py index 335495b..c9011df 100644 --- a/pytensor_ml/optim/alias.py +++ b/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, ) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 510efab..9272e1f 100644 --- a/pytensor_ml/optim/line_search.py +++ b/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.""" diff --git a/pytensor_ml/optim/rules.py b/pytensor_ml/optim/rules.py index bb922b6..7bf015f 100644 --- a/pytensor_ml/optim/rules.py +++ b/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,6 +961,15 @@ 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. @@ -964,12 +977,13 @@ def lbfgs_updates( 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 diff --git a/tests/dispatch/jax/test_line_search.py b/tests/dispatch/jax/test_line_search.py index 7048083..6ffb509 100644 --- a/tests/dispatch/jax/test_line_search.py +++ b/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) diff --git a/tests/dispatch/mlx/test_line_search.py b/tests/dispatch/mlx/test_line_search.py index c2b1571..92162f0 100644 --- a/tests/dispatch/mlx/test_line_search.py +++ b/tests/dispatch/mlx/test_line_search.py @@ -8,6 +8,7 @@ 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, @@ -15,6 +16,7 @@ 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) diff --git a/tests/optim/test_composition.py b/tests/optim/test_composition.py index d38734a..f2f90e7 100644 --- a/tests/optim/test_composition.py +++ b/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.""" diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py index 3303be5..9dce312 100644 --- a/tests/optim/test_rules.py +++ b/tests/optim/test_rules.py @@ -1,5 +1,7 @@ import inspect +from itertools import pairwise + import numpy as np import pytensor import pytensor.tensor as pt @@ -20,6 +22,8 @@ adamax_updates, adamw, adamw_updates, + chain, + clip_by_global_norm, compile_train, cosine_schedule, lbfgs, @@ -34,6 +38,7 @@ sgd_updates, ) from pytensor_ml.optim import alias as alias_module +from pytensor_ml.optim.line_search import zoom_line_search from pytensor_ml.pytensorf import function from tests.optim.lbfgs_reference import dense_inverse_hessian @@ -723,13 +728,17 @@ def test_lbfgs_rejects_a_pair_whose_curvature_is_positive_but_negligible(): assert int(pairs_written.get_value()) == 0 -def test_lbfgs_stays_finite_after_it_converges(): +@pytest.mark.parametrize( + "line_search", [None, zoom_line_search()], ids=["fixed_step", "line_search"] +) +def test_lbfgs_stays_finite_after_it_converges(line_search): """Past the minimum the gradient changes underflow, so ``y . y`` reaches zero while ``y . s`` is still a positive subnormal; a pair admitted then stores an infinite ``1 / (y . s)`` and the next - step is NaN. float32 at any floatX, where the underflow arrives within a few steps.""" + step is NaN. float32 at any floatX, where the underflow arrives within a few steps. A search there + starts from a zero slope and accepts its first trial.""" p = params.trainable(np.ones(2, dtype="float32"), name="w") loss = (p**2).sum() - step = function([], loss, updates=lbfgs_updates(loss, [p])) + step = function([], loss, updates=lbfgs_updates(loss, [p], line_search=line_search)) for _ in range(10): step() @@ -758,14 +767,17 @@ def test_lbfgs_parameters_of_different_dtypes_reach_the_minimum(): np.testing.assert_allclose([u.get_value()[0], v.get_value()[0]], [1.0, -2.0], rtol=1e-5) -def test_lbfgs_resumes_its_trajectory_from_a_checkpoint(tmp_path): +@pytest.mark.parametrize( + "line_search", [None, zoom_line_search()], ids=["fixed_step", "line_search"] +) +def test_lbfgs_resumes_its_trajectory_from_a_checkpoint(tmp_path, line_search): """Every piece of the rule's state is named and saved, so a run restored mid-way retraces the steps it took the first time; a ring index or curvature left behind would desynchronize the memory from its order.""" 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, dtype=floatX) @ p - updates = lbfgs_updates(loss, [p], learning_rate=0.5, memory_size=2) + updates = lbfgs_updates(loss, [p], learning_rate=0.5, memory_size=2, line_search=line_search) state = list(updates) step = function([], loss, updates=updates) for _ in range(3): @@ -780,6 +792,98 @@ def test_lbfgs_resumes_its_trajectory_from_a_checkpoint(tmp_path): np.testing.assert_array_equal(second, first) +# optax.lbfgs() on Rosenbrock from (-1.2, 1), whose default search is the same zoom with a first trial of +# one: the loss before each of its first 12 steps. +OPTAX_LBFGS_ROSENBROCK = [ + 24.199999999999992, + 11.298089684996773, + 4.400456781386164, + 4.113928620500794, + 4.104804048918484, + 4.100624910419455, + 4.060149693984276, + 3.9626941083836766, + 2.6301041311456412, + 2.6234786168036246, + 2.558343445108792, + 2.2005231515714794, +] + + +def test_lbfgs_with_a_line_search_follows_optax_to_the_minimum(): + """At the default rate the first trial is the full quasi-Newton step, and the search accepts or + shortens it as optax's does, so the trajectory is optax's down to the minimum.""" + p = trainable(np.array([-1.2, 1.0]), name="x") + loss = pt.sum(100.0 * (p[1:] - p[:-1] ** 2) ** 2 + (1 - p[:-1]) ** 2) + step = compile_train(loss, lbfgs(line_search=zoom_line_search()), inputs=[]) + + losses = [float(step()) for _ in range(37)] + + np.testing.assert_allclose(losses[:12], OPTAX_LBFGS_ROSENBROCK, rtol=1e-10) + assert losses[-1] < 1e-20 + + +def test_lbfgs_line_search_accepts_the_first_trial_on_a_quadratic(): + """On a quadratic every step the rule proposes, the capped first gradient step included, satisfies + both Wolfe conditions, so every search stops after one evaluation and none fails.""" + 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, dtype=floatX) @ p + updates = lbfgs_updates(loss, [p], line_search=zoom_line_search()) + state = {key.name: key for key in updates} + step = function([], loss, updates=updates) + + evaluations, step_sizes = [], [] + for _ in range(6): + step() + evaluations.append(int(state["lbfgs/line_search/evaluations"].get_value())) + step_sizes.append(float(state["lbfgs/line_search/step_size"].get_value())) + + assert evaluations == [1] * 6 + assert step_sizes == [1.0] * 6 + assert int(state["lbfgs/line_search/failures"].get_value()) == 0 + + +def test_lbfgs_line_search_shortens_a_step_the_rate_overshoots(): + """At ten times the natural rate the first trial jumps far past the minimum along the direction, so + the search shortens it and the loss still falls at every step.""" + 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, dtype=floatX) @ p + updates = lbfgs_updates(loss, [p], learning_rate=10.0, line_search=zoom_line_search()) + state = {key.name: key for key in updates} + step = function([], loss, updates=updates) + + # Three steps reach the minimum exactly; past it there is nothing left to shorten + losses, step_sizes = [], [] + for _ in range(3): + losses.append(float(step())) + step_sizes.append(float(state["lbfgs/line_search/step_size"].get_value())) + + assert all(size < 1.0 for size in step_sizes) + assert all(later < earlier for earlier, later in pairwise(losses)) + + +def test_lbfgs_line_search_refuses_input_other_than_the_loss(): + p = trainable(np.zeros(2), name="w") + loss = ((p - 1.0) ** 2).sum() + + with pytest.raises(ValueError, match="needs the loss graph"): + lbfgs_updates([2 * (p - 1.0)], [p], line_search=zoom_line_search()) + with pytest.raises(ValueError, match="needs the loss graph"): + chain(clip_by_global_norm(1.0), lbfgs(line_search=zoom_line_search()))(loss, [p]) + + +def test_lbfgs_searches_by_default(): + """The alias defaults to the zoom search, which needs the loss, so precomputed gradients are refused + unless the fixed step is asked for.""" + p = trainable(np.zeros(2), name="w") + + with pytest.raises(ValueError, match="needs the loss graph"): + lbfgs()([2 * (p - 1.0)], [p]) + lbfgs(line_search=None)([2 * (p - 1.0)], [p]) + + def test_lbfgs_rejects_a_zero_memory_size(): p = trainable(np.zeros(2), name="w") with pytest.raises(ValueError, match="memory_size must be at least 1"):