diff options
Diffstat (limited to 'kaft_mvp')
| -rw-r--r-- | kaft_mvp/__init__.py | 5 | ||||
| -rw-r--r-- | kaft_mvp/data.py | 100 | ||||
| -rw-r--r-- | kaft_mvp/experiment.py | 262 | ||||
| -rw-r--r-- | kaft_mvp/trainers.py | 380 |
4 files changed, 747 insertions, 0 deletions
diff --git a/kaft_mvp/__init__.py b/kaft_mvp/__init__.py new file mode 100644 index 0000000..864ff32 --- /dev/null +++ b/kaft_mvp/__init__.py @@ -0,0 +1,5 @@ +"""Minimal KAFT reproduction package.""" + +from .experiment import MVPConfig, run_mvp + +__all__ = ["MVPConfig", "run_mvp"] diff --git a/kaft_mvp/data.py b/kaft_mvp/data.py new file mode 100644 index 0000000..2f25d80 --- /dev/null +++ b/kaft_mvp/data.py @@ -0,0 +1,100 @@ +"""Dataset and graph-operator utilities for the MVP.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from urllib.request import urlopen + +import torch +import torch_geometric.transforms as transforms +from torch_geometric.datasets import Planetoid + + +PLANETOID_FILES = ( + "x", + "tx", + "allx", + "y", + "ty", + "ally", + "graph", + "test.index", +) + + +def _ensure_planetoid_raw(root: str | Path, name: str) -> None: + """Pre-fetch from raw.githubusercontent.com to avoid a flaky redirect.""" + raw_dir = Path(root) / name / "raw" + raw_dir.mkdir(parents=True, exist_ok=True) + lower_name = name.lower() + base_url = ( + "https://raw.githubusercontent.com/kimiyoung/" + "planetoid/master/data" + ) + for suffix in PLANETOID_FILES: + filename = f"ind.{lower_name}.{suffix}" + destination = raw_dir / filename + if destination.exists() and destination.stat().st_size > 0: + continue + temporary = destination.with_suffix(destination.suffix + ".part") + with urlopen(f"{base_url}/{filename}", timeout=60) as response: + temporary.write_bytes(response.read()) + temporary.replace(destination) + + +def spmm(matrix: torch.Tensor, dense: torch.Tensor) -> torch.Tensor: + """Multiply a native PyTorch sparse matrix by a dense matrix.""" + return torch.sparse.mm(matrix, dense) + + +def normalized_adjacency( + edge_index: torch.Tensor, + num_nodes: int, +) -> torch.Tensor: + """Return A-hat = D^-1/2 (A + I) D^-1/2 as a coalesced COO tensor.""" + row, col = edge_index + loops = torch.arange(num_nodes, device=edge_index.device) + row = torch.cat((row, loops)) + col = torch.cat((col, loops)) + + degree = torch.zeros(num_nodes, device=edge_index.device) + degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float32)) + inverse_sqrt = degree.pow(-0.5) + inverse_sqrt.masked_fill_(torch.isinf(inverse_sqrt), 0.0) + values = inverse_sqrt[row] * inverse_sqrt[col] + + return torch.sparse_coo_tensor( + torch.stack((row, col)), + values, + (num_nodes, num_nodes), + ).coalesce() + + +def load_cora( + root: str | Path = "data", + device: str | torch.device = "cpu", +) -> dict[str, Any]: + """Download Cora and return the tensors used by both training rules.""" + _ensure_planetoid_raw(root, "Cora") + dataset = Planetoid( + root=str(root), + name="Cora", + transform=transforms.NormalizeFeatures(), + ) + graph = dataset[0] + adjacency = normalized_adjacency( + graph.edge_index, + graph.num_nodes, + ).to(device) + return { + "x": graph.x.to(device), + "y": graph.y.to(device), + "adjacency": adjacency, + "train_mask": graph.train_mask.to(device), + "val_mask": graph.val_mask.to(device), + "test_mask": graph.test_mask.to(device), + "num_nodes": graph.num_nodes, + "num_features": dataset.num_features, + "num_classes": dataset.num_classes, + } 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 diff --git a/kaft_mvp/trainers.py b/kaft_mvp/trainers.py new file mode 100644 index 0000000..bff2d7c --- /dev/null +++ b/kaft_mvp/trainers.py @@ -0,0 +1,380 @@ +"""Minimal BP and KAFT trainers with an identical bias-free GCN forward pass.""" + +from __future__ import annotations + +import random +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as functional +from sklearn.linear_model import LogisticRegression +from sklearn.preprocessing import StandardScaler + +from .data import spmm + + +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) + + +def _weight_shapes( + input_dim: int, + hidden_dim: int, + output_dim: int, + depth: int, +) -> list[tuple[int, int]]: + dims = [input_dim] + [hidden_dim] * (depth - 1) + [output_dim] + return list(zip(dims[:-1], dims[1:])) + + +@dataclass +class StepResult: + loss: float + train_accuracy: float + + +class BPTrainer: + """Standard backpropagation through a plain deep GCN.""" + + def __init__( + self, + data: dict[str, Any], + depth: int, + hidden_dim: int, + learning_rate: float, + weight_decay: float, + ) -> None: + self.data = data + self.depth = depth + self.weights = torch.nn.ParameterList() + for input_dim, output_dim in _weight_shapes( + data["num_features"], + hidden_dim, + data["num_classes"], + depth, + ): + weight = torch.nn.Parameter( + torch.empty(input_dim, output_dim, device=data["x"].device) + ) + torch.nn.init.xavier_uniform_(weight) + self.weights.append(weight) + self.optimizer = torch.optim.Adam( + self.weights, + lr=learning_rate, + weight_decay=weight_decay, + ) + + def forward( + self, + retain_intermediate_gradients: bool = False, + ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]: + hidden = self.data["x"] + preactivations: list[torch.Tensor] = [] + hidden_states: list[torch.Tensor] = [hidden] + for layer, weight in enumerate(self.weights): + preactivation = spmm( + self.data["adjacency"], + hidden @ weight, + ) + if retain_intermediate_gradients: + preactivation.retain_grad() + preactivations.append(preactivation) + hidden = ( + functional.relu(preactivation) + if layer < self.depth - 1 + else preactivation + ) + hidden_states.append(hidden) + return hidden, { + "preactivations": preactivations, + "hidden_states": hidden_states, + } + + def train_step(self) -> StepResult: + self.optimizer.zero_grad(set_to_none=True) + logits, _ = self.forward() + mask = self.data["train_mask"] + loss = functional.cross_entropy( + logits[mask], + self.data["y"][mask], + ) + loss.backward() + self.optimizer.step() + accuracy = ( + logits[mask].argmax(dim=1) == self.data["y"][mask] + ).float().mean() + return StepResult(float(loss.item()), float(accuracy.item())) + + @torch.no_grad() + def accuracy(self, mask_name: str) -> float: + logits, _ = self.forward() + mask = self.data[mask_name] + return float( + (logits[mask].argmax(dim=1) == self.data["y"][mask]) + .float() + .mean() + .item() + ) + def gradient_diagnostic(self) -> dict[str, Any]: + """Measure parameter/error gradients and a standardized hidden probe.""" + self.optimizer.zero_grad(set_to_none=True) + logits, intermediates = self.forward( + retain_intermediate_gradients=True + ) + mask = self.data["train_mask"] + loss = functional.cross_entropy( + logits[mask], + self.data["y"][mask], + ) + loss.backward() + + weight_norms = [ + 0.0 if weight.grad is None else float(weight.grad.norm().item()) + for weight in self.weights + ] + error_norms = [ + 0.0 + if preactivation.grad is None + else float(preactivation.grad.norm().item()) + for preactivation in intermediates["preactivations"] + ] + penultimate_hidden = intermediates["hidden_states"][-2].detach() + train_x = penultimate_hidden[mask].cpu().numpy() + test_x = penultimate_hidden[self.data["test_mask"]].cpu().numpy() + train_y = self.data["y"][mask].cpu().numpy() + test_y = self.data["y"][self.data["test_mask"]].cpu().numpy() + scaler = StandardScaler().fit(train_x) + probe = LogisticRegression( + C=1.0, + max_iter=2000, + random_state=0, + ).fit(scaler.transform(train_x), train_y) + probe_accuracy = float( + probe.score(scaler.transform(test_x), test_y) + ) + return { + "all_weight_gradients_exact_zero": all( + value == 0.0 for value in weight_norms + ), + "weight_gradient_frobenius": weight_norms, + "preactivation_error_frobenius": error_norms, + "output_adjacent_error_frobenius": error_norms[-2], + "standardized_penultimate_probe_accuracy": probe_accuracy, + "diagnostic_loss": float(loss.item()), + } + + +class KAFTTrainer: + """Kronecker-Aligned Feedback Training for the same forward GCN.""" + + def __init__( + self, + data: dict[str, Any], + depth: int, + hidden_dim: int, + learning_rate: float, + weight_decay: float, + 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, + ) -> None: + self.data = data + self.depth = depth + self.diffusion_alpha = diffusion_alpha + self.diffusion_steps = diffusion_steps + self.hop_cap = hop_cap + self.num_probes = num_probes + self.feedback_learning_rate = feedback_learning_rate + self.align_every = align_every + self.step_index = 0 + + self.weights = torch.nn.ParameterList() + for input_dim, output_dim in _weight_shapes( + data["num_features"], + hidden_dim, + data["num_classes"], + depth, + ): + weight = torch.nn.Parameter( + torch.empty(input_dim, output_dim, device=data["x"].device), + requires_grad=False, + ) + torch.nn.init.xavier_uniform_(weight) + self.weights.append(weight) + + self.feedback = [ + torch.randn( + data["num_classes"], + hidden_dim, + device=data["x"].device, + ) + * 0.01 + for _ in range(depth - 1) + ] + self.optimizer = torch.optim.Adam( + self.weights, + lr=learning_rate, + weight_decay=weight_decay, + ) + + @torch.no_grad() + def forward( + self, + ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]: + hidden = self.data["x"] + preactivations: list[torch.Tensor] = [] + hidden_states: list[torch.Tensor] = [] + layer_inputs: list[torch.Tensor] = [] + for layer, weight in enumerate(self.weights): + layer_inputs.append(hidden) + preactivation = spmm( + self.data["adjacency"], + hidden @ weight, + ) + preactivations.append(preactivation) + hidden = ( + functional.relu(preactivation) + if layer < self.depth - 1 + else preactivation + ) + if layer < self.depth - 1: + hidden_states.append(hidden) + return hidden, { + "preactivations": preactivations, + "hidden_states": hidden_states, + "layer_inputs": layer_inputs, + } + + @torch.no_grad() + def _fixed_error_diffusion(self, error: torch.Tensor) -> torch.Tensor: + """Apply the fixed linear D(A-hat) used in the submitted rule.""" + diffused = error.clone() + for _ in range(self.diffusion_steps): + diffused = ( + (1.0 - self.diffusion_alpha) * error + + self.diffusion_alpha + * spmm(self.data["adjacency"], diffused) + ) + return diffused + + @torch.no_grad() + def _align_feature_factors(self) -> None: + """Move each R_l toward the chain-normalized suffix probe target.""" + hidden_dim = self.feedback[0].shape[1] + for layer in range(self.depth - 1): + probes = torch.randn( + hidden_dim, + self.num_probes, + device=self.data["x"].device, + ) + target_probes = probes + mean_probe_norm = probes.norm(dim=0).mean() + for suffix_layer in range(layer + 1, self.depth): + target_probes = ( + self.weights[suffix_layer].t() @ target_probes + ) + column_norm = target_probes.norm( + dim=0, + keepdim=True, + ).clamp(min=1e-8) + target_probes = ( + target_probes / column_norm * mean_probe_norm + ) + target = target_probes @ probes.t() / self.num_probes + current = self.feedback[layer] + cosine = functional.cosine_similarity( + current.reshape(1, -1), + target.reshape(1, -1), + ).item() + current_norm = current.norm().clamp(min=1e-8) + target_norm = target.norm().clamp(min=1e-8) + cosine_gradient = ( + target / (current_norm * target_norm) + - cosine * current / current_norm.square() + ) + updated = ( + current + + self.feedback_learning_rate * cosine_gradient + ) + self.feedback[layer] = updated / updated.norm( + dim=0, + keepdim=True, + ).clamp(min=1e-8) + + @torch.no_grad() + def train_step(self) -> StepResult: + self.optimizer.zero_grad(set_to_none=True) + logits, intermediates = self.forward() + mask = self.data["train_mask"] + labels = self.data["y"] + probabilities = functional.softmax(logits, dim=1) + one_hot = functional.one_hot( + labels, + self.data["num_classes"], + ).float() + output_error = torch.zeros_like(probabilities) + output_error[mask] = ( + probabilities[mask] - one_hot[mask] + ) / mask.sum().clamp(min=1) + diffused_error = self._fixed_error_diffusion(output_error) + + if self.step_index % self.align_every == 0: + self._align_feature_factors() + + gradients: list[torch.Tensor] = [] + for layer in range(self.depth - 1): + topology_error = diffused_error + hops = min(self.depth - 1 - layer, self.hop_cap) + for _ in range(hops): + topology_error = spmm( + self.data["adjacency"], + topology_error, + ) + hidden_error = ( + topology_error @ self.feedback[layer] + ) * (intermediates["preactivations"][layer] > 0) + local_error = spmm( + self.data["adjacency"], + hidden_error, + ) + gradients.append( + intermediates["layer_inputs"][layer].t() @ local_error + ) + + output_local_error = spmm( + self.data["adjacency"], + output_error, + ) + gradients.append( + intermediates["layer_inputs"][-1].t() @ output_local_error + ) + for weight, gradient in zip(self.weights, gradients): + weight.grad = gradient + self.optimizer.step() + self.step_index += 1 + + loss = functional.cross_entropy(logits[mask], labels[mask]) + accuracy = ( + logits[mask].argmax(dim=1) == labels[mask] + ).float().mean() + return StepResult(float(loss.item()), float(accuracy.item())) + + @torch.no_grad() + def accuracy(self, mask_name: str) -> float: + logits, _ = self.forward() + mask = self.data[mask_name] + return float( + (logits[mask].argmax(dim=1) == self.data["y"][mask]) + .float() + .mean() + .item() + ) |