path: root/kaft_mvp/trainers.py

"""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()
        )