Update logbook: Repro - Who Said Neural Networks Aren't Linear?

logbook.json

@@ -10,7 +10,7 @@
     "icml2026-repro",
     "paper-Eqm5VEaAom"
   ],
-  "updated_at": "2026-07-16T05:04:20+00:00",
+  "updated_at": "2026-07-16T06:59:55+00:00",
   "root": {
     "slug": "index",
     "title": "Repro - Who Said Neural Networks Aren't Linear?",
@@ -60,6 +60,6 @@
       }
     ]
   },
-  "agent_view_tokens": 6650,
-  "revision": "1784178260139660995"
+  "agent_view_tokens": 6941,
+  "revision": "1784185195171135424"
 }

pages/claim-3-idempotence-iff-a-2-a-lemma-4-2/page.md

@@ -278,3 +278,850 @@ ALL CHECKS PASSED
 ## Claim 3 — REPRODUCED in both directions (exact); architectural enforcement confirmed
 
 Lemma 4.2 (f idempotent iff A^2 = A) verified numerically: (i) a random symmetric projection P = Q diag(1,1,1,1,0,0,0,0) Q^T gives f(f(x)) = f(x) to 9e-16; (ii) the paper's Section-4.3 architectural parameterization — a binary diagonal core, which is what the straight-through-estimator training produces — is exactly idempotent by construction; (iii) the converse holds: a random non-idempotent A yields f(f(x)) != f(x) with gap ~47, i.e. idempotence genuinely requires A^2 = A, it is not an artifact of the construction. This confirms the paper's key point that idempotency can be enforced architecturally (choose an idempotent core) rather than via an optimization penalty as in IGN (Shocher et al. 2024). Note: the paper's own ign/ experiment code is an empty stub in the official repo (https://github.com/assafshocher/Linearizer), so the MNIST generative-projection experiment is not reproducible from released code — the lemma itself, and the mechanism that the experiment relies on, are verified exactly here.
+
+
+---
+<!-- trackio-cell
+{"type": "code", "id": "cell_6cc95a8b5e24", "created_at": "2026-07-16T06:59:41+00:00", "title": "HF GPU Job: train real-MNIST global image projector", "command": ["/home/ubuntu/samuel/repro/.venv/bin/hf", "jobs", "uv", "run", "--flavor", "l4x1", "--timeout", "30m", "--detach", "--label", "paper=Eqm5VEaAom", "--label", "claim=idempotent-projector", "--volume", "hf://buckets/SabaPivot/repro-linearizer-artifacts:/mnt/artifacts:rw", "--with", "torch==2.5.1", "--with", "torchvision==0.20.1", "--with", "matplotlib==3.10.3", "--with", "numpy==1.26.4", "claim3_image_projector.py", "--mode", "full", "--device", "cuda", "--output-dir", "/mnt/artifacts/experiments/claim3_image_projector_hf_job", "--data-dir", "/tmp/mnist", "--log-every", "100"], "exit_code": 1, "duration_s": 10.778}
+-->
+````bash
+$ /home/ubuntu/samuel/repro/.venv/bin/hf jobs uv run --flavor l4x1 --timeout 30m --detach --label paper=Eqm5VEaAom --label claim=idempotent-projector --volume hf://buckets/SabaPivot/repro-linearizer-artifacts:/mnt/artifacts:rw --with torch==2.5.1 --with torchvision==0.20.1 --with matplotlib==3.10.3 --with numpy==1.26.4 claim3_image_projector.py --mode full --device cuda --output-dir /mnt/artifacts/experiments/claim3_image_projector_hf_job --data-dir /tmp/mnist --log-every 100
+````
+
+exit 1 · 10.8s
+
+
+````python title=claim3_image_projector.py
+#!/usr/bin/env python3
+"""Real-data check of the Linearizer idempotent-network claim.
+
+This script trains an exact, invertible additive-coupling map ``g`` around a
+*fixed* diagonal projector ``P`` and evaluates
+
+    f(x) = g^{-1}(P g(x)).
+
+Because ``P**2 == P`` and ``g`` is bijective by construction, ``f`` is
+idempotent architecturally; no idempotence term appears in the training loss.
+The experiment is deliberately restorative rather than unconditional: it
+learns to project Gaussian-corrupted MNIST images toward clean MNIST images.
+
+The full preset uses 10,000 training digits, 1,000 untouched test digits, a
+rank-588 core, and a 454,512-parameter bijection.  The smoke preset exercises
+the complete train/evaluate/artifact path on a small subset.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import csv
+import hashlib
+import json
+import math
+import os
+import platform
+import random
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable
+
+# Required by PyTorch for deterministic CUDA GEMMs.  This must be set before
+# importing torch.
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Subset
+from torchvision import datasets, transforms
+
+INPUT_DIM = 28 * 28
+
+
+@dataclass(frozen=True)
+class Preset:
+    train_samples: int
+    test_samples: int
+    ood_samples: int
+    epochs: int
+    batch_size: int
+
+
+PRESETS = {
+    "smoke": Preset(
+        train_samples=512,
+        test_samples=128,
+        ood_samples=128,
+        epochs=1,
+        batch_size=128,
+    ),
+    "full": Preset(
+        train_samples=10_000,
+        test_samples=1_000,
+        ood_samples=1_000,
+        epochs=50,
+        batch_size=256,
+    ),
+}
+
+
+class Conditioner(nn.Module):
+    """Small translation network used by an additive coupling layer."""
+
+    def __init__(self, half_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(half_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, half_dim),
+        )
+
+        # Start g close to the identity.  The first layer is still randomly
+        # initialized, so symmetry breaks as soon as the final layer learns.
+        final = self.net[-1]
+        assert isinstance(final, nn.Linear)
+        nn.init.zeros_(final.weight)
+        nn.init.zeros_(final.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x)
+
+
+class AdditiveCoupling(nn.Module):
+    """Exactly invertible additive coupling with an alternating split."""
+
+    def __init__(self, dim: int, hidden_dim: int, update_first: bool) -> None:
+        super().__init__()
+        if dim % 2:
+            raise ValueError(f"coupling dimension must be even, got {dim}")
+        self.half_dim = dim // 2
+        self.update_first = update_first
+        self.conditioner = Conditioner(self.half_dim, hidden_dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        first, second = x.split(self.half_dim, dim=-1)
+        if self.update_first:
+            first = first + self.conditioner(second)
+        else:
+            second = second + self.conditioner(first)
+        return torch.cat((first, second), dim=-1)
+
+    def inverse(self, y: Tensor) -> Tensor:
+        first, second = y.split(self.half_dim, dim=-1)
+        if self.update_first:
+            first = first - self.conditioner(second)
+        else:
+            second = second - self.conditioner(first)
+        return torch.cat((first, second), dim=-1)
+
+
+class InvertibleMNISTMap(nn.Module):
+    """Four additive couplings plus an invertible latent translation.
+
+    With dim=784, hidden_dim=144, and four blocks, the coupling networks have
+    453,728 parameters.  The learned 784-vector translation brings the total
+    to exactly 454,512 parameters while preserving an analytic inverse.
+    """
+
+    def __init__(
+        self, dim: int = INPUT_DIM, hidden_dim: int = 144, blocks: int = 4
+    ) -> None:
+        super().__init__()
+        if blocks < 2:
+            raise ValueError("at least two alternating coupling blocks are required")
+        self.dim = dim
+        self.hidden_dim = hidden_dim
+        self.num_blocks = blocks
+        self.blocks = nn.ModuleList(
+            AdditiveCoupling(dim, hidden_dim, update_first=bool(index % 2))
+            for index in range(blocks)
+        )
+        self.latent_translation = nn.Parameter(torch.zeros(dim))
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = x
+        for block in self.blocks:
+            z = block(z)
+        return z + self.latent_translation
+
+    def decode(self, z: Tensor) -> Tensor:
+        x = z - self.latent_translation
+        for block in reversed(self.blocks):
+            x = block.inverse(x)
+        return x
+
+    def linearizer(self, x: Tensor, core_diag: Tensor) -> Tensor:
+        if core_diag.ndim != 1 or core_diag.numel() != self.dim:
+            raise ValueError(
+                f"core must have shape ({self.dim},), got {tuple(core_diag.shape)}"
+            )
+        return self.decode(self.encode(x) * core_diag)
+
+
+def parse_args(argv: Iterable[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train and audit a globally idempotent MNIST Linearizer projector."
+    )
+    parser.add_argument("--mode", choices=sorted(PRESETS), default="full")
+    parser.add_argument(
+        "--output-dir", type=Path, default=Path("outputs/claim3_image_projector")
+    )
+    parser.add_argument("--data-dir", type=Path, default=Path("data"))
+    parser.add_argument(
+        "--device", default="auto", help="auto, cpu, cuda, cuda:0, or mps"
+    )
+    parser.add_argument("--seed", type=int, default=25191)
+    parser.add_argument("--train-samples", type=int, default=None)
+    parser.add_argument("--test-samples", type=int, default=None)
+    parser.add_argument("--ood-samples", type=int, default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--eval-batch-size", type=int, default=256)
+    parser.add_argument("--learning-rate", type=float, default=2e-3)
+    parser.add_argument("--weight-decay", type=float, default=1e-6)
+    parser.add_argument("--identity-weight", type=float, default=0.25)
+    parser.add_argument("--noise-sigma", type=float, default=0.35)
+    parser.add_argument("--rank", type=int, default=588)
+    parser.add_argument("--hidden-dim", type=int, default=144)
+    parser.add_argument("--blocks", type=int, default=4)
+    parser.add_argument("--gradient-clip", type=float, default=5.0)
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--gallery-count", type=int, default=10)
+    parser.add_argument("--log-every", type=int, default=25)
+    parser.add_argument(
+        "--download",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="allow torchvision to download MNIST if it is not cached",
+    )
+    args = parser.parse_args(argv)
+
+    preset = PRESETS[args.mode]
+    for field in (
+        "train_samples",
+        "test_samples",
+        "ood_samples",
+        "epochs",
+        "batch_size",
+    ):
+        if getattr(args, field) is None:
+            setattr(args, field, getattr(preset, field))
+
+    if not 0 < args.rank <= INPUT_DIM:
+        parser.error(f"--rank must be in [1, {INPUT_DIM}]")
+    if args.train_samples <= 0 or args.test_samples <= 0 or args.ood_samples <= 0:
+        parser.error("sample counts must be positive")
+    if args.epochs <= 0 or args.batch_size <= 0 or args.eval_batch_size <= 0:
+        parser.error("epochs and batch sizes must be positive")
+    if args.noise_sigma < 0 or args.identity_weight < 0:
+        parser.error("noise sigma and identity weight must be non-negative")
+    return args
+
+
+def choose_device(requested: str) -> torch.device:
+    if requested == "auto":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but torch.cuda.is_available() is false")
+    if device.type == "mps" and not torch.backends.mps.is_available():
+        raise RuntimeError("MPS was requested but is not available")
+    return device
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.use_deterministic_algorithms(True, warn_only=True)
+    if hasattr(torch.backends, "cudnn"):
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+    if hasattr(torch, "set_float32_matmul_precision"):
+        torch.set_float32_matmul_precision("highest")
+    if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
+        torch.backends.cuda.matmul.allow_tf32 = False
+
+
+def tensor_sha256(values: Tensor) -> str:
+    array = values.detach().cpu().contiguous().numpy()
+    return hashlib.sha256(array.tobytes()).hexdigest()
+
+
+def file_sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def make_projector_diag(dim: int, rank: int, device: torch.device) -> Tensor:
+    """Return a spatially interleaved binary diagonal projector.
+
+    Evenly spreading the discarded coordinates prevents the untrained control
+    from receiving an accidental advantage by dropping only MNIST's blank
+    bottom rows.  The operator remains a fixed diagonal matrix.
+    """
+
+    diag = torch.ones(dim, dtype=torch.float32)
+    discarded = dim - rank
+    if discarded:
+        # For the default 588/784 split this selects 3, 7, ..., 783: exactly
+        # one discarded coordinate in every group of four.
+        discard_indices = (
+            torch.arange(1, discarded + 1, dtype=torch.int64) * dim // discarded - 1
+        )
+        if torch.unique(discard_indices).numel() != discarded:
+            raise RuntimeError(
+                "failed to construct unique interleaved projector coordinates"
+            )
+        diag[discard_indices] = 0.0
+    if int(diag.sum().item()) != rank:
+        raise RuntimeError("constructed projector has the wrong rank")
+    return diag.to(device)
+
+
+def load_mnist(
+    args: argparse.Namespace,
+) -> tuple[DataLoader, Tensor, Tensor, Tensor, Tensor, Tensor]:
+    transform = transforms.ToTensor()
+    train_dataset = datasets.MNIST(
+        root=str(args.data_dir), train=True, transform=transform, download=args.download
+    )
+    test_dataset = datasets.MNIST(
+        root=str(args.data_dir),
+        train=False,
+        transform=transform,
+        download=args.download,
+    )
+    if args.train_samples > len(train_dataset):
+        raise ValueError(
+            f"requested {args.train_samples} training items, only {len(train_dataset)} exist"
+        )
+    if args.test_samples > len(test_dataset):
+        raise ValueError(
+            f"requested {args.test_samples} test items, only {len(test_dataset)} exist"
+        )
+
+    split_generator = torch.Generator().manual_seed(args.seed + 10)
+    train_indices = torch.randperm(len(train_dataset), generator=split_generator)[
+        : args.train_samples
+    ]
+    test_indices = torch.randperm(len(test_dataset), generator=split_generator)[
+        : args.test_samples
+    ]
+
+    loader_generator = torch.Generator().manual_seed(args.seed + 20)
+    train_loader = DataLoader(
+        Subset(train_dataset, train_indices.tolist()),
+        batch_size=args.batch_size,
+        shuffle=True,
+        generator=loader_generator,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+        drop_last=False,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    # Materialize the held-out set so untrained and trained models see exactly
+    # the same examples and corruption tensor.
+    heldout_loader = DataLoader(
+        Subset(test_dataset, test_indices.tolist()),
+        batch_size=args.eval_batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        persistent_workers=args.num_workers > 0,
+    )
+    images: list[Tensor] = []
+    labels: list[Tensor] = []
+    for image, label in heldout_loader:
+        images.append(image.flatten(1))
+        labels.append(label)
+    clean = torch.cat(images, dim=0)
+    heldout_labels = torch.cat(labels, dim=0)
+
+    eval_generator = torch.Generator().manual_seed(args.seed + 30)
+    noise = torch.randn(clean.shape, generator=eval_generator)
+    noisy = (clean + args.noise_sigma * noise).clamp(0.0, 1.0)
+    return train_loader, clean, noisy, heldout_labels, train_indices, test_indices
+
+
+def make_noise_generator(device: torch.device, seed: int) -> torch.Generator:
+    try:
+        generator = torch.Generator(device=device)
+    except TypeError:
+        generator = torch.Generator(device=device.type)
+    generator.manual_seed(seed)
+    return generator
+
+
+def train_model(
+    model: InvertibleMNISTMap,
+    train_loader: DataLoader,
+    projector_diag: Tensor,
+    args: argparse.Namespace,
+    device: torch.device,
+) -> list[dict[str, float | int]]:
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=max(1, args.epochs), eta_min=args.learning_rate * 0.05
+    )
+    noise_generator = make_noise_generator(device, args.seed + 40)
+    history: list[dict[str, float | int]] = []
+    global_step = 0
+    model.train()
+
+    for epoch in range(1, args.epochs + 1):
+        sums = {"loss": 0.0, "restoration_loss": 0.0, "identity_loss": 0.0}
+        seen = 0
+        epoch_start = time.perf_counter()
+        for clean_images, _ in train_loader:
+            clean = clean_images.flatten(1).to(device, non_blocking=True)
+            noise = torch.randn(
+                clean.shape, dtype=clean.dtype, device=device, generator=noise_generator
+            )
+            noisy = (clean + args.noise_sigma * noise).clamp(0.0, 1.0)
+
+            optimizer.zero_grad(set_to_none=True)
+            restored = model.linearizer(noisy, projector_diag)
+            clean_projection = model.linearizer(clean, projector_diag)
+            restoration_loss = F.mse_loss(restored, clean)
+            identity_loss = F.mse_loss(clean_projection, clean)
+            loss = restoration_loss + args.identity_weight * identity_loss
+            loss.backward()
+            gradient_norm = torch.nn.utils.clip_grad_norm_(
+                model.parameters(), args.gradient_clip
+            )
+            optimizer.step()
+
+            batch_items = clean.shape[0]
+            seen += batch_items
+            sums["loss"] += float(loss.detach()) * batch_items
+            sums["restoration_loss"] += float(restoration_loss.detach()) * batch_items
+            sums["identity_loss"] += float(identity_loss.detach()) * batch_items
+            global_step += 1
+            if args.log_every and global_step % args.log_every == 0:
+                print(
+                    f"step={global_step} epoch={epoch}/{args.epochs} "
+                    f"loss={float(loss.detach()):.6f} "
+                    f"restore={float(restoration_loss.detach()):.6f} "
+                    f"identity={float(identity_loss.detach()):.6f} "
+                    f"grad_norm={float(gradient_norm):.4f}",
+                    flush=True,
+                )
+
+        scheduler.step()
+        row: dict[str, float | int] = {
+            "epoch": epoch,
+            "step": global_step,
+            "loss": sums["loss"] / seen,
+            "restoration_loss": sums["restoration_loss"] / seen,
+            "identity_loss": sums["identity_loss"] / seen,
+            "learning_rate": optimizer.param_groups[0]["lr"],
+            "seconds": time.perf_counter() - epoch_start,
+        }
+        history.append(row)
+        print(
+            f"epoch={epoch}/{args.epochs} mean_loss={row['loss']:.6f} "
+            f"mean_restore={row['restoration_loss']:.6f} "
+            f"mean_identity={row['identity_loss']:.6f} "
+            f"seconds={row['seconds']:.2f}",
+            flush=True,
+        )
+    return history
+
+
+@torch.inference_mode()
+def apply_in_batches(
+    model: InvertibleMNISTMap,
+    values: Tensor,
+    core_diag: Tensor,
+    device: torch.device,
+    batch_size: int,
+) -> Tensor:
+    outputs: list[Tensor] = []
+    model.eval()
+    for start in range(0, values.shape[0], batch_size):
+        batch = values[start : start + batch_size].to(device)
+        outputs.append(model.linearizer(batch, core_diag).cpu())
+    return torch.cat(outputs, dim=0)
+
+
+def reconstruction_metrics(prediction: Tensor, target: Tensor) -> dict[str, float]:
+    per_sample_mse = (prediction - target).square().mean(dim=1)
+    global_mse = float(per_sample_mse.mean())
+    per_sample_psnr = -10.0 * torch.log10(per_sample_mse.clamp_min(1e-12))
+    clipped = prediction.clamp(0.0, 1.0)
+    clipped_mse = float((clipped - target).square().mean())
+    return {
+        "mse": global_mse,
+        "mean_psnr_db": float(per_sample_psnr.mean()),
+        "psnr_from_global_mse_db": -10.0 * math.log10(max(global_mse, 1e-12)),
+        "clipped_mse": clipped_mse,
+        "clipped_psnr_from_global_mse_db": -10.0 * math.log10(max(clipped_mse, 1e-12)),
+    }
+
+
+def error_metrics(actual: Tensor, expected: Tensor) -> dict[str, float]:
+    difference = actual - expected
+    return {
+        "rms": float(torch.sqrt(difference.double().square().mean())),
+        "max_abs": float(difference.abs().max()),
+        "mean_abs": float(difference.abs().double().mean()),
+    }
+
+
+@torch.inference_mode()
+def structural_metrics(
+    model: InvertibleMNISTMap,
+    values: Tensor,
+    core_diag: Tensor,
+    device: torch.device,
+    batch_size: int,
+) -> dict[str, dict[str, float]]:
+    roundtrip_parts: list[Tensor] = []
+    idempotence_parts: list[Tensor] = []
+    model.eval()
+    for start in range(0, values.shape[0], batch_size):
+        original_cpu = values[start : start + batch_size]
+        original = original_cpu.to(device)
+        roundtrip = model.decode(model.encode(original))
+        projected_once = model.linearizer(original, core_diag)
+        projected_twice = model.linearizer(projected_once, core_diag)
+        roundtrip_parts.append((roundtrip - original).cpu())
+        idempotence_parts.append((projected_twice - projected_once).cpu())
+
+    roundtrip_difference = torch.cat(roundtrip_parts, dim=0)
+    idempotence_difference = torch.cat(idempotence_parts, dim=0)
+    zeros_roundtrip = torch.zeros_like(roundtrip_difference)
+    zeros_idempotence = torch.zeros_like(idempotence_difference)
+    return {
+        "g_inverse_roundtrip": error_metrics(roundtrip_difference, zeros_roundtrip),
+        "f_idempotence": error_metrics(idempotence_difference, zeros_idempotence),
+    }
+
+
+def make_ood_domains(sample_count: int, seed: int) -> dict[str, Tensor]:
+    generator = torch.Generator().manual_seed(seed)
+    domains = {
+        "ood_uniform_minus1_to_2": -1.0
+        + 3.0 * torch.rand((sample_count, INPUT_DIM), generator=generator),
+    }
+    for scale in (0.1, 1.0, 3.0):
+        label = str(scale).replace(".", "p")
+        domains[f"ood_gaussian_scale_{label}"] = scale * torch.randn(
+            (sample_count, INPUT_DIM), generator=generator
+        )
+    return domains
+
+
+def save_history(path: Path, history: list[dict[str, float | int]]) -> None:
+    if not history:
+        return
+    with path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(history[0]))
+        writer.writeheader()
+        writer.writerows(history)
+
+
+def save_gallery(
+    path: Path,
+    clean: Tensor,
+    noisy: Tensor,
+    untrained: Tensor,
+    trained: Tensor,
+    trained_twice: Tensor,
+    labels: Tensor,
+    count: int,
+) -> None:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    count = min(count, clean.shape[0])
+    rows = [
+        ("clean", clean),
+        ("corrupt", noisy),
+        ("untrained f(x)", untrained),
+        ("trained f(x)", trained),
+        ("trained f(f(x))", trained_twice),
+    ]
+    figure, axes = plt.subplots(
+        len(rows), count, figsize=(1.45 * count, 1.45 * len(rows)), squeeze=False
+    )
+    for row_index, (row_name, images) in enumerate(rows):
+        for column in range(count):
+            axis = axes[row_index, column]
+            axis.imshow(
+                images[column].reshape(28, 28).clamp(0, 1), cmap="gray", vmin=0, vmax=1
+            )
+            axis.axis("off")
+            if row_index == 0:
+                axis.set_title(f"label {int(labels[column])}", fontsize=8)
+            if column == 0:
+                axis.set_ylabel(row_name, fontsize=9)
+    figure.suptitle(
+        "Architecturally idempotent MNIST restoration (outputs displayed clipped to [0,1])",
+        fontsize=12,
+    )
+    figure.tight_layout(rect=(0, 0, 1, 0.97))
+    figure.savefig(path, dpi=180, bbox_inches="tight")
+    plt.close(figure)
+
+
+def environment_metadata(device: torch.device) -> dict[str, Any]:
+    cuda_name = None
+    if device.type == "cuda":
+        cuda_name = torch.cuda.get_device_name(device)
+    return {
+        "python": sys.version.split()[0],
+        "platform": platform.platform(),
+        "torch": torch.__version__,
+        "torchvision": getattr(sys.modules.get("torchvision"), "__version__", None),
+        "numpy": np.__version__,
+        "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_version": torch.version.cuda,
+        "cuda_device_name": cuda_name,
+        "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
+        "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"),
+    }
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    args = parse_args(argv)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    seed_everything(args.seed)
+    device = choose_device(args.device)
+    started = time.perf_counter()
+    print(f"mode={args.mode} device={device} seed={args.seed}", flush=True)
+
+    (
+        train_loader,
+        clean_test,
+        noisy_test,
+        heldout_labels,
+        train_indices,
+        test_indices,
+    ) = load_mnist(args)
+    np.savez_compressed(
+        args.output_dir / "split_indices.npz",
+        train_indices=train_indices.cpu().numpy(),
+        test_indices=test_indices.cpu().numpy(),
+    )
+
+    model = InvertibleMNISTMap(
+        dim=INPUT_DIM, hidden_dim=args.hidden_dim, blocks=args.blocks
+    ).to(device)
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    projector_diag = make_projector_diag(INPUT_DIM, args.rank, device)
+    nonprojector_diag = projector_diag.clone()
+    nonprojector_diag[nonprojector_diag == 0] = 0.5
+    initial_state = copy.deepcopy(model.state_dict())
+
+    print(
+        f"parameters={parameter_count} projector_rank={int(projector_diag.sum())} "
+        f"train_samples={args.train_samples} heldout_samples={args.test_samples}",
+        flush=True,
+    )
+
+    untrained_output = apply_in_batches(
+        model, noisy_test, projector_diag, device, args.eval_batch_size
+    )
+    untrained_metrics = reconstruction_metrics(untrained_output, clean_test)
+    corrupt_metrics = reconstruction_metrics(noisy_test, clean_test)
+
+    history = train_model(model, train_loader, projector_diag, args, device)
+    save_history(args.output_dir / "training_history.csv", history)
+
+    trained_output = apply_in_batches(
+        model, noisy_test, projector_diag, device, args.eval_batch_size
+    )
+    trained_twice = apply_in_batches(
+        model, trained_output, projector_diag, device, args.eval_batch_size
+    )
+    trained_metrics = reconstruction_metrics(trained_output, clean_test)
+
+    domains = {
+        "heldout_clean_mnist": clean_test,
+        "heldout_noisy_mnist": noisy_test,
+        **make_ood_domains(args.ood_samples, args.seed + 50),
+    }
+    projector_checks: dict[str, Any] = {}
+    nonprojector_checks: dict[str, Any] = {}
+    for name, values in domains.items():
+        print(f"structural_check={name} samples={values.shape[0]}", flush=True)
+        projector_checks[name] = structural_metrics(
+            model, values, projector_diag, device, args.eval_batch_size
+        )
+        nonprojector_checks[name] = structural_metrics(
+            model, values, nonprojector_diag, device, args.eval_batch_size
+        )["f_idempotence"]
+
+    save_gallery(
+        args.output_dir / "gallery.png",
+        clean_test,
+        noisy_test,
+        untrained_output,
+        trained_output,
+        trained_twice,
+        heldout_labels,
+        args.gallery_count,
+    )
+
+    checkpoint_path = args.output_dir / "checkpoint.pt"
+    config_for_artifact = {
+        key: str(value) if isinstance(value, Path) else value
+        for key, value in vars(args).items()
+    }
+    torch.save(
+        {
+            "model_class": "InvertibleMNISTMap",
+            "model_state_dict": model.state_dict(),
+            "initial_state_dict": initial_state,
+            "projector_diag": projector_diag.cpu(),
+            "nonprojector_diag": nonprojector_diag.cpu(),
+            "config": config_for_artifact,
+            "parameter_count": parameter_count,
+            "training_history": history,
+        },
+        checkpoint_path,
+    )
+
+    projector_core_residual = float(
+        (projector_diag.square() - projector_diag).abs().max()
+    )
+    nonprojector_core_residual = float(
+        (nonprojector_diag.square() - nonprojector_diag).abs().max()
+    )
+    worst_idempotence_max = max(
+        check["f_idempotence"]["max_abs"] for check in projector_checks.values()
+    )
+    worst_roundtrip_max = max(
+        check["g_inverse_roundtrip"]["max_abs"] for check in projector_checks.values()
+    )
+    minimum_control_rms = min(check["rms"] for check in nonprojector_checks.values())
+    elapsed = time.perf_counter() - started
+
+    results: dict[str, Any] = {
+        "schema_version": 1,
+        "claim": "A Linearizer f=g^-1 P g is idempotent iff its core is idempotent.",
+        "scope": (
+            "Real-data restorative projection on MNIST; this is not an unconditional "
+            "image-generation experiment."
+        ),
+        "architecture": {
+            "map": "four alternating additive-coupling blocks plus latent translation",
+            "dimension": INPUT_DIM,
+            "hidden_dimension": args.hidden_dim,
+            "coupling_blocks": args.blocks,
+            "parameter_count": parameter_count,
+            "analytic_inverse": True,
+            "core": "fixed interleaved binary diagonal projector",
+            "core_rank": int(projector_diag.sum().item()),
+            "core_idempotence_max_abs": projector_core_residual,
+        },
+        "protocol": {
+            "mode": args.mode,
+            "seed": args.seed,
+            "train_samples": args.train_samples,
+            "heldout_samples": args.test_samples,
+            "ood_samples_per_domain": args.ood_samples,
+            "noise_sigma": args.noise_sigma,
+            "epochs": args.epochs,
+            "batch_size": args.batch_size,
+            "learning_rate": args.learning_rate,
+            "identity_weight": args.identity_weight,
+            "training_loss": "MSE(f(noisy), clean) + identity_weight*MSE(f(clean), clean)",
+            "idempotence_penalty_used": False,
+            "train_indices_sha256": tensor_sha256(train_indices),
+            "test_indices_sha256": tensor_sha256(test_indices),
+        },
+        "heldout_restoration": {
+            "corrupt_input": corrupt_metrics,
+            "untrained_projector": untrained_metrics,
+            "trained_projector": trained_metrics,
+            "mse_improvement_vs_corrupt_fraction": (
+                corrupt_metrics["mse"] - trained_metrics["mse"]
+            )
+            / corrupt_metrics["mse"],
+            "mse_improvement_vs_untrained_fraction": (
+                untrained_metrics["mse"] - trained_metrics["mse"]
+            )
+            / untrained_metrics["mse"],
+        },
+        "global_projector_checks": projector_checks,
+        "nonidempotent_core_control": {
+            "construction": "replace every zero diagonal entry of P with 0.5",
+            "core_a2_minus_a_max_abs": nonprojector_core_residual,
+            "domain_checks": nonprojector_checks,
+            "minimum_domain_rms": minimum_control_rms,
+        },
+        "summary": {
+            "worst_projector_idempotence_max_abs": worst_idempotence_max,
+            "worst_g_inverse_roundtrip_max_abs": worst_roundtrip_max,
+            "minimum_nonprojector_control_rms": minimum_control_rms,
+            "trained_heldout_mse": trained_metrics["mse"],
+            "trained_mean_psnr_db": trained_metrics["mean_psnr_db"],
+            "elapsed_seconds": elapsed,
+        },
+        "training_history": history,
+        "environment": environment_metadata(device),
+        "artifacts": {
+            "checkpoint": checkpoint_path.name,
+            "checkpoint_sha256": file_sha256(checkpoint_path),
+            "gallery": "gallery.png",
+            "training_history": "training_history.csv",
+            "split_indices": "split_indices.npz",
+        },
+    }
+    metrics_path = args.output_dir / "metrics.json"
+    with metrics_path.open("w", encoding="utf-8") as handle:
+        json.dump(results, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+
+    print(
+        "RESULT "
+        f"corrupt_mse={corrupt_metrics['mse']:.6f} "
+        f"untrained_mse={untrained_metrics['mse']:.6f} "
+        f"trained_mse={trained_metrics['mse']:.6f} "
+        f"trained_psnr_db={trained_metrics['mean_psnr_db']:.3f} "
+        f"worst_idempotence_max={worst_idempotence_max:.3e} "
+        f"worst_roundtrip_max={worst_roundtrip_max:.3e} "
+        f"minimum_control_rms={minimum_control_rms:.3e}",
+        flush=True,
+    )
+    print(f"METRICS_JSON={metrics_path.resolve()}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+
+````
+
+
+````output
+/home/ubuntu/samuel/repro/.venv/lib/python3.10/site-packages/huggingface_hub/utils/_experimental.py:59: UserWarning: 'HfApi.run_uv_job' is experimental and might be subject to breaking changes in the future without prior notice. You can disable this warning by setting `HF_HUB_DISABLE_EXPERIMENTAL_WARNING=1` as environment variable.
+  warnings.warn(
+Error: Client error '402 Payment Required' for url 'https://huggingface.co/api/jobs/SabaPivot' (Request ID: Root=1-6a58815c-0a19259b24932e3178a76fae;aa469bb5-baa1-4b4a-aa8e-e72007b079bd)
+For more information check: https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/402
+
+Pre-paid credit balance is insufficient - add more credits to your account to use Jobs.
+Hint: set HF_DEBUG=1 as environment variable for full traceback.
+Your script and Job artifacts will be saved in this bucket: https://huggingface.co/buckets/SabaPivot/jobs-artifacts
+
+````