#!/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()