qqn-torch 0.1.0__py3-none-any.whl

qqn_torch/__init__.py

@@ -0,0 +1,10 @@
+"""QQN (Quadratic Quasi-Newton) optimizer for PyTorch.
+
+A drop-in replacement for ``torch.optim.LBFGS`` that searches a quadratic
+path blending the steepest-descent and quasi-Newton directions.
+"""
+
+from qqn_torch.optimizer import QQN
+
+__all__ = ["QQN"]
+__version__ = "0.1.1"

qqn_torch/_flatten.py

@@ -0,0 +1,53 @@
+"""Flatten / unflatten helpers for a list of parameter tensors.
+
+All math in QQN happens on a single 1-D vector. These helpers mirror the
+way ``torch.optim.LBFGS`` gathers parameters and gradients.
+"""
+
+from typing import List
+
+import torch
+from torch import Tensor
+
+
+def gather_flat_grad(params: List[Tensor]) -> Tensor:
+    """Concatenate the gradients of ``params`` into one 1-D tensor.
+
+    Parameters whose ``.grad`` is ``None`` contribute zeros (matching the
+    convention used by ``torch.optim.LBFGS._gather_flat_grad``).
+    """
+    views = []
+    for p in params:
+        if p.grad is None:
+            view = p.new_zeros(p.numel())
+        elif p.grad.is_sparse:
+            view = p.grad.to_dense().view(-1)
+        else:
+            view = p.grad.view(-1)
+        views.append(view)
+    return torch.cat(views, dim=0)
+
+
+def gather_flat_params(params: List[Tensor]) -> Tensor:
+    """Concatenate the current parameter values into one 1-D tensor."""
+    return torch.cat([p.detach().reshape(-1) for p in params], dim=0)
+
+
+def add_flat_update(params: List[Tensor], update: Tensor, alpha: float = 1.0) -> None:
+    """In-place ``params += alpha * update`` for the flattened ``update``."""
+    offset = 0
+    for p in params:
+        numel = p.numel()
+        p.add_(update[offset : offset + numel].view_as(p), alpha=alpha)
+        offset += numel
+    assert offset == update.numel(), "update size mismatch"
+
+
+def set_flat_params(params: List[Tensor], flat: Tensor) -> None:
+    """In-place overwrite of ``params`` from the flattened ``flat``."""
+    offset = 0
+    for p in params:
+        numel = p.numel()
+        p.copy_(flat[offset : offset + numel].view_as(p))
+        offset += numel
+    assert offset == flat.numel(), "flat size mismatch"

qqn_torch/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Vendored third-party code. See individual modules for attribution."""

qqn_torch/_vendor/strong_wolfe.py

@@ -0,0 +1,159 @@
+"""Vendored copy of PyTorch's ``_strong_wolfe`` helper.
+
+Source: ``torch/optim/lbfgs.py`` (PyTorch project, BSD-style license).
+Reproduced here, lightly adapted, to avoid depending on a private symbol
+that may change or move across PyTorch versions.
+
+Original copyright (c) PyTorch Contributors. Licensed under the BSD license;
+see the PyTorch LICENSE file. Attribution retained per Section 7 of
+``pytorch.plan.md``.
+"""
+
+from typing import Callable, Tuple
+
+
+def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
+    if bounds is not None:
+        xmin_bound, xmax_bound = bounds
+    else:
+        xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
+    d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
+    d2_square = d1**2 - g1 * g2
+    if d2_square >= 0:
+        d2 = d2_square**0.5
+        if x1 <= x2:
+            min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
+        else:
+            min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
+        return min(max(min_pos, xmin_bound), xmax_bound)
+    else:
+        return (xmin_bound + xmax_bound) / 2.0
+
+
+def strong_wolfe(
+    obj_func: Callable[[float], Tuple[float, float]],
+    d_dot,
+    f0: float,
+    g0_dot_d: float,
+    c1: float = 1e-4,
+    c2: float = 0.9,
+    tolerance_change: float = 1e-9,
+    max_ls: int = 25,
+    t_init: float = 1.0,
+):
+    """One-dimensional strong-Wolfe search along a fixed direction.
+
+    ``obj_func(t)`` must return ``(f, gtd)`` where ``gtd`` is the directional
+    derivative ``<grad(x + t*dir), dir>`` at the probed point.
+
+    Returns ``(t, f, gtd, n_evals)``.
+    """
+    d_norm = d_dot
+    f_prev, g_prev = f0, g0_dot_d
+    t = t_init
+    t_prev = 0.0
+
+    f_new, g_new = obj_func(t)
+    ls_func_evals = 1
+    gtd_new = g_new
+    ls_iter = 0
+
+    # bracket phase
+    bracket: list = [0.0, t]
+    bracket_f: list = [f0, f_new]
+    bracket_g: list = [g0_dot_d, gtd_new]
+    bracketed = False
+    while ls_iter < max_ls:
+        if f_new > (f0 + c1 * t * g0_dot_d) or (ls_iter > 1 and f_new >= f_prev):
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, gtd_new]
+            bracketed = True
+            break
+        if abs(gtd_new) <= -c2 * g0_dot_d:
+            return t, f_new, gtd_new, ls_func_evals
+        if gtd_new >= 0:
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, gtd_new]
+            bracketed = True
+            break
+
+        min_step = t + 0.01 * (t - t_prev)
+        max_step = t * 10
+        tmp = t
+        t = _cubic_interpolate(
+            t_prev,
+            f_prev,
+            g_prev,
+            t,
+            f_new,
+            gtd_new,
+            bounds=(min_step, max_step),
+        )
+        t_prev = tmp
+        f_prev = f_new
+        g_prev = gtd_new
+        f_new, g_new = obj_func(t)
+        ls_func_evals += 1
+        gtd_new = g_new
+        ls_iter += 1
+
+    if not bracketed:
+        bracket = [0.0, t]
+        bracket_f = [f0, f_new]
+        bracket_g = [g0_dot_d, gtd_new]
+
+    # zoom phase
+    insuf_progress = False
+    low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
+    while ls_iter < max_ls:
+        if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
+            break
+        t = _cubic_interpolate(
+            bracket[0],
+            bracket_f[0],
+            bracket_g[0],
+            bracket[1],
+            bracket_f[1],
+            bracket_g[1],
+        )
+
+        eps = 0.1 * (max(bracket) - min(bracket))
+        if min(max(bracket) - t, t - min(bracket)) < eps:
+            if insuf_progress or t >= max(bracket) or t <= min(bracket):
+                if abs(t - max(bracket)) < abs(t - min(bracket)):
+                    t = max(bracket) - eps
+                else:
+                    t = min(bracket) + eps
+                insuf_progress = False
+            else:
+                insuf_progress = True
+        else:
+            insuf_progress = False
+
+        f_new, g_new = obj_func(t)
+        ls_func_evals += 1
+        gtd_new = g_new
+        ls_iter += 1
+
+        if f_new > (f0 + c1 * t * g0_dot_d) or f_new >= bracket_f[low_pos]:
+            bracket[high_pos] = t
+            bracket_f[high_pos] = f_new
+            bracket_g[high_pos] = gtd_new
+            low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
+        else:
+            if abs(gtd_new) <= -c2 * g0_dot_d:
+                break
+            elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
+                bracket[high_pos] = bracket[low_pos]
+                bracket_f[high_pos] = bracket_f[low_pos]
+                bracket_g[high_pos] = bracket_g[low_pos]
+            bracket[low_pos] = t
+            bracket_f[low_pos] = f_new
+            bracket_g[low_pos] = gtd_new
+
+    t = bracket[low_pos]
+    f_new = bracket_f[low_pos]
+    gtd_new = bracket_g[low_pos]
+    return t, f_new, gtd_new, ls_func_evals

qqn_torch/line_search/__init__.py

@@ -0,0 +1,35 @@
+"""Line-search registry: name -> LineSearch factory."""
+
+from typing import Union
+
+from qqn_torch.line_search.armijo import ArmijoLineSearch
+from qqn_torch.line_search.backtracking import BacktrackingLineSearch
+from qqn_torch.line_search.base import LineSearch, SearchResult
+from qqn_torch.line_search.fixed import FixedLineSearch
+from qqn_torch.line_search.strong_wolfe import StrongWolfeLineSearch
+
+__all__ = [
+    "LineSearch",
+    "SearchResult",
+    "ArmijoLineSearch",
+    "BacktrackingLineSearch",
+    "StrongWolfeLineSearch",
+    "FixedLineSearch",
+    "build_line_search",
+]
+
+
+def build_line_search(spec: Union[str, LineSearch]) -> LineSearch:
+    """Resolve a string shortcut or pass through a ``LineSearch`` instance."""
+    if not isinstance(spec, str):
+        return spec
+    name = spec.lower()
+    if name == "armijo":
+        return ArmijoLineSearch()
+    if name == "backtracking":
+        return BacktrackingLineSearch()
+    if name == "strong_wolfe":
+        return StrongWolfeLineSearch()
+    if name == "fixed":
+        return FixedLineSearch()
+    raise ValueError(f"unknown line search: {spec!r}")

qqn_torch/line_search/armijo.py

@@ -0,0 +1,121 @@
+"""Backtracking line search enforcing the Armijo sufficient-decrease rule.
+
+Starting from ``t = t_init`` we reduce ``t`` until
+
+    f(x + d(t)) <= f0 + c1 * t * <g, d'(0)>
+
+holds, or ``max_iter`` is exhausted. Because ``d'(0) = -g``, the directional
+derivative at the origin ``g0_dot_d0 = <g, d'(0)> = -||g||^2 <= 0`` is the
+slope used by the Armijo test, scaled by ``t`` as a proxy for arc length.
+
+Rather than fixed geometric contraction we use a safeguarded *quadratic*
+interpolation of the observed function values, which typically reduces the
+number of objective evaluations dramatically on the quadratic path. We also
+permit an initial *expansion* phase: because the QQN path is quadratic in
+``t`` and anchored to the (often well-scaled) quasi-Newton endpoint at
+``t = 1``, the best step frequently lies at or beyond ``t = 1``.
+"""
+
+from typing import Optional
+
+from qqn_torch.line_search.base import EvalFn, SearchResult
+from qqn_torch.path import path_point
+
+
+class ArmijoLineSearch:
+    def search(
+        self,
+        evaluate: EvalFn,
+        f0: float,
+        g0_dot_d0: float,
+        grad_dir,
+        qn_dir,
+        options: Optional[dict],
+    ) -> SearchResult:
+        opts = options or {}
+        c1 = opts.get("c1", 1e-4)
+        shrink = opts.get("shrink", 0.5)
+        t_init = opts.get("t_init", 1.0)
+        max_iter = opts.get("max_iter", 25)
+        # Whether to try growing past t_init when t_init already satisfies
+        # Armijo (the quadratic path often has its best point at t >= 1).
+        allow_expand = opts.get("allow_expand", True)
+        grow = opts.get("grow", 2.0)
+        t_max = opts.get("t_max", 16.0)
+        # Lower clamp on interpolation contraction to avoid tiny steps.
+        min_shrink = opts.get("min_shrink", 0.1)
+        max_shrink = opts.get("max_shrink", 0.5)
+
+        def armijo_ok(t: float, f_t: float) -> bool:
+            return f_t <= f0 + c1 * t * g0_dot_d0
+
+        n_evals = 0
+        best = None  # (t, f, g): best (lowest f) probe seen overall
+
+        def record(t, f_t, g_t):
+            nonlocal best
+            if best is None or f_t < best[1]:
+                best = (t, f_t, g_t)
+
+        # --- Evaluate the initial step ---
+        t = t_init
+        f_t, g_t = evaluate(t)
+        n_evals += 1
+        record(t, f_t, g_t)
+
+        if armijo_ok(t, f_t):
+            # Optionally try to expand: keep growing while Armijo holds and
+            # the function keeps decreasing. This lets the quadratic path
+            # reach its natural (often super-unit) optimum.
+            if allow_expand:
+                best_t, best_f, best_g = t, f_t, g_t
+                t_e = t
+                while t_e * grow <= t_max and n_evals < max_iter:
+                    t_e = t_e * grow
+                    f_e, g_e = evaluate(t_e)
+                    n_evals += 1
+                    record(t_e, f_e, g_e)
+                    if armijo_ok(t_e, f_e) and f_e < best_f:
+                        best_t, best_f, best_g = t_e, f_e, g_e
+                    else:
+                        break
+                update = path_point(best_t, grad_dir, qn_dir)
+                return SearchResult(best_t, best_f, best_g, update, n_evals, True)
+            update = path_point(t, grad_dir, qn_dir)
+            return SearchResult(t, f_t, g_t, update, n_evals, True)
+
+        # --- Backtracking with safeguarded quadratic interpolation ---
+        # We have f(0) = f0, f'(0) = g0_dot_d0, and f(t) = f_t.
+        # The quadratic q(s) matching these has its minimizer at
+        #   s* = -g0_dot_d0 * t^2 / (2 * (f_t - f0 - g0_dot_d0 * t)).
+        t_prev = t
+        f_prev = f_t
+        while n_evals < max_iter:
+            denom = 2.0 * (f_prev - f0 - g0_dot_d0 * t_prev)
+            if denom > 0:
+                t_quad = -g0_dot_d0 * t_prev * t_prev / denom
+                # Safeguard the interpolated step into a sane contraction
+                # band relative to the previous t.
+                lo = min_shrink * t_prev
+                hi = max_shrink * t_prev
+                t = min(max(t_quad, lo), hi)
+            else:
+                t = shrink * t_prev
+
+            f_t, g_t = evaluate(t)
+            n_evals += 1
+            record(t, f_t, g_t)
+
+            if armijo_ok(t, f_t):
+                update = path_point(t, grad_dir, qn_dir)
+                return SearchResult(t, f_t, g_t, update, n_evals, True)
+
+            t_prev = t
+            f_prev = f_t
+
+        # No Armijo-acceptable step found; return the best probe so the
+        # caller can still make (possibly non-decreasing) progress / stop.
+        assert best is not None
+        t_b, f_b, g_b = best
+        update = path_point(t_b, grad_dir, qn_dir)
+        return SearchResult(t_b, f_b, g_b, update, n_evals, f_b < f0)

qqn_torch/line_search/backtracking.py

@@ -0,0 +1,48 @@
+"""Plain backtracking line search (Armijo with fixed geometric contraction).
+
+Unlike :class:`ArmijoLineSearch` (which uses safeguarded quadratic
+interpolation and an expansion phase), this strategy performs simple
+geometric backtracking from ``t = t_init`` with a fixed ``shrink`` factor.
+Kept as a minimal, predictable baseline for benchmarking and as the
+explicit ``"backtracking"`` registry entry, matching the JAX reference.
+"""
+
+from typing import Optional
+
+from qqn_torch.line_search.base import EvalFn, SearchResult
+from qqn_torch.path import path_point
+
+
+class BacktrackingLineSearch:
+    def search(
+        self,
+        evaluate: EvalFn,
+        f0: float,
+        g0_dot_d0: float,
+        grad_dir,
+        qn_dir,
+        options: Optional[dict],
+    ) -> SearchResult:
+        opts = options or {}
+        c1 = opts.get("c1", 1e-4)
+        shrink = opts.get("shrink", 0.5)
+        t_init = opts.get("t_init", 1.0)
+        max_iter = opts.get("max_iter", 25)
+
+        t = t_init
+        n_evals = 0
+        best = None
+        for _ in range(max_iter):
+            f_t, g_t = evaluate(t)
+            n_evals += 1
+            if f_t <= f0 + c1 * t * g0_dot_d0:
+                update = path_point(t, grad_dir, qn_dir)
+                return SearchResult(t, f_t, g_t, update, n_evals, True)
+            if best is None or f_t < best[1]:
+                best = (t, f_t, g_t)
+            t *= shrink
+
+        assert best is not None
+        t_b, f_b, g_b = best
+        update = path_point(t_b, grad_dir, qn_dir)
+        return SearchResult(t_b, f_b, g_b, update, n_evals, f_b < f0)

qqn_torch/line_search/base.py

@@ -0,0 +1,38 @@
+"""LineSearch protocol and the shared result container.
+
+A line search traverses the quadratic path ``d(t)`` over ``t in [0, 1]``,
+evaluating candidate *states* ``x + d(t)`` and selecting the ``t`` (and step)
+that satisfies a sufficient-decrease condition.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Optional, Protocol, runtime_checkable
+
+from torch import Tensor
+
+
+@dataclass
+class SearchResult:
+    t_star: float  # accepted path parameter
+    f_new: float  # objective value at the accepted state
+    g_new: Tensor  # flat gradient at the accepted state
+    update: Tensor  # d(t_star): displacement from x
+    n_evals: int  # number of objective evaluations consumed
+    success: bool  # whether sufficient decrease was achieved
+
+
+# ``evaluate(t)`` returns (f, flat_grad) at the state x + d(t).
+EvalFn = Callable[[float], "tuple[float, Tensor]"]
+
+
+@runtime_checkable
+class LineSearch(Protocol):
+    def search(
+        self,
+        evaluate: EvalFn,
+        f0: float,
+        g0_dot_d0: float,
+        grad_dir: Tensor,
+        qn_dir: Tensor,
+        options: Optional[dict],
+    ) -> SearchResult: ...

qqn_torch/line_search/fixed.py

@@ -0,0 +1,24 @@
+"""Fixed-step "line search": always take t = t_fixed. Debug/baseline only."""
+
+from typing import Optional
+
+from qqn_torch.line_search.base import EvalFn, SearchResult
+
+
+class FixedLineSearch:
+    def search(
+        self,
+        evaluate: EvalFn,
+        f0: float,
+        g0_dot_d0: float,
+        grad_dir,
+        qn_dir,
+        options: Optional[dict],
+    ) -> SearchResult:
+        opts = options or {}
+        t = opts.get("t_fixed", 1.0)
+        f_t, g_t = evaluate(t)
+        from qqn_torch.path import path_point
+
+        update = path_point(t, grad_dir, qn_dir)
+        return SearchResult(t, f_t, g_t, update, 1, f_t < f0)

qqn_torch/line_search/strong_wolfe.py

@@ -0,0 +1,74 @@
+"""Strong-Wolfe line search wrapping the vendored ``strong_wolfe`` helper.
+
+The path is non-linear in ``t``, so the directional derivative along the
+path is ``<grad(x + d(t)), d'(t)>`` (not a constant direction). We feed the
+vendored 1-D search this state-dependent slope, which keeps the L-BFGS
+curvature updates well-conditioned.
+
+Note (per algorithm.md): strong-Wolfe can *over-restrict* the quadratic-path
+step on some problems; it is offered but not the default.
+"""
+
+from typing import Optional
+
+import torch
+
+from qqn_torch.line_search.base import EvalFn, SearchResult
+from qqn_torch.path import path_derivative, path_point
+from qqn_torch._vendor.strong_wolfe import strong_wolfe
+
+
+class StrongWolfeLineSearch:
+    def search(
+        self,
+        evaluate: EvalFn,
+        f0: float,
+        g0_dot_d0: float,
+        grad_dir,
+        qn_dir,
+        options: Optional[dict],
+    ) -> SearchResult:
+        opts = options or {}
+        c1 = opts.get("c1", 1e-4)
+        c2 = opts.get("c2", 0.9)
+        max_ls = opts.get("max_iter", 25)
+        tol_change = opts.get("tolerance_change", 1e-9)
+
+        # Cache the most recent full-gradient evaluation so we can recover
+        # the accepted gradient without an extra objective call.
+        cache: dict = {"t": None, "f": None, "g": None}
+
+        def obj_func(t):
+            f_t, g_t = evaluate(t)
+            cache["t"] = t
+            cache["f"] = f_t
+            cache["g"] = g_t
+            # Directional derivative along the (curved) path: <g, d'(t)>.
+            dt = path_derivative(t, grad_dir, qn_dir)
+            gtd = float(torch.dot(g_t, dt))
+            return f_t, gtd
+
+        # ``d_dot`` scales the bracket tolerance; use the path tangent norm.
+        d0 = path_derivative(0.0, grad_dir, qn_dir)
+        d_norm = float(d0.abs().max())
+
+        t, f_new, _gtd, n_evals = strong_wolfe(
+            obj_func,
+            d_norm,
+            f0,
+            g0_dot_d0,
+            c1=c1,
+            c2=c2,
+            tolerance_change=tol_change,
+            max_ls=max_ls,
+        )
+
+        # Recover the gradient at the accepted t. Reuse the cached value if
+        # the last probe matches; otherwise evaluate once more.
+        if cache["t"] == t and cache["g"] is not None:
+            f_acc, g_acc = cache["f"], cache["g"]
+        else:
+            f_acc, g_acc = evaluate(t)
+            n_evals += 1
+        update = path_point(t, grad_dir, qn_dir)
+        return SearchResult(t, f_acc, g_acc, update, n_evals, f_acc < f0)