diff options
Diffstat (limited to 'kaft_mvp/trainers.py')
| -rw-r--r-- | kaft_mvp/trainers.py | 380 |
1 files changed, 380 insertions, 0 deletions
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() + ) |