code/audit_piecewise_claim4.py

#!/usr/bin/env python3
"""CPU audit of the CAffNet piecewise-constraint claim.

The paper specifies the benchmark functions and the reported Table 2 means,
but not enough transformer implementation detail to call a new training run a
faithful reproduction.  This program therefore keeps two outputs separate:

* an exact Decimal audit of the published rounded means and the claimed
  73.33% reduction; and
* an independent, explicitly labelled one-token, three-head transformer
  experiment using the published functions and training regime.

No result is inferred from the source table: the script prints the source
arithmetic and the executed CPU experiment independently.
"""

from __future__ import annotations

import argparse
import json
import math
import random
from decimal import Decimal, getcontext
from typing import Callable

import numpy as np
import torch
from torch import nn


torch.set_num_threads(1)
torch.set_num_interop_threads(1)
getcontext().prec = 40


def pw(x: np.ndarray | torch.Tensor, *, kind: str):
    """Published Appendix D.1 functions, evaluated by the four branches."""
    if isinstance(x, torch.Tensor):
        sin = torch.sin
        pi = torch.pi
        where = torch.where
        z = torch.zeros_like(x)
    else:
        sin = np.sin
        pi = np.pi
        where = np.where
        z = np.zeros_like(x)

    a = x <= -1
    b = (x > -1) & (x <= 0)
    c = (x > 0) & (x <= 1)

    if kind == "f":
        vals = (-5 * sin(pi / 2 * (x + 1)) - 2,
                -2 + z,
                2 - 9 * (x - 2 / 3) ** 2,
                3 / x ** 2 - 2)
    elif kind == "u1":
        vals = (-3 * sin(pi / 2 * (x + 1)) + 1 / 5,
                -2 + z,
                3 - 4 * (x - 1 / 2) ** 2,
                2 + z)
    elif kind == "u2":
        vals = (-3 * sin(pi / 2 * (x + 1)) ** 3 + 1,
                2 + z,
                3 - 4 * (x - 4 / 5) ** 2,
                5 / 2 + z)
    elif kind == "l1":
        vals = (5 * sin(pi / 2 * (x + 1)) ** 2 - 3,
                -2 + z,
                (4 - 9 * (x - 2 / 3) ** 2) * x - 5 / 2,
                3 / x ** 3 - 5 / 2)
    elif kind == "l2":
        vals = (5 * sin(pi / 2 * (x + 1)) ** 8 - 2,
                -3 + z,
                (5 - 4 * (x - 1 / 6) ** 2) * x - 5 / 2,
                3 / (2 * x ** 3) - 16 / 9)
    else:
        raise ValueError(kind)
    return where(a, vals[0], where(b, vals[1], where(c, vals[2], vals[3])))


def source_table_audit() -> dict[str, object]:
    nn_mse = Decimal("0.0045")
    tf_mse = Decimal("0.0012")
    reduction = (Decimal(1) - tf_mse / nn_mse) * Decimal(100)
    # Rounded four-decimal means admit an interval; report it rather than
    # silently treating the displayed values as hidden unrounded results.
    nn_lo, nn_hi = Decimal("0.00445"), Decimal("0.00455")
    tf_lo, tf_hi = Decimal("0.00115"), Decimal("0.00125")
    lo = (Decimal(1) - tf_hi / nn_lo) * Decimal(100)
    hi = (Decimal(1) - tf_lo / nn_hi) * Decimal(100)
    return {
        "published_nn_mse": str(nn_mse),
        "published_tf_mse": str(tf_mse),
        "reduction_from_displayed_means_percent": str(reduction),
        "reduction_interval_from_four_decimal_rounding_percent": [str(lo), str(hi)],
        "displayed_zero_violation_is_not_proof_of_exact_zero": True,
    }


def domain_audit() -> dict[str, float]:
    x = np.linspace(-2.0, 2.0, 400_001, dtype=np.float64)
    target = pw(x, kind="f")
    upper = np.minimum(pw(x, kind="u1"), pw(x, kind="u2"))
    lower = np.maximum(pw(x, kind="l1"), pw(x, kind="l2"))
    violations = np.maximum(target - upper, 0) + np.maximum(lower - target, 0)
    return {
        "grid_points": float(x.size),
        "target_max_constraint_residual": float(np.max(violations)),
        "target_min_feasible_width": float(np.min(upper - lower)),
        "target_max_feasible_width": float(np.max(upper - lower)),
    }


class OneTokenTransformer(nn.Module):
    """Small explicit interpretation of 3 heads of size 40, width 120."""

    def __init__(self):
        super().__init__()
        self.embed = nn.Linear(1, 120)
        layer = nn.TransformerEncoderLayer(
            d_model=120,
            nhead=3,
            dim_feedforward=120,
            dropout=0.0,
            activation="gelu",
            batch_first=True,
            norm_first=False,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=1, enable_nested_tensor=False)
        self.out = nn.Linear(120, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.embed(x[:, None, None])
        return self.out(self.encoder(h))[:, 0, 0]


def train(seed: int, epochs: int) -> dict[str, float | int]:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    dtype = torch.float64
    x_train = torch.from_numpy(np.random.default_rng(seed).uniform(-2, 2, 50)).to(dtype)
    y_train = pw(x_train, kind="f")
    x_test = torch.linspace(-2, 2, 400, dtype=dtype)
    y_test = pw(x_test, kind="f")
    upper_train = torch.minimum(pw(x_train, kind="u1"), pw(x_train, kind="u2"))
    lower_train = torch.maximum(pw(x_train, kind="l1"), pw(x_train, kind="l2"))
    upper_test = torch.minimum(pw(x_test, kind="u1"), pw(x_test, kind="u2"))
    lower_test = torch.maximum(pw(x_test, kind="l1"), pw(x_test, kind="l2"))

    def fit(project: bool) -> tuple[float, float]:
        model = OneTokenTransformer().to(dtype)
        opt = torch.optim.Adam(model.parameters(), lr=1e-4)
        for _ in range(epochs):
            raw = model(x_train)
            if project:
                pred = torch.maximum(torch.minimum(raw, upper_train), lower_train)
                loss = torch.mean((pred - y_train) ** 2)
            else:
                residual = torch.stack((raw - upper_train, raw - upper_train, lower_train - raw, lower_train - raw), dim=1)
                # The paper specifies the 2-norm in Eq. (6), not a sum of
                # row penalties.  Keep the four rows explicit so the exact
                # source constraint convention is visible in the run.
                penalty = torch.linalg.vector_norm(torch.relu(residual), ord=2, dim=1)
                loss = torch.mean((raw - y_train) ** 2 + 100 * penalty)
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
        with torch.no_grad():
            raw = model(x_test)
            pred = torch.maximum(torch.minimum(raw, upper_test), lower_test) if project else raw
            residual = torch.stack((pred - upper_test, pred - upper_test, lower_test - pred, lower_test - pred), dim=1)
            violation = torch.relu(residual)
            mse = torch.mean((pred - y_test) ** 2).item()
            vmax = torch.max(violation).item()
        return mse, vmax

    nn_mse, nn_vmax = fit(project=False)
    tf_mse, tf_vmax = fit(project=True)
    return {
        "seed": seed,
        "epochs": epochs,
        "nn_mse": nn_mse,
        "nn_max_violation": nn_vmax,
        "caffnet_tf_mse": tf_mse,
        "caffnet_tf_max_violation": tf_vmax,
        "mse_reduction_percent": 100 * (nn_mse - tf_mse) / nn_mse,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=5_000)
    parser.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2, 3, 4])
    args = parser.parse_args()
    result: dict[str, object] = {
        "source_table_audit": source_table_audit(),
        "domain_audit": domain_audit(),
        "training": [train(seed, args.epochs) for seed in args.seeds],
    }
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()