polystep 0.7.0__py3-none-any.whl

polystep/projection/__init__.py

@@ -0,0 +1,5 @@
+"""Projection module for memory-efficient large-scale subspace compression."""
+
+from .sparse import SparseRandomProjection
+
+__all__ = ["SparseRandomProjection"]

polystep/projection/sparse.py

@@ -0,0 +1,291 @@
+"""Sparse random projection for memory-efficient large-scale subspace.
+
+Implements an Achlioptas-style sparse projection (Rademacher entries at
+randomly sampled positions) at the very-sparse density of Li, Hastie &
+Church (2006), preserving Johnson-Lindenstrauss distance guarantees
+while using O(nnz) memory.
+
+References:
+    Achlioptas, D. (2003). "Database-friendly random projections:
+        Johnson-Lindenstrauss with binary coins."
+    Li, P., Hastie, T. & Church, K. (2006). "Very sparse random
+        projections."
+"""
+
+import math
+import warnings
+from typing import Optional
+
+import torch
+
+
+# Below this ratio of subspace-dim to source-dim, the Johnson-
+# Lindenstrauss distance guarantees stop holding for typical
+# optimization workloads. Empirical floor: GPT-2 124M projected to
+# 128-dim (ratio ~ 1e-6) collapsed to random predictions in early
+# experiments.
+_EXTREME_COMPRESSION_RATIO = 1e-5
+
+
+class SparseRandomProjection:
+    """Sparse random projection for memory-efficient large-scale subspace.
+
+    Achlioptas-style sparse projection: each column has Rademacher (+1/-1)
+    entries at randomly sampled positions, scaled by 1/sqrt(nnz_per_col).
+    The default density ``1/sqrt(full_dim)`` follows Li, Hastie & Church
+    (2006), which preserves JL distance guarantees while keeping memory
+    at O(nnz).
+
+    The projection matrix P is (full_dim x subspace_dim) with:
+    - Each column has exactly nnz_per_col = density * full_dim nonzeros
+    - Values are +1 or -1 (Rademacher) scaled by 1/sqrt(nnz_per_col)
+
+    Memory usage: O(density * full_dim * subspace_dim) = O(nnz)
+    For 1M params, rank-256, density 1/sqrt(1M) = 0.1%:
+    - Dense: 1M * 256 * 4 bytes = 1GB
+    - Sparse: 1M * 0.001 * 256 * 4 bytes = 1MB
+
+    Attributes:
+        full_dim: Dimension of the full parameter space.
+        subspace_dim: Dimension of the low-rank subspace.
+        density: Fraction of nonzeros per column (default: 1/sqrt(full_dim)).
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        full_dim: int,
+        subspace_dim: int,
+        density: Optional[float] = None,
+        seed: int = 0,
+    ):
+        """Initialize sparse random projection.
+
+        Args:
+            full_dim: Dimension of the full parameter space.
+            subspace_dim: Dimension of the low-rank subspace.
+            density: Fraction of nonzeros per column. Default: 1/sqrt(full_dim).
+            seed: Random seed for reproducibility.
+        """
+        self.full_dim = full_dim
+        self.subspace_dim = subspace_dim
+        self.density = density if density is not None else 1.0 / math.sqrt(full_dim)
+        self.seed = seed
+
+        # Warn loudly at extreme compression ratios -- see
+        # ``_EXTREME_COMPRESSION_RATIO`` above.
+        if full_dim > 0:
+            ratio = subspace_dim / full_dim
+            if ratio < _EXTREME_COMPRESSION_RATIO:
+                warnings.warn(
+                    f"SparseRandomProjection: subspace_dim={subspace_dim} / "
+                    f"full_dim={full_dim} = {ratio:.2e} is below the empirical "
+                    f"floor ({_EXTREME_COMPRESSION_RATIO:.0e}) where the "
+                    f"projection preserves distances meaningfully. The paper "
+                    f"reports GPT-2 124M with 128-dim projection collapsed to "
+                    f"random predictions in this regime. Consider increasing "
+                    f"subspace_dim.",
+                    stacklevel=2,
+                )
+
+        # Compute nnz per column (at least 1)
+        self._nnz_per_col = max(1, int(self.density * full_dim))
+
+        # Lazy initialization: indices and values created on first use
+        self._indices: Optional[torch.Tensor] = None
+        self._values: Optional[torch.Tensor] = None
+        self._device: Optional[torch.device] = None
+        self._dtype: Optional[torch.dtype] = None
+        self._sparse_matrix: Optional[torch.Tensor] = None
+
+    def _init_sparse_matrix(self, device: torch.device, dtype: torch.dtype) -> None:
+        """Initialize sparse projection matrix on demand.
+
+        Creates row indices via vectorized random sampling and Rademacher
+        values (+1/-1) scaled by 1/sqrt(nnz_per_col).
+
+        Uses ``torch.randint`` for row index generation which samples with
+        replacement. At typical sparsity (density = 1/sqrt(n)), the collision
+        rate is negligible: for n=786K and nnz_per_col=886, expected collisions
+        per column is ~0.5 (0.06% of entries). Any duplicate indices in a
+        column are summed by ``coalesce()``, which marginally reduces that
+        column's effective nonzeros but does not affect JL distance
+        preservation guarantees in practice.
+
+        Args:
+            device: Device for the sparse matrix.
+            dtype: Data type for values.
+        """
+        # Use CPU generator for consistent behavior (avoids CUDA generator issues)
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(self.seed)
+
+        total_nnz = self._nnz_per_col * self.subspace_dim
+
+        # Scale factor for unit variance
+        scale = 1.0 / math.sqrt(self._nnz_per_col)
+
+        # Vectorized row index sampling (with replacement -- negligible collisions)
+        row_indices = torch.randint(
+            0,
+            self.full_dim,
+            (total_nnz,),
+            generator=generator,
+        )
+
+        # Column indices: each column gets nnz_per_col consecutive entries
+        col_indices = (
+            torch.arange(self.subspace_dim)
+            .unsqueeze(1)
+            .expand(
+                -1,
+                self._nnz_per_col,
+            )
+            .reshape(-1)
+        )
+
+        # Rademacher values: +1 or -1 with equal probability, vectorized
+        signs = (
+            torch.randint(
+                0,
+                2,
+                (total_nnz,),
+                generator=generator,
+            )
+            * 2
+            - 1
+        )
+        values = signs.to(dtype) * scale
+
+        # Store indices as (2, nnz) for sparse_coo_tensor
+        self._indices = torch.stack([row_indices, col_indices]).to(device)
+        self._values = values.to(device)
+        self._device = device
+        self._dtype = dtype
+
+    def _get_sparse_matrix(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        """Get sparse projection matrix, initializing if needed.
+
+        Args:
+            device: Target device.
+            dtype: Target dtype.
+
+        Returns:
+            Sparse COO tensor of shape (full_dim, subspace_dim).
+        """
+        # Initialize on first use or if device/dtype changed.
+        if self._indices is None or self._device != device or self._dtype != dtype:
+            self._init_sparse_matrix(device, dtype)
+            self._sparse_matrix = None
+
+        if self._sparse_matrix is None:
+            self._sparse_matrix = torch.sparse_coo_tensor(
+                self._indices,
+                self._values,
+                size=(self.full_dim, self.subspace_dim),
+                device=device,
+                dtype=dtype,
+            ).coalesce()
+
+        return self._sparse_matrix
+
+    def project(self, coords: torch.Tensor) -> torch.Tensor:
+        """Project subspace coordinates to full parameter space.
+
+        Computes: full = P @ coords where P is (full_dim x subspace_dim).
+
+        Args:
+            coords: Subspace coordinates of shape (subspace_dim,) or (batch, subspace_dim).
+
+        Returns:
+            Full parameter space vector of shape (full_dim,) or (batch, full_dim).
+        """
+        is_1d = coords.dim() == 1
+        if is_1d:
+            coords = coords.unsqueeze(0)  # (1, subspace_dim)
+
+        # Get sparse matrix
+        P = self._get_sparse_matrix(coords.device, coords.dtype)
+
+        # Sparse-dense matmul: P @ coords.T -> (full_dim, batch)
+        # coords is (batch, subspace_dim), need (subspace_dim, batch)
+        result = torch.sparse.mm(P, coords.T)  # (full_dim, batch)
+        result = result.T  # (batch, full_dim)
+
+        if is_1d:
+            result = result.squeeze(0)
+
+        return result
+
+    def project_transpose(self, full: torch.Tensor) -> torch.Tensor:
+        """Project full parameters back to subspace (transpose).
+
+        Computes: coords = P^T @ full where P is (full_dim x subspace_dim).
+
+        Args:
+            full: Full parameter space vector of shape (full_dim,) or (batch, full_dim).
+
+        Returns:
+            Subspace coordinates of shape (subspace_dim,) or (batch, subspace_dim).
+        """
+        is_1d = full.dim() == 1
+        if is_1d:
+            full = full.unsqueeze(0)  # (1, full_dim)
+
+        # Get sparse matrix and transpose
+        P = self._get_sparse_matrix(full.device, full.dtype)
+        P_T = P.t()  # (subspace_dim, full_dim)
+
+        # Sparse-dense matmul: P^T @ full.T -> (subspace_dim, batch)
+        result = torch.sparse.mm(P_T, full.T)  # (subspace_dim, batch)
+        result = result.T  # (batch, subspace_dim)
+
+        if is_1d:
+            result = result.squeeze(0)
+
+        return result
+
+    @property
+    def dtype(self) -> Optional[torch.dtype]:
+        """Data type of the projection matrix (None if not yet initialized)."""
+        return self._dtype
+
+    @property
+    def device(self) -> Optional[torch.device]:
+        """Device of the projection matrix (None if not yet initialized)."""
+        return self._device
+
+    @property
+    def memory_bytes(self) -> int:
+        """Estimated memory usage in bytes.
+
+        Memory breakdown:
+        - Indices: 2 * nnz * 8 bytes (int64 for row and col)
+        - Values: nnz * element_size bytes (depends on dtype)
+
+        Returns:
+            Estimated memory usage in bytes.
+        """
+        total_nnz = self._nnz_per_col * self.subspace_dim
+        # 2 index arrays (row, col) with int64
+        indices_bytes = 2 * total_nnz * 8
+        element_size = self._dtype.itemsize if self._dtype else 4
+        values_bytes = total_nnz * element_size
+        return indices_bytes + values_bytes
+
+    @property
+    def nnz(self) -> int:
+        """Total number of nonzeros in projection matrix."""
+        return self._nnz_per_col * self.subspace_dim
+
+    @property
+    def actual_density(self) -> float:
+        """Actual density (nnz / total_elements)."""
+        return self.nnz / (self.full_dim * self.subspace_dim)
+
+    def __repr__(self) -> str:
+        return (
+            f"SparseRandomProjection(full_dim={self.full_dim}, "
+            f"subspace_dim={self.subspace_dim}, density={self.density:.4f}, "
+            f"nnz_per_col={self._nnz_per_col}, memory={self.memory_bytes / 1e6:.2f}MB)"
+        )

polystep/quadratic_model.py

@@ -0,0 +1,285 @@
+"""Quadratic model extraction from orthoplex cost evaluations.
+
+Extracts finite-difference gradient, Hessian diagonal, and Newton step
+from the (P, V, K) loss tensor produced by orthoplex probe evaluations.
+All functions are pure tensor operations with no side effects.
+
+Vertex ordering convention (orthoplex):
+  Vertices 0..d-1: +e_0, +e_1, ..., +e_{d-1}
+  Vertices d..2d-1: -e_0, -e_1, ..., -e_{d-1}
+  Pair for direction i: vertex i (+) and vertex i+d (-)
+"""
+
+import torch
+
+
+def extract_fd_gradient(
+    losses_3d: torch.Tensor,
+    scales: torch.Tensor,
+    probe_radius: float,
+    pdim: int,
+) -> torch.Tensor:
+    """Extract finite-difference gradient from orthoplex cost evaluations.
+
+    Central difference at each probe scale, averaged for robustness:
+        g_i = mean_k[(L(+s_k*d_i) - L(-s_k*d_i)) / (2 * s_k * r)]
+
+    Args:
+        losses_3d: Loss values of shape (P, V, K) where V = 2*pdim.
+        scales: Probe scale factors of shape (K,).
+        probe_radius: Probe distance multiplier.
+        pdim: Particle dimension (number of orthoplex directions).
+
+    Returns:
+        Gradient in rotated frame of shape (P, pdim).
+    """
+    assert losses_3d.shape[1] == 2 * pdim, (
+        f"FD gradient needs orthoplex vertices (V == 2*pdim); got V={losses_3d.shape[1]}, pdim={pdim}"
+    )
+    fwd = losses_3d[:, :pdim, :]  # (P, pdim, K) at +directions
+    bwd = losses_3d[:, pdim:, :]  # (P, pdim, K) at -directions
+
+    # Central difference at each scale
+    denom = (2.0 * scales * probe_radius).unsqueeze(0).unsqueeze(0)  # (1, 1, K)
+    grad_per_scale = (fwd - bwd) / denom.clamp(min=1e-10)  # (P, pdim, K)
+
+    return grad_per_scale.mean(dim=-1)  # (P, pdim)
+
+
+def extract_fd_hessian_diag(
+    losses_3d: torch.Tensor,
+    scales: torch.Tensor,
+    probe_radius: float,
+    pdim: int,
+) -> torch.Tensor:
+    """Extract diagonal Hessian from orthoplex cost evaluations.
+
+    Uses the symmetry property of the quadratic model:
+        L(+s) + L(-s) = 2*L(0) + H_ii * s^2
+    Regresses the symmetric sum on s^2 to estimate H_ii.
+
+    Args:
+        losses_3d: Loss values of shape (P, V, K) where V = 2*pdim.
+        scales: Probe scale factors of shape (K,).
+        probe_radius: Probe distance multiplier.
+        pdim: Particle dimension.
+
+    Returns:
+        Diagonal Hessian in rotated frame of shape (P, pdim).
+    """
+    assert losses_3d.shape[1] == 2 * pdim, (
+        f"FD Hessian needs orthoplex vertices (V == 2*pdim); got V={losses_3d.shape[1]}, pdim={pdim}"
+    )
+    fwd = losses_3d[:, :pdim, :]  # (P, pdim, K)
+    bwd = losses_3d[:, pdim:, :]  # (P, pdim, K)
+
+    # Symmetric sum: L(+s) + L(-s) = 2a + H*s^2
+    sym_sum = fwd + bwd  # (P, pdim, K)
+
+    s_sq = (scales * probe_radius) ** 2  # (K,)
+
+    # Linear regression of sym_sum on s_sq to get slope = H
+    # Center for numerical stability
+    s_sq_mean = s_sq.mean()
+    sym_mean = sym_sum.mean(dim=-1, keepdim=True)  # (P, pdim, 1)
+
+    s_centered = s_sq - s_sq_mean  # (K,)
+    y_centered = sym_sum - sym_mean  # (P, pdim, K)
+
+    # slope = sum(s_centered * y_centered) / sum(s_centered^2)
+    numerator = (s_centered.unsqueeze(0).unsqueeze(0) * y_centered).sum(dim=-1)  # (P, pdim)
+    denominator = (s_centered**2).sum().clamp(min=1e-10)
+
+    return numerator / denominator  # (P, pdim) = Hessian diagonal
+
+
+def compute_newton_step(
+    gradient: torch.Tensor,
+    hessian_diag: torch.Tensor,
+    max_step_norm: float = 10.0,
+    hessian_reg: float = 1e-4,
+) -> torch.Tensor:
+    """Compute a diagonal Newton step in the rotated frame.
+
+    Where curvature is above ``hessian_reg`` returns ``-g_i / H_i``; where it
+    is at or below ``hessian_reg`` (small-positive or nonpositive) the Hessian
+    is floored to ``hessian_reg`` (a small gradient step, never an ascent step).
+    The full step is then clipped to ``max_step_norm``.
+
+    Args:
+        gradient: FD gradient of shape ``(P, pdim)``.
+        hessian_diag: Diagonal Hessian of shape ``(P, pdim)``.
+        max_step_norm: Maximum step norm (trust region bound).
+        hessian_reg: Floor applied to curvature entries at or below it.
+
+    Returns:
+        Newton step in rotated frame of shape ``(P, pdim)``.
+    """
+    # Where curvature is positive use H_i directly; otherwise fall back to
+    # the regulariser so the direction stays a descent step.
+    H_safe = torch.where(
+        hessian_diag > hessian_reg,
+        hessian_diag,
+        torch.full_like(hessian_diag, hessian_reg),
+    )
+    delta = -gradient / H_safe  # (P, pdim)
+
+    # Clamp step norm
+    norms = torch.norm(delta, dim=-1, keepdim=True).clamp(min=1e-10)
+    scale = torch.clamp(max_step_norm / norms, max=1.0)
+    return delta * scale
+
+
+def compute_predicted_improvement(
+    gradient: torch.Tensor,
+    hessian_diag: torch.Tensor,
+    step: torch.Tensor,
+) -> torch.Tensor:
+    """Predict improvement from quadratic model: dL = g.delta + 0.5*delta.H.delta.
+
+    Args:
+        gradient: FD gradient of shape (P, pdim).
+        hessian_diag: Diagonal Hessian of shape (P, pdim).
+        step: Step vector in rotated frame of shape (P, pdim).
+
+    Returns:
+        Predicted loss change per particle of shape (P,). Negative = improvement.
+    """
+    linear = (gradient * step).sum(dim=-1)
+    quadratic = 0.5 * (hessian_diag * step**2).sum(dim=-1)
+    return linear + quadratic
+
+
+def apply_newton_refinement(
+    X_bary: torch.Tensor,
+    losses_3d: torch.Tensor,
+    scales: torch.Tensor,
+    probe_radius: float,
+    pdim: int,
+    rot_mats: torch.Tensor,
+    X_current: torch.Tensor,
+    alpha: float = 0.3,
+    max_step_norm: float = 1.0,
+    hessian_reg: float = 1e-4,
+) -> torch.Tensor:
+    """Apply post-OT Newton refinement using the quadratic model from probe evaluations.
+
+    After the OT barycentric projection gives a coarse update direction, this
+    function uses second-order information already available from the probe
+    evaluations to compute a Newton correction step, improving per-step
+    convergence without additional forward passes.
+
+    Steps:
+    1. Extract FD gradient and diagonal Hessian from losses_3d (rotated frame)
+    2. Compute Newton step in rotated frame: delta_rot = -g / (H + reg)
+    3. Transform Newton step to original space: delta_orig = rot_mats @ delta_rot
+    4. Newton-predicted minimum from the probe center: X_newton = X_current + delta_orig
+    5. Blend: X_refined = (1 - alpha) * X_bary + alpha * X_newton
+
+    Args:
+        X_bary: Post-OT barycentric position, shape (P, pdim).
+        losses_3d: Probe loss values, shape (P, V, K) where V = 2*pdim.
+        scales: Probe scale factors, shape (K,).
+        probe_radius: Probe distance multiplier.
+        pdim: Particle dimension (number of orthoplex directions).
+        rot_mats: Rotation matrices, shape (P, pdim, pdim).
+        X_current: Probe-center position the quadratic model is built around,
+            shape (P, pdim). The Newton step is measured from here, so the
+            predicted minimum is X_current + delta_orig.
+        alpha: Blending weight. 0 = pure OT, 1 = pure Newton correction.
+        max_step_norm: Maximum Newton correction norm (trust region bound).
+        hessian_reg: Regularization for near-zero Hessian entries.
+
+    Returns:
+        Refined position of shape (P, pdim).
+    """
+    # 1. Extract gradient and Hessian in rotated frame
+    gradient = extract_fd_gradient(losses_3d, scales, probe_radius, pdim)
+    hessian_diag = extract_fd_hessian_diag(losses_3d, scales, probe_radius, pdim)
+
+    # 2. Compute Newton step in rotated frame
+    delta_rot = compute_newton_step(
+        gradient,
+        hessian_diag,
+        max_step_norm=max_step_norm,
+        hessian_reg=hessian_reg,
+    )
+
+    # 3. Transform Newton step to original space: delta_orig = rot_mats @ delta_rot
+    # rot_mats: (P, pdim, pdim), delta_rot: (P, pdim)
+    delta_orig = torch.einsum("bij,bj->bi", rot_mats, delta_rot)
+
+    # 4. Newton-predicted minimum, anchored at the probe center X_current (the
+    # point the quadratic is built around). X_bary already carries the transport
+    # step, so anchoring the Newton step there would double-count that move.
+    X_newton = X_current + delta_orig
+
+    # 5. Blend the OT result with the Newton prediction
+    X_refined = (1.0 - alpha) * X_bary + alpha * X_newton
+
+    # 6. Descent gate: keep the blend only where the quadratic model predicts it
+    # is no worse than the pure-OT step, else fall back to X_bary. The model is
+    # diagonal in the rotated frame, so measure both steps there from X_current.
+    rot_mats_t = rot_mats.transpose(-1, -2)
+    ot_rot = torch.einsum("bij,bj->bi", rot_mats_t, X_bary - X_current)
+    refined_rot = torch.einsum("bij,bj->bi", rot_mats_t, X_refined - X_current)
+    pred_ot = (gradient * ot_rot + 0.5 * hessian_diag * ot_rot**2).sum(dim=-1)
+    pred_refined = (gradient * refined_rot + 0.5 * hessian_diag * refined_rot**2).sum(dim=-1)
+    accept = (pred_refined <= pred_ot).unsqueeze(-1)
+    return torch.where(accept, X_refined, X_bary)
+
+
+def update_trust_region(
+    predicted_improvement: torch.Tensor,
+    actual_improvement: torch.Tensor,
+    current_radius: float,
+    expand_threshold: float = 0.75,
+    shrink_threshold: float = 0.25,
+    expand_factor: float = 1.5,
+    shrink_factor: float = 0.5,
+    min_radius: float = 0.1,
+    max_radius: float = 3.0,
+) -> float:
+    """Update trust region multiplier based on predicted vs actual improvement.
+
+    Both predicted and actual improvement use the same sign convention:
+    negative = loss decreased (improvement). The ratio actual/predicted
+    should be positive and near 1.0 when the quadratic model is accurate.
+
+    Args:
+        predicted_improvement: Predicted loss change (P,). Negative = improvement.
+        actual_improvement: Actual loss change scalar or (1,). Negative = improvement.
+        current_radius: Current trust region multiplier.
+        expand_threshold: Ratio above which to expand.
+        shrink_threshold: Ratio below which to shrink.
+        expand_factor: Radius expansion multiplier.
+        shrink_factor: Radius shrink multiplier.
+        min_radius: Minimum allowed multiplier.
+        max_radius: Maximum allowed multiplier.
+
+    Returns:
+        Updated trust region multiplier.
+    """
+    pred = predicted_improvement.mean().item()
+    actual = actual_improvement.mean().item()
+
+    if abs(pred) < 1e-10:
+        return current_radius
+
+    ratio = actual / pred
+
+    # Clamp ratio to prevent extreme updates from noisy estimates
+    ratio = max(-2.0, min(ratio, 5.0))
+
+    # Negative ratio: actual went opposite direction from prediction - aggressive shrink
+    if ratio < 0:
+        return max(current_radius * shrink_factor * 0.5, min_radius)
+
+    # Only expand when the model predicted an improvement (pred < 0) and reality
+    # matched it. An accurate but worsening step (pred > 0, actual > 0) also gives
+    # ratio ~1 and must not grow the region.
+    if ratio > expand_threshold and pred < 0:
+        return min(current_radius * expand_factor, max_radius)
+    elif ratio < shrink_threshold:
+        return max(current_radius * shrink_factor, min_radius)
+    return current_radius