path: root/sdil/transformer.py

"""Matched decoder-Transformer components for local-learning crossovers.

The forward graph is deliberately identical for BP, ordinary FA, clean KP,
and SDIL.  Only the vector transported through parameterized affine maps is
changed.  Parameter-free Jacobians (residual addition, LayerNorm, GELU, and
softmax attention) remain local and exact.
"""
from dataclasses import dataclass
import math
from typing import Dict, Iterable, Optional

import torch
from torch import nn
import torch.nn.functional as F


_FEEDBACK_METHODS = {"fa", "clean_kp", "sdil"}
_SUPPORTED_METHODS = {"bp", "dfa"} | _FEEDBACK_METHODS


class _FeedbackLinearFunction(torch.autograd.Function):
    """Linear map with fixed or locally plastic feedback.

    ``feedback`` has the same orientation as ``weight``.  Consequently the
    transported vector is ``delta @ feedback``, while both plastic matrices
    can be updated from the locally available ``delta.T @ input`` correlation.
    The feedback correlation is recomputed rather than copied from the
    forward-weight gradient.
    """

    @staticmethod
    def forward(
            ctx, x, weight, feedback, bias, method_code, traffic_ratio,
            raw_rms, innovation_rms, traffic_rms):
        output = F.linear(x, weight, bias)
        ctx.save_for_backward(x, feedback, output)
        ctx.method_code = int(method_code)
        ctx.traffic_ratio = float(traffic_ratio)
        ctx.has_bias = bias is not None
        ctx.raw_rms = raw_rms
        ctx.innovation_rms = innovation_rms
        ctx.traffic_rms = traffic_rms
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, feedback, soma = ctx.saved_tensors
        task_instruction = grad_output
        traffic = torch.zeros_like(task_instruction)
        if ctx.method_code == 2 and ctx.traffic_ratio:
            # A paired neutral observation exposes the component predictable
            # from the local somatic response.  Match it to a frozen multiple
            # of task-instruction RMS without changing its somatic direction.
            task_rms = task_instruction.square().mean().sqrt()
            centered_soma = soma - soma.mean(dim=-1, keepdim=True)
            soma_rms = centered_soma.square().mean().sqrt().clamp_min(1e-30)
            traffic = (
                centered_soma * task_rms * ctx.traffic_ratio / soma_rms)
        raw_apical = task_instruction + traffic
        neutral_prediction = traffic
        innovation = raw_apical - neutral_prediction

        input_flat = x.reshape(-1, x.shape[-1])
        delta_flat = innovation.reshape(-1, innovation.shape[-1])
        grad_input = innovation @ feedback
        grad_weight = delta_flat.t() @ input_flat
        grad_feedback = None
        if ctx.method_code in (1, 2):
            # This is intentionally a second evaluation of the local
            # correlation, not an assignment from grad_weight.
            grad_feedback = delta_flat.t() @ input_flat
        grad_bias = None
        if ctx.has_bias:
            grad_bias = delta_flat.sum(dim=0)

        with torch.no_grad():
            ctx.raw_rms.copy_(raw_apical.square().mean().sqrt())
            ctx.innovation_rms.copy_(innovation.square().mean().sqrt())
            ctx.traffic_rms.copy_(traffic.square().mean().sqrt())
        return (
            grad_input, grad_weight, grad_feedback, grad_bias,
            None, None, None, None, None)


class FeedbackLinear(nn.Module):
    """A forward-matched affine map for BP, FA, clean KP, or SDIL."""

    def __init__(
            self, in_features: int, out_features: int, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator,
            bias: bool = False, init_std: float = 0.02,
            traffic_ratio: float = 4.0, dtype=torch.float32):
        super().__init__()
        if method not in _SUPPORTED_METHODS:
            raise ValueError(f"unsupported feedback-linear method: {method}")
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.method = method
        self.traffic_ratio = float(traffic_ratio if method == "sdil" else 0.0)
        self.weight = nn.Parameter(torch.empty(
            out_features, in_features, dtype=dtype))
        nn.init.normal_(
            self.weight, mean=0.0, std=init_std,
            generator=forward_generator)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype))
        else:
            self.register_parameter("bias", None)

        if method in _FEEDBACK_METHODS:
            feedback = torch.empty(
                out_features, in_features, dtype=dtype)
            nn.init.normal_(
                feedback, mean=0.0, std=init_std,
                generator=feedback_generator)
            if method == "fa":
                self.register_buffer("feedback", feedback)
            else:
                self.feedback = nn.Parameter(feedback)
        else:
            self.register_buffer("feedback", None)
        self.register_buffer("last_raw_rms", torch.zeros((), dtype=dtype))
        self.register_buffer(
            "last_innovation_rms", torch.zeros((), dtype=dtype))
        self.register_buffer("last_traffic_rms", torch.zeros((), dtype=dtype))

    def forward(self, x):
        if self.method in {"bp", "dfa"}:
            return F.linear(x, self.weight, self.bias)
        method_code = {"fa": 0, "clean_kp": 1, "sdil": 2}[self.method]
        return _FeedbackLinearFunction.apply(
            x, self.weight, self.feedback, self.bias, method_code,
            self.traffic_ratio, self.last_raw_rms,
            self.last_innovation_rms, self.last_traffic_rms)

    def extra_repr(self):
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, method={self.method}")


@dataclass(frozen=True)
class LocalTransformerConfig:
    vocab_size: int = 65
    context_length: int = 64
    depth: int = 4
    width: int = 128
    heads: int = 4
    mlp_ratio: int = 4
    dropout: float = 0.0
    bias: bool = False
    init_std: float = 0.02
    traffic_ratio: float = 4.0
    seed: int = 2027

    def __post_init__(self):
        if self.width % self.heads:
            raise ValueError("width must be divisible by heads")
        if self.depth < 1 or self.context_length < 1:
            raise ValueError("depth and context_length must be positive")


class LocalCausalSelfAttention(nn.Module):

    def __init__(
            self, config: LocalTransformerConfig, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator, dtype=torch.float32):
        super().__init__()
        self.heads = config.heads
        self.head_width = config.width // config.heads
        self.width = config.width
        common = dict(
            method=method, forward_generator=forward_generator,
            feedback_generator=feedback_generator, bias=config.bias,
            init_std=config.init_std, traffic_ratio=config.traffic_ratio,
            dtype=dtype)
        self.q = FeedbackLinear(config.width, config.width, **common)
        self.k = FeedbackLinear(config.width, config.width, **common)
        self.v = FeedbackLinear(config.width, config.width, **common)
        self.output = FeedbackLinear(config.width, config.width, **common)
        causal = torch.tril(torch.ones(
            config.context_length, config.context_length, dtype=torch.bool))
        self.register_buffer("causal_mask", causal, persistent=False)
        self.dropout = float(config.dropout)

    def forward(self, x):
        batch, time, width = x.shape

        def split_heads(value):
            return value.view(
                batch, time, self.heads, self.head_width).transpose(1, 2)

        query = split_heads(self.q(x))
        key = split_heads(self.k(x))
        value = split_heads(self.v(x))
        scores = query @ key.transpose(-2, -1)
        scores = scores * self.head_width ** -0.5
        mask = self.causal_mask[:time, :time]
        scores = scores.masked_fill(~mask, float("-inf"))
        attention = F.softmax(scores, dim=-1)
        attention = F.dropout(
            attention, p=self.dropout, training=self.training)
        mixed = attention @ value
        mixed = mixed.transpose(1, 2).contiguous().view(batch, time, width)
        return self.output(mixed)


class LocalTransformerBlock(nn.Module):

    def __init__(
            self, config: LocalTransformerConfig, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator, dtype=torch.float32):
        super().__init__()
        self.ln_attention = nn.LayerNorm(config.width, dtype=dtype)
        self.attention = LocalCausalSelfAttention(
            config, method, forward_generator, feedback_generator, dtype)
        self.ln_mlp = nn.LayerNorm(config.width, dtype=dtype)
        hidden = config.mlp_ratio * config.width
        common = dict(
            method=method, forward_generator=forward_generator,
            feedback_generator=feedback_generator, bias=config.bias,
            init_std=config.init_std, traffic_ratio=config.traffic_ratio,
            dtype=dtype)
        self.mlp_in = FeedbackLinear(config.width, hidden, **common)
        self.mlp_out = FeedbackLinear(hidden, config.width, **common)
        self.dropout = float(config.dropout)

    def forward(self, x):
        x = x + F.dropout(
            self.attention(self.ln_attention(x)),
            p=self.dropout, training=self.training)
        x = x + F.dropout(
            self.mlp_out(F.gelu(self.mlp_in(self.ln_mlp(x)))),
            p=self.dropout, training=self.training)
        return x


class LocalDecoderTransformer(nn.Module):
    """Depth-scaled, forward-matched character decoder."""

    def __init__(
            self, config: LocalTransformerConfig, method: str = "bp",
            dtype=torch.float32):
        super().__init__()
        if method not in _SUPPORTED_METHODS:
            raise ValueError(f"unsupported Transformer method: {method}")
        self.config = config
        self.method = method
        forward_generator = torch.Generator().manual_seed(config.seed)
        feedback_generator = torch.Generator().manual_seed(config.seed + 1)
        self.token_embedding = nn.Embedding(
            config.vocab_size, config.width, dtype=dtype)
        self.position_embedding = nn.Parameter(torch.empty(
            config.context_length, config.width, dtype=dtype))
        nn.init.normal_(
            self.token_embedding.weight, mean=0.0, std=config.init_std,
            generator=forward_generator)
        nn.init.normal_(
            self.position_embedding, mean=0.0, std=config.init_std,
            generator=forward_generator)
        self.blocks = nn.ModuleList([
            LocalTransformerBlock(
                config, method, forward_generator, feedback_generator, dtype)
            for _ in range(config.depth)])
        self.final_norm = nn.LayerNorm(config.width, dtype=dtype)
        self.head = FeedbackLinear(
            config.width, config.vocab_size, method, forward_generator,
            feedback_generator, bias=False, init_std=config.init_std,
            traffic_ratio=config.traffic_ratio, dtype=dtype)
        if method == "dfa":
            dfa_generator = torch.Generator().manual_seed(config.seed + 2)
            feedback_scale = config.init_std
            self.register_buffer("dfa_block_feedback", torch.empty(
                config.depth, config.vocab_size, config.width, dtype=dtype))
            self.register_buffer("dfa_embedding_feedback", torch.empty(
                config.vocab_size, config.width, dtype=dtype))
            self.register_buffer("dfa_final_norm_feedback", torch.empty(
                config.vocab_size, config.width, dtype=dtype))
            nn.init.normal_(
                self.dfa_block_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
            nn.init.normal_(
                self.dfa_embedding_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
            nn.init.normal_(
                self.dfa_final_norm_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
        else:
            self.register_buffer("dfa_block_feedback", None)
            self.register_buffer("dfa_embedding_feedback", None)
            self.register_buffer("dfa_final_norm_feedback", None)

    def forward(
            self, tokens, targets: Optional[torch.Tensor] = None,
            return_cache: bool = False):
        if tokens.ndim != 2:
            raise ValueError("tokens must have shape (batch, time)")
        if tokens.shape[1] > self.config.context_length:
            raise ValueError("sequence exceeds configured context length")
        positions = self.position_embedding[:tokens.shape[1]]
        hidden = self.token_embedding(tokens) + positions
        hidden = F.dropout(
            hidden, p=self.config.dropout, training=self.training)
        block_inputs = []
        for block in self.blocks:
            if return_cache:
                block_inputs.append(hidden.detach())
            hidden = block(hidden)
        final_input = hidden
        normalized = self.final_norm(hidden)
        logits = self.head(normalized)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                targets.reshape(-1))
        result = {"logits": logits, "loss": loss, "hidden": hidden}
        if return_cache:
            result["block_inputs"] = block_inputs
            result["final_input"] = final_input.detach()
            result["normalized"] = normalized.detach()
        return result

    @staticmethod
    def _assign_gradients(parameters, gradients):
        for parameter, gradient in zip(parameters, gradients):
            if parameter.grad is None:
                parameter.grad = gradient.detach().clone()
            else:
                parameter.grad.copy_(gradient.detach())

    def dfa_gradients(
            self, tokens, targets: Optional[torch.Tensor] = None,
            cache: Optional[Dict[str, torch.Tensor]] = None,
            output_error: Optional[torch.Tensor] = None):
        """Populate strict block-DFA gradients without a task-loss backward.

        Every decoder block receives a separate fixed projection of the
        analytical output error.  Inputs are detached at block boundaries;
        autograd is used only for each explicitly local block objective.
        Embeddings and the final normalization receive their own fixed direct
        projections, while the vocabulary head uses its exact local delta.
        """
        if self.method != "dfa":
            raise ValueError("dfa_gradients is only valid for method='dfa'")
        if cache is None:
            with torch.no_grad():
                cache = self.forward(tokens, return_cache=True)
        if output_error is None:
            if targets is None:
                raise ValueError("targets or output_error must be provided")
            probabilities = torch.softmax(cache["logits"], dim=-1)
            one_hot = F.one_hot(
                targets, self.config.vocab_size).to(probabilities.dtype)
            output_error = (
                probabilities - one_hot) / targets.numel()
        output_error = output_error.detach()
        for parameter in self.parameters():
            parameter.grad = None

        error_flat = output_error.reshape(-1, self.config.vocab_size)
        normalized_flat = cache["normalized"].reshape(
            -1, self.config.width)
        self.head.weight.grad = error_flat.t() @ normalized_flat
        if self.head.bias is not None:
            self.head.bias.grad = error_flat.sum(dim=0)

        final_input = cache["final_input"].detach()
        normalized = self.final_norm(final_input)
        final_field = output_error @ self.dfa_final_norm_feedback
        final_parameters = tuple(self.final_norm.parameters())
        final_objective = torch.sum(normalized * final_field)
        final_gradients = torch.autograd.grad(
            final_objective, final_parameters)
        self._assign_gradients(final_parameters, final_gradients)

        for index, (block, block_input) in enumerate(zip(
                self.blocks, cache["block_inputs"])):
            local_input = block_input.detach()
            local_output = block(local_input)
            local_field = output_error @ self.dfa_block_feedback[index]
            local_parameters = tuple(block.parameters())
            local_objective = torch.sum(local_output * local_field)
            local_gradients = torch.autograd.grad(
                local_objective, local_parameters)
            self._assign_gradients(local_parameters, local_gradients)

        embedding_field = output_error @ self.dfa_embedding_feedback
        token_activity = self.token_embedding(tokens)
        position_activity = self.position_embedding[:tokens.shape[1]]
        embedding_objective = (
            torch.sum(token_activity * embedding_field)
            + torch.sum(position_activity * embedding_field.sum(dim=0)))
        embedding_parameters = (
            self.token_embedding.weight, self.position_embedding)
        embedding_gradients = torch.autograd.grad(
            embedding_objective, embedding_parameters)
        self._assign_gradients(
            embedding_parameters, embedding_gradients)
        return {
            "output_error_rms": float(
                output_error.square().mean().sqrt()),
            "block_field_rms": [
                float((output_error @ feedback).square().mean().sqrt())
                for feedback in self.dfa_block_feedback],
            "uses_task_loss_backward": False,
            "detached_block_boundaries": len(self.blocks),
        }

    def feedback_linears(self) -> Iterable[FeedbackLinear]:
        return (
            module for module in self.modules()
            if isinstance(module, FeedbackLinear))

    def forward_parameters(self) -> Iterable[nn.Parameter]:
        feedback_ids = {
            id(module.feedback)
            for module in self.feedback_linears()
            if isinstance(module.feedback, nn.Parameter)}
        return (
            parameter for parameter in self.parameters()
            if id(parameter) not in feedback_ids)

    @property
    def n_forward_parameters(self):
        return sum(parameter.numel() for parameter in self.forward_parameters())

    @property
    def n_feedback_parameters(self):
        affine_feedback = sum(
            module.feedback.numel()
            for module in self.feedback_linears()
            if module.feedback is not None)
        dfa_feedback = sum(
            tensor.numel() for tensor in (
                self.dfa_block_feedback, self.dfa_embedding_feedback,
                self.dfa_final_norm_feedback)
            if tensor is not None)
        return affine_feedback + dfa_feedback

    def teaching_statistics(self) -> Dict[str, float]:
        modules = list(self.feedback_linears())
        if not modules:
            return {
                "raw_rms": 0.0, "innovation_rms": 0.0,
                "traffic_rms": 0.0}
        return {
            name: float(torch.stack([
                getattr(module, f"last_{name}")
                for module in modules]).mean())
            for name in ("raw_rms", "innovation_rms", "traffic_rms")}

    @torch.no_grad()
    def set_feedback_equal_to_forward(self):
        for module in self.feedback_linears():
            if module.feedback is None:
                raise ValueError("BP modules do not contain feedback tensors")
            module.feedback.copy_(module.weight)