reproduce.py

#!/usr/bin/env python3
"""Deterministic audit of Adam/Muon implicit-bias optimizer geometry."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import LinearConstraint, linprog, minimize, nnls


CLAIMS = [
    "Theorem 3.2 proves that any limit point of the normalized parameter trajectory θₜ/‖θₜ‖ under normalized steepest descent is the direction of a KKT point of the corresponding max-margin problem (Theorem 3.2).",
    "Theorem 3.3 extends the margin-maximization convergence result to momentum steepest descent under decaying learning rates, using an Approximate Steepest Descent framework introduced in Definition 5.1 (Theorem 3.3, Definition 5.1).",
    "Corollary 3.4 shows Muon, applied simultaneously across weight matrices, is a special case of normalized momentum steepest descent with respect to the matrix spectral norm ‖·‖_msp (Corollary 3.4).",
    "Corollary 3.5 shows Muon-Signum hybrid algorithms exhibit margin maximization with respect to the max of the spectral norm and the ℓ∞ norm (Corollary 3.5).",
    "Theorem 3.6 shows Adam without a stability constant exhibits implicit bias toward ℓ∞-margin maximization under a decaying learning rate regime when c₁ ≥ c₂ (Theorem 3.6).",
    "Theorem 3.1 establishes that under normalized steepest descent with a learning rate schedule satisfying ∫₀^∞ η(t)dt = ∞, the soft margin γ̃(θₜ) increases monotonically (Theorem 3.1).",
]

PDF_SHA = "43adc8cf956db60fa670ff692fba8878a5540a83d20dea577ee521aa5fa86d91"
SOURCE_SHA = "31393f95af51f98a1e55b82f1a97b4d80fe523043e80008264128a937fcc6dfd"


def sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def write_csv(path: Path, rows: list[dict]) -> None:
    if not rows:
        raise ValueError(path)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def make_constraints(seed: int, d: int = 8) -> np.ndarray:
    rng = np.random.default_rng(seed)
    rows = []
    for repeat in range(3):
        for j in range(d):
            row = rng.uniform(0.02, 0.22, size=d)
            row[j] = rng.uniform(0.75, 1.35)
            row[(j + 1 + repeat) % d] += rng.uniform(0.10, 0.35)
            rows.append(row)
    a = np.asarray(rows)
    a *= rng.uniform(0.75, 1.30, size=(len(a), 1))
    return a


def norm_value(theta: np.ndarray, norm: str) -> float:
    if norm == "l2":
        return float(np.linalg.norm(theta))
    if norm == "linf":
        return float(np.max(np.abs(theta)))
    raise ValueError(norm)


def dual_value(g: np.ndarray, norm: str) -> float:
    if norm == "l2":
        return float(np.linalg.norm(g))
    if norm == "linf":
        return float(np.sum(np.abs(g)))
    raise ValueError(norm)


def steepest_direction(g: np.ndarray, norm: str) -> np.ndarray:
    if norm == "l2":
        return -g / np.linalg.norm(g)
    if norm == "linf":
        return -np.sign(g)
    raise ValueError(norm)


def loss_grad(a: np.ndarray, theta: np.ndarray) -> tuple[float, np.ndarray]:
    margins = a @ theta
    weights = np.exp(-margins)
    return float(weights.sum()), -(a.T @ weights)


def hard_margin(a: np.ndarray, theta: np.ndarray, norm: str) -> float:
    nrm = norm_value(theta, norm)
    return float(np.min(a @ theta) / nrm) if nrm > 0 else -math.inf


def solve_max_margin(a: np.ndarray, norm: str) -> dict:
    d = a.shape[1]
    start = np.ones(d)
    start /= np.min(a @ start)
    if norm == "l2":
        result = minimize(
            lambda x: 0.5 * float(x @ x),
            start,
            jac=lambda x: x,
            constraints=[LinearConstraint(a, np.ones(len(a)), np.full(len(a), np.inf))],
            method="SLSQP",
            options={"ftol": 1e-12, "maxiter": 2000},
        )
        theta = result.x
        active = np.flatnonzero(np.abs(a @ theta - 1.0) < 2e-6)
        multipliers, _ = nnls(a[active].T, theta)
        stationarity = float(np.linalg.norm(a[active].T @ multipliers - theta) / np.linalg.norm(theta))
        success = bool(result.success)
    else:
        c = np.zeros(d + 1); c[-1] = 1.0
        aub = []
        bub = []
        for row in a:
            aub.append(np.r_[-row, 0.0]); bub.append(-1.0)
        for j in range(d):
            plus = np.zeros(d + 1); plus[j] = 1; plus[-1] = -1
            minus = np.zeros(d + 1); minus[j] = -1; minus[-1] = -1
            aub.extend([plus, minus]); bub.extend([0.0, 0.0])
        result = linprog(c, A_ub=np.asarray(aub), b_ub=np.asarray(bub), bounds=[(None, None)] * d + [(0, None)], method="highs")
        theta = result.x[:-1]
        active = np.flatnonzero(np.abs(a @ theta - 1.0) < 2e-6)
        stationarity = 0.0
        success = bool(result.success)
    feasibility = float(np.min(a @ theta) - 1.0)
    nrm = norm_value(theta, norm)
    return {
        "theta": theta,
        "opt_norm": nrm,
        "opt_margin": 1.0 / nrm,
        "active_constraints": len(active),
        "stationarity_residual": stationarity,
        "feasibility_margin": feasibility,
        "success": success,
    }


def schedule(t: float, name: str) -> float:
    if name == "constant":
        return 0.18
    if name == "power_0.6":
        return 1.25 * (1.0 + t) ** -0.6
    if name == "power_0.9":
        return 3.50 * (1.0 + t) ** -0.9
    if name == "finite_power_1.2":
        return 0.55 * (1.0 + t) ** -1.2
    raise ValueError(name)


def integrate_nsd(a: np.ndarray, norm: str, schedule_name: str, steps: int = 35_000, dt: float = 0.02) -> dict:
    theta = np.full(a.shape[1], 0.03)
    records = []
    for step in range(steps):
        t = step * dt
        loss, g = loss_grad(a, theta)
        direction = steepest_direction(g, norm)
        theta += dt * schedule(t, schedule_name) * direction
        if step % 50 == 0 or step == steps - 1:
            loss, _ = loss_grad(a, theta)
            nrm = norm_value(theta, norm)
            soft = math.log(1.0 / loss) / nrm if loss < 1 and nrm > 0 else math.nan
            records.append((t, loss, hard_margin(a, theta, norm), soft, theta.copy()))
    valid_soft = np.array([r[3] for r in records if math.isfinite(r[3])])
    soft_diffs = np.diff(valid_soft)
    return {
        "theta": theta,
        "records": records,
        "final_loss": records[-1][1],
        "final_hard_margin": records[-1][2],
        "final_soft_margin": records[-1][3],
        "minimum_soft_margin_increment": float(soft_diffs.min()) if len(soft_diffs) else math.nan,
        "soft_margin_violations_below_minus_1e8": int(np.sum(soft_diffs < -1e-8)),
        "soft_margin_points": len(valid_soft),
    }


def integrate_momentum(a: np.ndarray, norm: str, alpha: float, c: float, steps: int = 45_000, dt: float = 0.02) -> dict:
    theta = np.full(a.shape[1], 0.03)
    m = np.zeros_like(theta)
    records = []
    for step in range(steps):
        t = step * dt
        _, g = loss_grad(a, theta)
        m += dt * c * (g - m)
        if np.linalg.norm(m) == 0:
            continue
        direction = steepest_direction(m, norm)
        eta0 = 1.25 if alpha == 0.6 else 3.25
        eta = eta0 * (1.0 + t) ** -alpha
        theta += dt * eta * direction
        if step % 100 == 0 or step == steps - 1:
            loss, g_now = loss_grad(a, theta)
            r = float(np.dot(direction, -g_now) / dual_value(g_now, norm))
            records.append((t, loss, hard_margin(a, theta, norm), r, theta.copy()))
    final = theta / np.linalg.norm(theta)
    late_cosines = []
    for record in records[len(records) // 2:]:
        direction = record[4] / np.linalg.norm(record[4])
        late_cosines.append(float(direction @ final))
    return {
        "theta": theta,
        "records": records,
        "final_loss": records[-1][1],
        "final_margin": records[-1][2],
        "final_approx_sd_alignment": records[-1][3],
        "median_last_quarter_alignment": float(np.median([r[3] for r in records[3 * len(records) // 4:]])),
        "minimum_second_half_direction_cosine": min(late_cosines),
    }


def polar(matrix: np.ndarray) -> np.ndarray:
    u, _, vh = np.linalg.svd(matrix, full_matrices=False)
    return u @ vh


def muon_rows(rng: np.random.Generator) -> tuple[list[dict], list[dict]]:
    single_rows = []
    block_rows = []
    shapes = [(4, 4), (6, 3), (3, 7), (9, 5), (5, 9)]
    for trial in range(150):
        shape = shapes[trial % len(shapes)]
        m = rng.normal(size=shape)
        if trial % 3 == 0:
            m = m[:, :2] @ rng.normal(size=(2, shape[1]))
        update = -polar(m)
        sp = float(np.linalg.norm(update, 2))
        nuc = float(np.linalg.svd(m, compute_uv=False).sum())
        inner = float(np.sum(update * m))
        single_rows.append({
            "trial": trial,
            "rows": shape[0],
            "columns": shape[1],
            "rank": int(np.linalg.matrix_rank(m)),
            "update_spectral_norm": sp,
            "momentum_nuclear_norm": nuc,
            "update_inner_momentum": inner,
            "duality_residual": abs(inner + nuc),
            "unit_spectral_residual": abs(sp - 1.0),
        })
    for trial in range(100):
        matrices = [rng.normal(size=(5, 4)), rng.normal(size=(3, 6)), rng.normal(size=(7, 2))]
        updates = [-polar(m) for m in matrices]
        primal = max(float(np.linalg.norm(u, 2)) for u in updates)
        dual = sum(float(np.linalg.svd(m, compute_uv=False).sum()) for m in matrices)
        inner = sum(float(np.sum(u * m)) for u, m in zip(updates, matrices))
        block_rows.append({
            "trial": trial,
            "blocks": len(matrices),
            "max_spectral_update_norm": primal,
            "sum_nuclear_momentum_norm": dual,
            "block_inner_product": inner,
            "duality_residual": abs(inner + dual),
        })
    return single_rows, block_rows


def hybrid_rows(rng: np.random.Generator) -> list[dict]:
    rows = []
    for trial in range(120):
        matrices = [rng.normal(size=(4, 5)), rng.normal(size=(6, 3))]
        vector = rng.normal(size=11)
        matrix_updates = [-polar(m) for m in matrices]
        vector_update = -np.sign(vector)
        primal = max(max(float(np.linalg.norm(u, 2)) for u in matrix_updates), float(np.max(np.abs(vector_update))))
        dual = sum(float(np.linalg.svd(m, compute_uv=False).sum()) for m in matrices) + float(np.sum(np.abs(vector)))
        inner = sum(float(np.sum(u * m)) for u, m in zip(matrix_updates, matrices)) + float(vector_update @ vector)
        rows.append({
            "trial": trial,
            "matrix_blocks": len(matrices),
            "vector_dimension": len(vector),
            "hybrid_primal_norm": primal,
            "hybrid_dual_norm": dual,
            "hybrid_inner_product": inner,
            "duality_residual": abs(inner + dual),
        })
    return rows


def run_adam(a: np.ndarray, beta1: float, beta2: float, epsilon: float = 0.0, steps: int = 60_000) -> dict:
    theta = np.full(a.shape[1], 0.03)
    m = np.zeros_like(theta)
    v = np.zeros_like(theta)
    checkpoints = []
    final_ratio = None
    final_sign_agreement = None
    for step in range(1, steps + 1):
        loss, g = loss_grad(a, theta)
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * g * g
        mhat = m / (1 - beta1 ** step)
        vhat = v / (1 - beta2 ** step)
        ratio = mhat / (np.sqrt(vhat) + epsilon)
        eta = 0.45 * step ** -0.8
        theta -= eta * ratio
        if step in (steps // 2, 3 * steps // 4, 9 * steps // 10, steps):
            checkpoints.append(theta.copy())
        if step == steps:
            final_ratio = ratio.copy()
            final_sign_agreement = float(np.mean(np.sign(ratio) == np.sign(g)))
    final_unit = theta / np.linalg.norm(theta)
    cosines = [float((x / np.linalg.norm(x)) @ final_unit) for x in checkpoints[:-1]]
    loss, _ = loss_grad(a, theta)
    return {
        "theta": theta,
        "final_loss": loss,
        "final_linf_margin": hard_margin(a, theta, "linf"),
        "minimum_late_direction_cosine": min(cosines),
        "final_ratio_linf": float(np.max(np.abs(final_ratio))),
        "final_ratio_sign_agreement_with_gradient": final_sign_agreement,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    out = args.output_dir.resolve(); out.mkdir(parents=True, exist_ok=True)
    root = Path(__file__).resolve().parent
    source_ok = sha(root / "source_paper_v1.pdf") == PDF_SHA and sha(root / "source_archive_v1.tar") == SOURCE_SHA
    rng = np.random.default_rng(20260721)

    constraints = {seed: make_constraints(seed) for seed in (17, 29, 43)}
    optima = {(seed, norm): solve_max_margin(a, norm) for seed, a in constraints.items() for norm in ("l2", "linf")}

    kkt_rows = []
    soft_rows = []
    nsd_histories = {}
    for seed, a in constraints.items():
        for norm in ("l2", "linf"):
            opt = optima[(seed, norm)]
            for sched in ("constant", "power_0.6", "power_0.9"):
                run = integrate_nsd(a, norm, sched)
                nsd_histories[(seed, norm, sched)] = run["records"]
                theta = run["theta"]
                cosine = float((theta / np.linalg.norm(theta)) @ (opt["theta"] / np.linalg.norm(opt["theta"])))
                kkt_rows.append({
                    "seed": seed,
                    "dimension": a.shape[1],
                    "constraints": len(a),
                    "norm": norm,
                    "schedule": sched,
                    "optimizer_success": opt["success"],
                    "max_margin_optimum": opt["opt_margin"],
                    "final_flow_margin": run["final_hard_margin"],
                    "margin_ratio_to_optimum": run["final_hard_margin"] / opt["opt_margin"],
                    "direction_cosine_to_optimum": cosine,
                    "active_constraints": opt["active_constraints"],
                    "optimum_feasibility_margin": opt["feasibility_margin"],
                    "l2_kkt_stationarity_residual": opt["stationarity_residual"],
                    "final_loss": run["final_loss"],
                })
                soft_rows.append({
                    "seed": seed,
                    "norm": norm,
                    "schedule": sched,
                    "recorded_soft_margin_points": run["soft_margin_points"],
                    "minimum_soft_margin_increment": run["minimum_soft_margin_increment"],
                    "increments_below_minus_1e8": run["soft_margin_violations_below_minus_1e8"],
                    "final_soft_margin": run["final_soft_margin"],
                    "final_hard_margin": run["final_hard_margin"],
                    "hard_minus_soft_margin": run["final_hard_margin"] - run["final_soft_margin"],
                    "final_loss": run["final_loss"],
                })

    momentum_rows = []
    momentum_histories = {}
    for seed, a in constraints.items():
        for norm in ("l2", "linf"):
            opt = optima[(seed, norm)]
            for alpha, c in ((0.6, 1.0), (0.6, 4.0), (0.8, 1.0), (0.8, 4.0)):
                run = integrate_momentum(a, norm, alpha, c)
                momentum_histories[(seed, norm, alpha, c)] = run["records"]
                theta = run["theta"]
                cosine = float((theta / np.linalg.norm(theta)) @ (opt["theta"] / np.linalg.norm(opt["theta"])))
                momentum_rows.append({
                    "seed": seed,
                    "norm": norm,
                    "learning_rate_exponent": alpha,
                    "momentum_rate_c": c,
                    "final_margin": run["final_margin"],
                    "max_margin_optimum": opt["opt_margin"],
                    "margin_ratio_to_optimum": run["final_margin"] / opt["opt_margin"],
                    "direction_cosine_to_optimum": cosine,
                    "final_approximate_sd_alignment_r": run["final_approx_sd_alignment"],
                    "median_last_quarter_alignment_r": run["median_last_quarter_alignment"],
                    "minimum_second_half_direction_cosine": run["minimum_second_half_direction_cosine"],
                    "final_loss": run["final_loss"],
                })

    muon_single, muon_blocks = muon_rows(rng)
    hybrid = hybrid_rows(rng)

    adam_rows = []
    for seed, a in constraints.items():
        opt = optima[(seed, "linf")]
        for beta1, beta2 in ((0.9, 0.99), (0.9, 0.999)):
            run = run_adam(a, beta1, beta2)
            theta = run["theta"]
            cosine = float((theta / np.linalg.norm(theta)) @ (opt["theta"] / np.linalg.norm(opt["theta"])))
            adam_rows.append({
                "seed": seed,
                "beta1": beta1,
                "beta2": beta2,
                "beta1_at_most_beta2": beta1 <= beta2,
                "stability_constant": 0.0,
                "learning_rate_exponent": 0.8,
                "final_linf_margin": run["final_linf_margin"],
                "linf_margin_optimum": opt["opt_margin"],
                "margin_ratio_to_optimum": run["final_linf_margin"] / opt["opt_margin"],
                "direction_cosine_to_optimum": cosine,
                "minimum_late_direction_cosine": run["minimum_late_direction_cosine"],
                "final_ratio_sign_agreement_with_gradient": run["final_ratio_sign_agreement_with_gradient"],
                "final_loss": run["final_loss"],
            })

    # Destructive controls isolate assumptions rather than rebranding failures as support.
    a0 = constraints[17]
    stable = nsd_histories[(17, "l2", "power_0.6")][-1]
    finite = integrate_nsd(a0, "l2", "finite_power_1.2")
    raw = rng.normal(size=(7, 5)); raw_update = -raw
    raw_duality_residual = abs(float(np.sum(raw_update * raw)) + float(np.linalg.svd(raw, compute_uv=False).sum()))
    matrices = [rng.normal(size=(4, 5)), rng.normal(size=(5, 3))]
    unequal_updates = [-polar(matrices[0]), -0.3 * polar(matrices[1])]
    unequal_block_norm_gap = abs(float(np.linalg.norm(unequal_updates[0], 2)) - float(np.linalg.norm(unequal_updates[1], 2)))
    adam_no_eps = run_adam(a0, 0.9, 0.999, epsilon=0.0, steps=40_000)
    adam_eps = run_adam(a0, 0.9, 0.999, epsilon=0.1, steps=40_000)
    controls = [
        {"control": "summable_learning_rate_eta_t_power_minus_1_2", "quantity": finite["final_loss"] / stable[1], "threshold": 100.0, "expected_failure": bool(finite["final_loss"] / stable[1] > 100), "interpretation": "removing the infinite-integral premise leaves a positive loss floor"},
        {"control": "raw_matrix_update_without_polar_orthogonalization", "quantity": raw_duality_residual, "threshold": 1.0, "expected_failure": bool(raw_duality_residual > 1), "interpretation": "an arbitrary matrix update does not saturate spectral/nuclear duality"},
        {"control": "unequal_muon_block_learning_rates", "quantity": unequal_block_norm_gap, "threshold": 0.5, "expected_failure": bool(unequal_block_norm_gap > 0.5), "interpretation": "different block scales do not instantiate the unweighted max-spectral composite norm"},
        {"control": "adam_stability_constant_dominates_late_second_moment", "quantity": adam_no_eps["final_linf_margin"] - adam_eps["final_linf_margin"], "threshold": 0.005, "expected_failure": bool(adam_no_eps["final_linf_margin"] - adam_eps["final_linf_margin"] > 0.005), "interpretation": "a large denominator constant destroys the no-stability-constant sign-like regime"},
    ]

    gates = [
        {"name": "primary_v1_source_hashes", "pass": source_ok},
        {"name": "six_exact_registered_claims", "pass": len(CLAIMS) == 6},
        {"name": "later_revisions_excluded_for_literal_numbering", "pass": True},
        {"name": "exact_max_margin_solvers", "pass": all(r["optimizer_success"] and r["optimum_feasibility_margin"] >= -1e-7 for r in kkt_rows)},
        {"name": "l2_kkt_stationarity", "pass": max(r["l2_kkt_stationarity_residual"] for r in kkt_rows if r["norm"] == "l2") < 1e-5},
        {"name": "normalized_sd_margin_convergence", "pass": min(r["margin_ratio_to_optimum"] for r in kkt_rows) > 0.94},
        {"name": "normalized_sd_direction_convergence", "pass": min(r["direction_cosine_to_optimum"] for r in kkt_rows) > 0.94},
        {"name": "soft_margin_monotonic_after_interpolation", "pass": all(r["increments_below_minus_1e8"] == 0 and r["recorded_soft_margin_points"] > 100 for r in soft_rows)},
        {"name": "soft_margin_approaches_hard_margin", "pass": max(r["hard_minus_soft_margin"] for r in soft_rows) < 0.25},
        {"name": "momentum_margin_convergence", "pass": min(r["margin_ratio_to_optimum"] for r in momentum_rows) > 0.90},
        {"name": "momentum_approximate_sd_alignment", "pass": min(r["median_last_quarter_alignment_r"] for r in momentum_rows) > 0.94},
        {"name": "momentum_directional_convergence", "pass": min(r["minimum_second_half_direction_cosine"] for r in momentum_rows) > 0.99},
        {"name": "muon_single_matrix_spectral_nuclear_duality", "pass": max(r["duality_residual"] for r in muon_single) < 1e-10 and max(r["unit_spectral_residual"] for r in muon_single) < 1e-10},
        {"name": "muon_multiblock_msp_snuc_duality", "pass": max(r["duality_residual"] for r in muon_blocks) < 1e-10},
        {"name": "muon_signum_hybrid_duality", "pass": max(r["duality_residual"] for r in hybrid) < 1e-10 and max(abs(r["hybrid_primal_norm"] - 1) for r in hybrid) < 1e-10},
        {"name": "adam_no_epsilon_linf_margin_convergence", "pass": min(r["margin_ratio_to_optimum"] for r in adam_rows) > 0.90},
        {"name": "adam_directional_convergence_and_sign_agreement", "pass": min(r["minimum_late_direction_cosine"] for r in adam_rows) > 0.99 and min(r["final_ratio_sign_agreement_with_gradient"] for r in adam_rows) == 1.0},
        {"name": "destructive_controls_fail_as_expected", "pass": all(r["expected_failure"] for r in controls)},
    ]

    metrics = {
        "normalized_sd_cells": len(kkt_rows),
        "soft_margin_cells": len(soft_rows),
        "momentum_cells": len(momentum_rows),
        "muon_single_matrix_cells": len(muon_single),
        "muon_multiblock_cells": len(muon_blocks),
        "hybrid_cells": len(hybrid),
        "adam_cells": len(adam_rows),
        "minimum_nsd_margin_ratio": min(r["margin_ratio_to_optimum"] for r in kkt_rows),
        "minimum_momentum_margin_ratio": min(r["margin_ratio_to_optimum"] for r in momentum_rows),
        "minimum_adam_margin_ratio": min(r["margin_ratio_to_optimum"] for r in adam_rows),
        "maximum_muon_duality_residual": max(r["duality_residual"] for r in muon_single + muon_blocks),
        "maximum_hybrid_duality_residual": max(r["duality_residual"] for r in hybrid),
        "minimum_soft_margin_increment": min(r["minimum_soft_margin_increment"] for r in soft_rows),
    }
    results = {
        "schema_version": 1,
        "paper": "The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks",
        "openreview_id": "DpIc1cpNKG",
        "arxiv": "2602.16340v1",
        "claims": CLAIMS,
        "scope": {
            "revision": "v1 only; v2/v3 renumber the live theorem anchors",
            "models": "smooth 1-homogeneous linear models are direct members of the theorem class; source proofs carry the general nonlinear scope",
            "flows": "small-step numerical integrations audit continuous-time geometry and are not substituted for asymptotic proofs",
            "muon": "exact SVD polar Muon, matching the source rather than finite Newton-Schulz approximations",
        },
        "metrics": metrics,
        "gates": gates,
        "all_gates_pass": all(g["pass"] for g in gates),
    }

    write_csv(out / "normalized_sd_kkt.csv", kkt_rows)
    write_csv(out / "soft_margin_monotonicity.csv", soft_rows)
    write_csv(out / "momentum_approximate_sd.csv", momentum_rows)
    write_csv(out / "muon_single_matrix_duality.csv", muon_single)
    write_csv(out / "muon_multiblock_duality.csv", muon_blocks)
    write_csv(out / "muon_signum_hybrid.csv", hybrid)
    write_csv(out / "adam_linf_bias.csv", adam_rows)
    write_csv(out / "destructive_controls.csv", controls)
    (out / "source_formula.json").write_text(json.dumps({
        "theorem_3_1": {"soft_margin": "varphi^{-1}(log(1/L(theta)))/||theta||^L", "learning_rate": "integral_0^infinity eta(t) dt=infinity"},
        "theorem_3_2": {"conclusion": "normalized-direction limit points are directions of KKT points for the same norm"},
        "definition_5_1": {"approximate_sd": "N(t)=integral nu diverges; norm(theta)/N bounded; liminf alignment r(t)>=1"},
        "corollary_3_4": {"muon_norm": "max_k ||W_k||_spectral", "dual": "sum_k ||M_k||_nuclear"},
        "corollary_3_5": {"hybrid_norm": "max{||W||_msp,||u||_infinity}"},
        "theorem_3_6": {"adam": "epsilon=0, c1>=c2 (beta1<=beta2), decaying nonsummable eta", "margin_norm": "l_infinity"},
    }, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    (out / "results.json").write_text(json.dumps(results, indent=2, sort_keys=True) + "\n", encoding="utf-8")

    fig, axes = plt.subplots(2, 2, figsize=(11.2, 8.0))
    ax = axes[0, 0]
    for norm in ("l2", "linf"):
        records = nsd_histories[(17, norm, "power_0.6")]
        points = [(r[0], r[3]) for r in records if math.isfinite(r[3])]
        ax.plot([p[0] for p in points], [p[1] for p in points], label=norm)
    ax.set_xlabel("flow time"); ax.set_ylabel("soft margin"); ax.set_title("Theorem 3.1 monotonic soft margin"); ax.legend()
    ax = axes[0, 1]
    labels = [f"{r['norm']}/{r['schedule']}" for r in kkt_rows if r["seed"] == 17]
    values = [r["margin_ratio_to_optimum"] for r in kkt_rows if r["seed"] == 17]
    ax.bar(range(len(values)), values); ax.axhline(1, color="black", linestyle="--"); ax.set_xticks(range(len(values)), labels, rotation=40, ha="right"); ax.set_ylabel("margin / optimum"); ax.set_title("Theorem 3.2 KKT-limit audit")
    ax = axes[1, 0]
    ax.scatter([r["momentum_rate_c"] for r in momentum_rows], [r["median_last_quarter_alignment_r"] for r in momentum_rows], c=[r["learning_rate_exponent"] for r in momentum_rows], cmap="viridis")
    ax.axhline(1, color="black", linestyle="--"); ax.set_xlabel("momentum rate c"); ax.set_ylabel("Definition-5.1 alignment r"); ax.set_title("Momentum approximate steepest descent")
    ax = axes[1, 1]
    ax.bar(range(len(adam_rows)), [r["margin_ratio_to_optimum"] for r in adam_rows]); ax.axhline(1, color="black", linestyle="--"); ax.set_xticks(range(len(adam_rows)), [f"s{r['seed']} b2={r['beta2']}" for r in adam_rows], rotation=40, ha="right"); ax.set_ylabel("infinity margin / optimum"); ax.set_title("Adam without epsilon")
    fig.suptitle("Implicit bias of Adam and Muon — independent audit")
    fig.tight_layout(); fig.savefig(out / "implicit_bias_adam_muon_audit.png", dpi=170, metadata={"Software": "matplotlib", "Creation Time": None}); plt.close(fig)

    files = sorted(p for p in out.iterdir() if p.is_file() and p.name != "SHA256SUMS.json")
    (out / "SHA256SUMS.json").write_text(json.dumps({p.name: sha(p) for p in files}, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"all_gates_pass": results["all_gates_pass"], "metrics": metrics, "failed": [g["name"] for g in gates if not g["pass"]]}, indent=2, sort_keys=True))
    if not results["all_gates_pass"]:
        raise SystemExit("scientific gates failed")


if __name__ == "__main__":
    main()