Add anonymous KAFT MVP reproduction

1 files changed, 262 insertions, 0 deletions

diff --git a/kaft_mvp/experiment.py b/kaft_mvp/experiment.py
new file mode 100644
index 0000000..f867884
--- /dev/null
+++ b/kaft_mvp/experiment.py
@@ -0,0 +1,262 @@
+"""End-to-end MVP experiment and artifact generation."""
+
+from __future__ import annotations
+
+import json
+import platform
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch_geometric
+
+from .data import load_cora
+from .trainers import BPTrainer, KAFTTrainer, seed_everything
+
+
+@dataclass(frozen=True)
+class MVPConfig:
+    dataset: str = "Cora"
+    seeds: tuple[int, ...] = (0, 1, 2)
+    depth: int = 6
+    diagnostic_depth: int = 10
+    hidden_dim: int = 64
+    epochs: int = 200
+    diagnostic_epochs: int = 100
+    eval_every: int = 5
+    learning_rate: float = 0.01
+    weight_decay: float = 5e-4
+    diffusion_alpha: float = 0.5
+    diffusion_steps: int = 10
+    hop_cap: int = 3
+    num_probes: int = 64
+    feedback_learning_rate: float = 0.5
+    align_every: int = 10
+    device: str = "cpu"
+
+
+def _train(
+    trainer: BPTrainer | KAFTTrainer,
+    epochs: int,
+    eval_every: int,
+) -> tuple[dict[str, Any], dict[str, list[float | None]]]:
+    history: dict[str, list[float | None]] = {
+        "train_loss": [],
+        "train_accuracy": [],
+        "validation_accuracy": [],
+        "test_accuracy": [],
+    }
+    best_validation = float("-inf")
+    test_at_best_validation = float("nan")
+    best_epoch = -1
+    for epoch in range(epochs):
+        step = trainer.train_step()
+        history["train_loss"].append(step.loss)
+        history["train_accuracy"].append(step.train_accuracy)
+        if epoch % eval_every == 0 or epoch == epochs - 1:
+            validation = trainer.accuracy("val_mask")
+            test = trainer.accuracy("test_mask")
+            history["validation_accuracy"].append(validation)
+            history["test_accuracy"].append(test)
+            if validation > best_validation:
+                best_validation = validation
+                test_at_best_validation = test
+                best_epoch = epoch
+        else:
+            history["validation_accuracy"].append(None)
+            history["test_accuracy"].append(None)
+    return {
+        "best_validation_accuracy": best_validation,
+        "test_accuracy_at_validation_peak": test_at_best_validation,
+        "best_epoch": best_epoch,
+        "final_train_loss": history["train_loss"][-1],
+        "final_train_accuracy": history["train_accuracy"][-1],
+    }, history
+
+
+def _make_trainer(
+    method: str,
+    data: dict[str, Any],
+    config: MVPConfig,
+    depth: int,
+) -> BPTrainer | KAFTTrainer:
+    common = {
+        "data": data,
+        "depth": depth,
+        "hidden_dim": config.hidden_dim,
+        "learning_rate": config.learning_rate,
+        "weight_decay": config.weight_decay,
+    }
+    if method == "BP":
+        return BPTrainer(**common)
+    if method == "KAFT":
+        return KAFTTrainer(
+            **common,
+            diffusion_alpha=config.diffusion_alpha,
+            diffusion_steps=config.diffusion_steps,
+            hop_cap=config.hop_cap,
+            num_probes=config.num_probes,
+            feedback_learning_rate=config.feedback_learning_rate,
+            align_every=config.align_every,
+        )
+    raise ValueError(f"Unknown method: {method}")
+
+
+def _summary_frame(records: list[dict[str, Any]]) -> pd.DataFrame:
+    rows = []
+    for method in ("BP", "KAFT"):
+        values = np.asarray(
+            [
+                row["result"]["test_accuracy_at_validation_peak"]
+                for row in records
+                if row["method"] == method
+            ],
+            dtype=float,
+        )
+        rows.append(
+            {
+                "method": method,
+                "mean_test_accuracy_percent": 100.0 * values.mean(),
+                "std_test_accuracy_percent": (
+                    100.0 * values.std(ddof=1)
+                    if len(values) > 1
+                    else 0.0
+                ),
+                "seeds": len(values),
+            }
+        )
+    return pd.DataFrame(rows)
+
+
+def _plot_histories(
+    records: list[dict[str, Any]],
+    output_path: Path,
+) -> None:
+    fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
+    for method, color in (("BP", "#4C78A8"), ("KAFT", "#E45756")):
+        histories = [
+            row["history"]
+            for row in records
+            if row["method"] == method
+        ]
+        loss = np.asarray([history["train_loss"] for history in histories])
+        test = np.asarray(
+            [
+                [
+                    np.nan if value is None else value
+                    for value in history["test_accuracy"]
+                ]
+                for history in histories
+            ]
+        )
+        epochs = np.arange(1, loss.shape[1] + 1)
+        median_loss = np.nanmedian(loss, axis=0)
+        axes[0].plot(epochs, median_loss, color=color, label=method)
+        axes[0].fill_between(
+            epochs,
+            np.nanpercentile(loss, 25, axis=0),
+            np.nanpercentile(loss, 75, axis=0),
+            color=color,
+            alpha=0.15,
+        )
+        evaluated = ~np.isnan(test).all(axis=0)
+        axes[1].plot(
+            epochs[evaluated],
+            100.0 * np.nanmedian(test[:, evaluated], axis=0),
+            color=color,
+            label=method,
+        )
+    axes[0].set_xlabel("Epoch")
+    axes[0].set_ylabel("Training cross-entropy")
+    axes[0].set_yscale("log")
+    axes[1].set_xlabel("Epoch")
+    axes[1].set_ylabel("Test accuracy (%)")
+    for axis in axes:
+        axis.grid(alpha=0.2)
+        axis.legend(frameon=False)
+    fig.suptitle("Cora, identical 6-layer GCN forward model")
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=180, bbox_inches="tight")
+    plt.close(fig)
+
+
+def run_mvp(
+    config: MVPConfig | None = None,
+    output_dir: str | Path = "artifacts",
+    data_root: str | Path = "data",
+) -> dict[str, Any]:
+    """Run BP/KAFT accuracy and BP transport diagnostics end to end."""
+    config = config or MVPConfig()
+    if config.dataset != "Cora":
+        raise ValueError("The minimal package currently exposes only Cora")
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    data = load_cora(root=data_root, device=config.device)
+
+    records: list[dict[str, Any]] = []
+    for seed in config.seeds:
+        for method in ("BP", "KAFT"):
+            seed_everything(seed)
+            trainer = _make_trainer(method, data, config, config.depth)
+            result, history = _train(
+                trainer,
+                config.epochs,
+                config.eval_every,
+            )
+            records.append(
+                {
+                    "seed": seed,
+                    "method": method,
+                    "result": result,
+                    "history": history,
+                }
+            )
+
+    diagnostics: list[dict[str, Any]] = []
+    for seed in config.seeds:
+        seed_everything(seed)
+        trainer = _make_trainer(
+            "BP",
+            data,
+            config,
+            config.diagnostic_depth,
+        )
+        result, _ = _train(
+            trainer,
+            config.diagnostic_epochs,
+            config.eval_every,
+        )
+        diagnostics.append(
+            {
+                "seed": seed,
+                "training_result": result,
+                **trainer.gradient_diagnostic(),
+            }
+        )
+
+    summary = _summary_frame(records)
+    summary.to_csv(output_dir / "mvp_summary.csv", index=False)
+    _plot_histories(records, output_dir / "training_curves.png")
+
+    payload = {
+        "config": {
+            **asdict(config),
+            "seeds": list(config.seeds),
+        },
+        "environment": {
+            "python": platform.python_version(),
+            "torch": torch.__version__,
+            "torch_geometric": torch_geometric.__version__,
+            "device": config.device,
+        },
+        "summary": summary.to_dict(orient="records"),
+        "records": records,
+        "gradient_diagnostics": diagnostics,
+    }
+    with (output_dir / "mvp_results.json").open("w") as handle:
+        json.dump(payload, handle, indent=2, allow_nan=False)
+    return payload