code/claim5_stability.py

"""Paper-scale CPU reproduction of the critical-depth NTK experiment."""

import json
import time

import jax

jax.config.update("jax_platform_name", "cpu")

import jax.numpy as jnp
import numpy as np


WIDTH = 200
DEPTH = 30
N_NETWORKS = 1000
N_PROBES = 4
BATCH_SIZE = 4
SEED = 250811522
INPUTS = np.array(
    [
        [-0.9895229339599609, -0.5992491841316223],
        [-0.17877478897571564, 2.253682851791382],
    ],
    dtype=np.float32,
)
REGIMES = {"low": 0.25, "critical": 2.0, "high": 4.0}


def init_params(key, width: int, depth: int, c_w: float):
    keys = jax.random.split(key, depth)
    shapes = [(width, INPUTS.shape[1])] + [(width, width)] * (depth - 1)
    return tuple(
        jnp.sqrt(jnp.float32(c_w))
        * jax.random.normal(layer_key, shape, dtype=jnp.float32)
        for layer_key, shape in zip(keys, shapes, strict=True)
    )


def forward_all_layers(params, inputs):
    z = inputs
    outputs = []
    for layer, weights in enumerate(params):
        activations = z if layer == 0 else jax.nn.relu(z)
        z = activations @ weights.T / jnp.sqrt(jnp.float32(activations.shape[-1]))
        outputs.append(z)
    return jnp.stack(outputs)


def rademacher_tangent(key, params):
    keys = jax.random.split(key, len(params))
    return tuple(
        jax.random.rademacher(tangent_key, weights.shape, dtype=jnp.float32)
        for tangent_key, weights in zip(keys, params, strict=True)
    )


def one_network(key, c_w: float, width: int, depth: int, probes: int):
    parameter_key, probe_root = jax.random.split(key)
    params = init_params(parameter_key, width, depth, c_w)
    probe_keys = jax.random.split(probe_root, probes)

    def one_probe(probe_key):
        tangent = rademacher_tangent(probe_key, params)
        _, directional = jax.jvp(
            lambda p: forward_all_layers(p, jnp.asarray(INPUTS)),
            (params,),
            (tangent,),
        )
        theta_00 = jnp.mean(directional[:, 0] ** 2, axis=-1)
        theta_01 = jnp.mean(directional[:, 0] * directional[:, 1], axis=-1)
        theta_11 = jnp.mean(directional[:, 1] ** 2, axis=-1)
        return jnp.stack((theta_00, theta_01, theta_11), axis=-1)

    return jnp.mean(jax.vmap(one_probe)(probe_keys), axis=0)


def regression(x, y):
    design = np.column_stack((np.ones_like(x), x))
    intercept, slope = np.linalg.lstsq(design, y, rcond=None)[0]
    fitted = intercept + slope * x
    ss_res = float(np.sum((y - fitted) ** 2))
    ss_tot = float(np.sum((y - np.mean(y)) ** 2))
    return {
        "intercept": float(intercept),
        "slope": float(slope),
        "r_squared": float(1.0 - ss_res / ss_tot) if ss_tot else 1.0,
    }


def summarize(samples):
    n = samples.shape[0]
    mean = samples.mean(axis=0)
    standard_deviation = samples.std(axis=0, ddof=1)
    standard_error = standard_deviation / np.sqrt(n)
    layers = np.arange(1, samples.shape[1] + 1, dtype=np.float64)
    slopes = np.einsum("l,nlc->nc", layers, samples) / np.dot(layers, layers)
    return {
        "mean": mean.tolist(),
        "standard_deviation": standard_deviation.tolist(),
        "standard_error": standard_error.tolist(),
        "per_network_through_origin_slopes": slopes.tolist(),
        "through_origin_slope_mean": slopes.mean(axis=0).tolist(),
        "through_origin_slope_standard_error": (
            slopes.std(axis=0, ddof=1) / np.sqrt(n)
        ).tolist(),
    }


def critical_checks(summary, low_summary, high_summary):
    layers = np.arange(1, DEPTH + 1, dtype=np.float64)
    observed = np.asarray(summary["mean"], dtype=np.float64)
    standard_error = np.asarray(summary["standard_error"], dtype=np.float64)
    expected_slope = float(np.dot(INPUTS[0], INPUTS[0]) / INPUTS.shape[1])
    expected = expected_slope * layers
    residual = observed[:, 0] - expected
    safe_se = np.maximum(standard_error[:, 0], np.finfo(np.float64).eps)
    slope_through_origin = float(summary["through_origin_slope_mean"][0])
    slope_se = float(summary["through_origin_slope_standard_error"][0])
    slope_z = abs(slope_through_origin - expected_slope) / max(slope_se, np.finfo(float).eps)
    critical_diag_regression = regression(layers, observed[:, 0])
    critical_offdiag_regression = regression(layers[9:], observed[9:, 1])

    high = np.maximum(np.asarray(high_summary["mean"], dtype=np.float64)[:, 0], 1e-30)
    high_log_regression = regression(layers[9:], np.log(high[9:]))
    low = np.maximum(np.asarray(low_summary["mean"], dtype=np.float64)[:, 0], 1e-30)
    low_log_regression = regression(layers[9:], np.log(low[9:]))
    high_expected_residual = high - expected
    high_se = np.maximum(
        np.asarray(high_summary["standard_error"], dtype=np.float64)[:, 0],
        np.finfo(np.float64).eps,
    )
    high_falsely_critical = bool(
        abs(float(np.dot(layers, high, ) / np.dot(layers, layers)) - expected_slope)
        / max(float(np.sqrt(np.sum((layers * high_se) ** 2)) / np.dot(layers, layers)), np.finfo(float).eps)
        <= 3.0
        and np.max(np.abs(high_expected_residual) / np.maximum(expected, 1e-30)) <= 0.10
    )

    checks = {
        "exact_paper_scale": WIDTH == 200 and DEPTH == 30 and N_NETWORKS == 1000,
        "critical_slope_within_3_standard_errors": slope_z <= 3.0,
        "critical_max_relative_deviation_at_most_10_percent": (
            float(np.max(np.abs(residual) / expected)) <= 0.10
        ),
        "critical_at_least_27_of_30_points_within_99pct_pointwise_ci": (
            int(np.sum(np.abs(residual) <= 2.576 * safe_se)) >= 27
        ),
        "critical_diagonal_linear_r_squared_at_least_0_995": (
            critical_diag_regression["r_squared"] >= 0.995
        ),
        "critical_offdiagonal_asymptotic_linear_r_squared_at_least_0_98": (
            critical_offdiag_regression["r_squared"] >= 0.98
        ),
        "high_variance_control_has_positive_exponential_log_slope": (
            high_log_regression["slope"] >= 0.20
        ),
        "low_variance_control_has_negative_exponential_log_slope": (
            low_log_regression["slope"] <= -0.20
        ),
        "high_variance_curve_rejected_by_critical_contract": not high_falsely_critical,
    }
    return {
        "expected_critical_diagonal_slope": expected_slope,
        "observed_critical_diagonal_slope_through_origin": slope_through_origin,
        "observed_slope_standard_error": slope_se,
        "observed_slope_z": float(slope_z),
        "critical_max_relative_deviation": float(np.max(np.abs(residual) / expected)),
        "critical_points_within_99pct_pointwise_ci": int(
            np.sum(np.abs(residual) <= 2.576 * safe_se)
        ),
        "critical_diagonal_regression": critical_diag_regression,
        "critical_offdiagonal_regression_depths_10_to_30": critical_offdiag_regression,
        "high_diagonal_log_regression_depths_10_to_30": high_log_regression,
        "low_diagonal_log_regression_depths_10_to_30": low_log_regression,
        "negative_control_high_falsely_accepted_as_critical": high_falsely_critical,
        "checks": checks,
        "passed": all(checks.values()),
    }


def run_paper_scale():
    started = time.perf_counter()
    regime_summaries = {}
    regime_seeds = {}
    for regime_index, (name, c_w) in enumerate(REGIMES.items()):
        seed = SEED + 100_000 * regime_index
        regime_seeds[name] = seed
        keys = jax.random.split(jax.random.PRNGKey(seed), N_NETWORKS)
        batched = jax.jit(
            jax.vmap(
                lambda network_key: one_network(
                    network_key, c_w, WIDTH, DEPTH, N_PROBES
                )
            )
        )
        batches = []
        regime_started = time.perf_counter()
        for lower in range(0, N_NETWORKS, BATCH_SIZE):
            upper = min(lower + BATCH_SIZE, N_NETWORKS)
            batches.append(np.asarray(batched(keys[lower:upper])))
            if upper % 100 == 0:
                print(
                    f"CLAIM5_PROGRESS regime={name} networks={upper}/{N_NETWORKS} "
                    f"seconds={time.perf_counter() - regime_started:.1f}",
                    flush=True,
                )
        samples = np.concatenate(batches, axis=0).astype(np.float64)
        regime_summaries[name] = summarize(samples)
        regime_summaries[name]["runtime_seconds"] = time.perf_counter() - regime_started

    verification = critical_checks(
        regime_summaries["critical"],
        regime_summaries["low"],
        regime_summaries["high"],
    )
    return {
        "claim": (
            "At C_W=2, a bias-free width-200 ReLU MLP has linearly scaling "
            "mean NTK through depth 30; away from C_W=2 it is exponentially unstable"
        ),
        "estimator": (
            "Appendix C output-channel trace average with unbiased Rademacher "
            "Hutchinson parameter-space probes"
        ),
        "inputs": INPUTS.tolist(),
        "width": WIDTH,
        "depth": DEPTH,
        "network_initializations_per_regime": N_NETWORKS,
        "hutchinson_probes_per_network": N_PROBES,
        "regimes": REGIMES,
        "seeds": regime_seeds,
        "columns": ["theta_00", "theta_01", "theta_11"],
        "layers": list(range(1, DEPTH + 1)),
        "summaries": regime_summaries,
        "verification": verification,
        "runtime_seconds": time.perf_counter() - started,
        "passed": verification["passed"],
    }


def main():
    result = run_paper_scale()
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0 if result["passed"] else 1


if __name__ == "__main__":
    raise SystemExit(main())