repro-bundle/v1/scripts/ste_repro.py

"""Core implementation for the STE-quantization reproduction.

Re-implements the quantizer, the macroscopic ODE (Theorem V.3 of
Ichikawa et al., arXiv:2510.10693), and a batched STE simulation of the
underlying linear-quantized regression problem.

Reference paper (verbatim from arXiv abstract):
    "Quantized neural network training optimizes a discrete, non-differentiable
    objective. The straight-through estimator (STE) enables backpropagation
    through surrogate gradients... We theoretically show that in the
    high-dimensional limit, STE dynamics converge to a deterministic ordinary
    differential equation."

Notation follows the paper:
  rho      := ||w*||^2 / d                            (per-component 2nd moment)
  sigma2   := noise variance  (Var[xi])
  (b, omega, L, Delta, theta, v) describe the uniform quantizer psi.
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
from scipy.stats import norm

# ---------------------------------------------------------------------------
# 1. Quantizer definition
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Quantizer:
    """Uniform mid-tread quantizer as defined in Eq. (1) of the paper.

    Levels  v_k = -omega + k*Delta,  k = 0..L        (L + 1 = 2^b - 1 levels)
    Step    Delta = 2*omega / L
    Thresholds  theta_k = -omega + (k - 1/2) * Delta, k = 1..L
    """

    b: int
    omega: float

    def __post_init__(self):
        assert self.b >= 2, "bit width b>=2 required (paper Eq. defines L=2^b-2)"
        L = (1 << self.b) - 2  # L + 1 = 2^b - 1
        Delta = 2.0 * self.omega / L
        levels = -self.omega + np.arange(L + 1) * Delta
        # thresholds theta_1..theta_L ;  (theta_0 = -inf, theta_{L+1} = +inf)
        theta = -self.omega + (np.arange(1, L + 1) - 0.5) * Delta
        # expose as object attrs bypassing frozen
        object.__setattr__(self, "L", L)
        object.__setattr__(self, "Delta", Delta)
        object.__setattr__(self, "levels", levels)
        object.__setattr__(self, "theta", theta)

    # --- forward maps ---------------------------------------------------
    def psi(self, x: np.ndarray) -> np.ndarray:
        """Hard quantizer psi(x), Eq. (1)."""
        # psi(x) = -omega + Delta * sum_{k=1..L} Heaviside(x - theta_k)
        # = -omega + Delta * (number of thresholds <= x)
        out = np.full_like(x, -self.omega, dtype=float)
        for k in range(self.L):
            out = out + self.Delta * (x >= self.theta[k]).astype(float)
        return out

    def psi_t(self, x: np.ndarray, T: float) -> np.ndarray:
        """Differentiable relaxation psi_T(x), Eq. after (1).

            psi_T(x) = -omega + Delta * sum_{k=1..L} Phi((x - theta_k)/T)

        As T -> +0, psi_T -> psi (hard quantizer).
        For T = 0, falls back to psi(x).
        """
        if T <= 0.0:
            return self.psi(x)
        out = np.full_like(x, -self.omega, dtype=float)
        for k in range(self.L):
            out = out + self.Delta * norm.cdf((x - self.theta[k]) / T)
        return out

    # --- moments --------------------------------------------------------
    def sigma2_psi(self) -> float:
        """E[psi(X)^2], X ~ N(0, 1).  Prop. II.1."""
        L = self.L
        v = self.levels
        th = self.theta
        # boundaries: theta_0 = -inf, theta_{L+1} = +inf
        # p_k = Phi(theta_{k+1}) - Phi(theta_k),  k = 0..L
        cdf = lambda x: 0.0 if x == -np.inf else (1.0 if x == np.inf else norm.cdf(x))
        # use list of bounds
        bounds = [-np.inf] + list(th) + [np.inf]
        s = 0.0
        for k in range(L + 1):
            p_k = cdf(bounds[k + 1]) - cdf(bounds[k])
            s += v[k] ** 2 * p_k
        return float(s)

    def kappa_psi(self) -> float:
        """E[X * psi(X)], X ~ N(0, 1).  Prop. II.1."""
        L = self.L
        v = self.levels
        th = self.theta
        # kappa = sum_{k=1..L} (v_k - v_{k-1}) * phi(theta_k)
        s = 0.0
        for k in range(L):
            s += (v[k + 1] - v[k]) * norm.pdf(th[k])
        return float(s)

    def describe(self) -> dict:
        return {
            "b": self.b,
            "omega": self.omega,
            "L": self.L,
            "Delta": self.Delta,
            "levels": self.levels.tolist(),
            "theta": self.theta.tolist(),
            "sigma2_psi": self.sigma2_psi(),
            "kappa_psi": self.kappa_psi(),
        }


def identity_quantizer_moments():
    """For the unquantized baseline: psi(x) = x gives kappa=1, sigma2=1."""
    return 1.0, 1.0


# ---------------------------------------------------------------------------
# 2. Macroscopic state functions (Prop. V.2)
# ---------------------------------------------------------------------------


def macro_m_psi(m: float, s: float, qw: Quantizer, rho: float = 1.0) -> float:
    """m_psi(m, s) for w* with E[(w*)^2] = rho.  For w* = 1_d (rho=1) the
    expectation E_{w*} is trivial: f(1)."""
    if rho != 1.0:
        raise NotImplementedError("Only w* = 1_d (rho=1) is implemented.")
    # m_psi = -omega + Delta * sum_{i=1..L} Phi((m/rho - theta_i)/s)
    if s <= 0:
        # degenerate: psi(m) for a constant value m
        return float(qw.psi(np.array([m]))[0])
    z = (m - qw.theta) / s
    return float(-qw.omega + qw.Delta * np.sum(norm.cdf(z)))


def macro_q_psi(m: float, s: float, qw: Quantizer, rho: float = 1.0) -> float:
    """q_psi(m, s)."""
    if rho != 1.0:
        raise NotImplementedError("Only w* = 1_d (rho=1) is implemented.")
    if s <= 0:
        v = float(qw.psi(np.array([m]))[0])
        return v * v
    z = (m - qw.theta) / s
    v = qw.levels
    # q_psi = v_0^2 + sum_{i=1..L} (v_i^2 - v_{i-1}^2) Phi((m - theta_i)/s)
    diffs = np.diff(v ** 2)  # length L
    return float(v[0] ** 2 + np.sum(diffs * norm.cdf(z)))


def macro_r_psi(m: float, s: float, qw: Quantizer, rho: float = 1.0) -> float:
    """r_psi(m, s) = (m/rho) * m_psi + Delta * s * sum_i phi((m/rho - theta_i)/s)."""
    if rho != 1.0:
        raise NotImplementedError("Only w* = 1_d (rho=1) is implemented.")
    m_psi = macro_m_psi(m, s, qw, rho)
    if s <= 0:
        return float(m * m_psi)
    z = (m - qw.theta) / s
    return float((m / rho) * m_psi + qw.Delta * s * np.sum(norm.pdf(z)))


def eps_g(m: float, s: float, qw: Quantizer, kappa_x: float, sigma2_x: float,
          rho: float = 1.0, sigma2: float = 0.0) -> float:
    """Generalization error eps_g = sigma2 + rho + sigma2_x * q_psi - 2 * kappa_x * m_psi."""
    return float(
        sigma2 + rho
        + sigma2_x * macro_q_psi(m, s, qw, rho)
        - 2.0 * kappa_x * macro_m_psi(m, s, qw, rho)
    )


# ---------------------------------------------------------------------------
# 3. Macroscopic ODE  (Theorem V.3)
# ---------------------------------------------------------------------------


def ode_rhs(tau, y, qw: Quantizer, kappa_x: float, sigma2_x: float,
            eta: float, lam: float, rho: float = 1.0, sigma2: float = 0.0):
    """Right-hand side of the macroscopic ODE system.

    State y = (m, q).  s = sqrt(q - m^2 / rho).
    """
    m, q = y
    # numerical guard: q >= m^2 / rho always (Cauchy-Schwarz); clip to keep s real.
    s2 = q - m * m / rho
    if s2 < 1e-12:
        s = 1e-6
    else:
        s = math.sqrt(s2)
    m_p = macro_m_psi(m, s, qw, rho)
    r_p = macro_r_psi(m, s, qw, rho)
    e = eps_g(m, s, qw, kappa_x, sigma2_x, rho, sigma2)
    dmd = -eta * ((sigma2_x + lam) * m_p - kappa_x * rho)
    dqd = -2.0 * eta * ((sigma2_x + lam) * r_p - kappa_x * m) + eta * eta * sigma2_x * e
    return [dmd, dqd]


# ---------------------------------------------------------------------------
# 4. Input-only quantization: closed form fixed point (Prop. VI.1)
# ---------------------------------------------------------------------------


def input_only_fixed_point(kappa_x: float, sigma2_x: float, eta: float,
                           lam: float, rho: float = 1.0, sigma2: float = 0.0):
    """Closed form (m*, q*) for input-only quantization (Prop. VI.1, Supp. V-A).

        m* = rho * kappa_x / (sigma2_x + lambda)
        q* = (2 kappa_x^2 + eta sigma2_x ((rho+sigma2)(sigma2_x+lambda) - 2 kappa_x^2))
             / ((sigma2_x+lambda) (2(sigma2_x+lambda) - eta sigma2_x^2))
    """
    m_star = rho * kappa_x / (sigma2_x + lam)
    num = (2.0 * kappa_x ** 2
           + eta * sigma2_x * ((rho + sigma2) * (sigma2_x + lam) - 2.0 * kappa_x ** 2))
    den = (sigma2_x + lam) * (2.0 * (sigma2_x + lam) - eta * sigma2_x ** 2)
    q_star = num / den
    eps_star = rho + sigma2 + sigma2_x * q_star - 2.0 * kappa_x * m_star
    return float(m_star), float(q_star), float(eps_star)


def input_only_stability_bound(sigma2_x: float, lam: float) -> float:
    """Stability boundary: 0 < eta < 2 (sigma2_x + lambda) / sigma2_x^2."""
    return 2.0 * (sigma2_x + lam) / (sigma2_x ** 2)


# ---------------------------------------------------------------------------
# 5. Small-eta fixed-point prediction for joint weight+input quantization
#    (Prop. V.6, Theorem V.8)
# ---------------------------------------------------------------------------


def small_eta_fixed_point_prediction(qw: Quantizer, qx: Quantizer | None,
                                     lam: float, rho: float = 1.0,
                                     sigma2: float = 0.0):
    """Closed-form small-eta prediction of eps_g*.

    If qx is None, kappa_x = sigma2_x = 1 (unquantized inputs).
    Else uses input quantizer moments.
    """
    if qx is None:
        kappa_x, sigma2_x = 1.0, 1.0
    else:
        kappa_x = qx.kappa_psi()
        sigma2_x = qx.sigma2_psi()

    c = kappa_x * rho / (sigma2_x + lam)  # Definition VI.2
    # locate index i* in {0..L-1} such that v_{i*} <= c <= v_{i*+1}
    levels = qw.levels
    if c <= -qw.omega or c >= qw.omega:
        # |c| >= omega case
        eps_star = rho + sigma2 - 2.0 * kappa_x * qw.omega + sigma2_x * qw.omega ** 2
        regime = "boundary"
        return {
            "c": c,
            "regime": regime,
            "m_star_pred": float(np.sign(c) * qw.omega),
            "eps_star_pred": float(eps_star),
            "delta2_p1p": 0.0,
            "p": None,
            "i_star": None,
        }
    # find i*
    i_star = int(np.searchsorted(levels, c, side="right") - 1)
    i_star = max(0, min(qw.L - 1, i_star))
    p = (c - levels[i_star]) / qw.Delta
    if p <= 0.0 or p >= 1.0:
        # p in {0, 1}: logarithmic regime (rare boundary case)
        # eps_g* = eps_g0 + o(1/sqrt(log(1/eta)))  -- just return eps_g0
        eps0 = (rho + sigma2 - 2.0 * kappa_x * c
                + sigma2_x * qw.omega ** 2)  # NOTE: actually c^2 if c is on a level
        # when p in {0,1}, c == v_{i*} exactly; treat as c^2:
        eps0 = rho + sigma2 - 2.0 * kappa_x * c + sigma2_x * (c ** 2)
        return {
            "c": c,
            "regime": "log",
            "m_star_pred": float(c),
            "eps_star_pred": float(eps0),
            "delta2_p1p": 0.0,
            "p": float(p),
            "i_star": int(i_star),
        }
    # p in (0, 1): eps_g* = eps_g^(0) + sigma2_x * Delta^2 * p * (1-p) + o(eta)
    eps0 = rho + sigma2 - 2.0 * kappa_x * c + sigma2_x * (c ** 2)
    correction = sigma2_x * qw.Delta ** 2 * p * (1.0 - p)
    return {
        "c": c,
        "regime": "interior",
        "m_star_pred": float(c),
        "eps_star_pred": float(eps0 + correction),
        "eps_g0_pred": float(eps0),
        "delta2_p1p": float(correction),
        "p": float(p),
        "i_star": int(i_star),
    }