baseline: add matched Transformer PEPITA

1 files changed, 118 insertions, 4 deletions

diff --git a/sdil/transformer.py b/sdil/transformer.py
index 608f43f..ee1e69f 100644
--- a/sdil/transformer.py
+++ b/sdil/transformer.py
@@ -15,7 +15,7 @@ import torch.nn.functional as F
 
 
 _FEEDBACK_METHODS = {"fa", "clean_kp", "sdil"}
-_SUPPORTED_METHODS = {"bp", "dfa"} | _FEEDBACK_METHODS
+_SUPPORTED_METHODS = {"bp", "dfa", "pepita"} | _FEEDBACK_METHODS
 
 
 class _FeedbackLinearFunction(torch.autograd.Function):
@@ -126,7 +126,7 @@ class FeedbackLinear(nn.Module):
         self.register_buffer("last_traffic_rms", torch.zeros((), dtype=dtype))
 
     def forward(self, x):
-        if self.method in {"bp", "dfa"}:
+        if self.method in {"bp", "dfa", "pepita"}:
             return F.linear(x, self.weight, self.bias)
         method_code = {"fa": 0, "clean_kp": 1, "sdil": 2}[self.method]
         return _FeedbackLinearFunction.apply(
@@ -292,23 +292,44 @@ class LocalDecoderTransformer(nn.Module):
             self.register_buffer("dfa_block_feedback", None)
             self.register_buffer("dfa_embedding_feedback", None)
             self.register_buffer("dfa_final_norm_feedback", None)
+        if method == "pepita":
+            pepita_generator = torch.Generator().manual_seed(
+                config.seed + 3)
+            limit = math.sqrt(6.0 / config.width) * 0.05
+            projection = (
+                2.0 * torch.rand(
+                    config.vocab_size, config.width,
+                    generator=pepita_generator, dtype=dtype) - 1.0
+            ) * limit
+            self.register_buffer("pepita_input_feedback", projection)
+        else:
+            self.register_buffer("pepita_input_feedback", None)
 
     def forward(
             self, tokens, targets: Optional[torch.Tensor] = None,
-            return_cache: bool = False):
+            return_cache: bool = False,
+            embedding_offset: Optional[torch.Tensor] = None):
         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
+        if embedding_offset is not None:
+            if embedding_offset.shape != hidden.shape:
+                raise ValueError("embedding_offset shape does not match tokens")
+            hidden = hidden + embedding_offset
+        embedded = hidden
         hidden = F.dropout(
             hidden, p=self.config.dropout, training=self.training)
         block_inputs = []
+        block_outputs = []
         for block in self.blocks:
             if return_cache:
                 block_inputs.append(hidden.detach())
             hidden = block(hidden)
+            if return_cache:
+                block_outputs.append(hidden.detach())
         final_input = hidden
         normalized = self.final_norm(hidden)
         logits = self.head(normalized)
@@ -320,6 +341,8 @@ class LocalDecoderTransformer(nn.Module):
         result = {"logits": logits, "loss": loss, "hidden": hidden}
         if return_cache:
             result["block_inputs"] = block_inputs
+            result["block_outputs"] = block_outputs
+            result["embedded"] = embedded.detach()
             result["final_input"] = final_input.detach()
             result["normalized"] = normalized.detach()
         return result
@@ -410,6 +433,94 @@ class LocalDecoderTransformer(nn.Module):
             "detached_block_boundaries": len(self.blocks),
         }
 
+    def pepita_gradients(self, tokens, targets):
+        """Populate two-presentation PEPITA/ERIN local gradients.
+
+        The analytical output error is projected into the continuous token
+        embedding stream.  Each block then uses its first-minus-second output
+        difference and the second-presentation input.  Block boundaries are
+        detached.  Since discrete token IDs cannot themselves be perturbed,
+        the embedding table and positional code use the directly observable
+        embedding difference as their local field.
+        """
+        if self.method != "pepita":
+            raise ValueError(
+                "pepita_gradients is only valid for method='pepita'")
+        with torch.no_grad():
+            clean = self.forward(tokens, return_cache=True)
+            one_hot = F.one_hot(
+                targets, self.config.vocab_size).to(clean["logits"].dtype)
+            clean_error = torch.softmax(clean["logits"], dim=-1) - one_hot
+            embedding_offset = clean_error @ self.pepita_input_feedback
+            modulated = self.forward(
+                tokens, return_cache=True,
+                embedding_offset=embedding_offset)
+            modulated_error = (
+                torch.softmax(modulated["logits"], dim=-1) - one_hot)
+        for parameter in self.parameters():
+            parameter.grad = None
+        observations = targets.numel()
+
+        error_flat = modulated_error.reshape(
+            -1, self.config.vocab_size) / observations
+        normalized_flat = modulated["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 = modulated["final_input"].detach()
+        normalized = self.final_norm(final_input)
+        final_field = (
+            clean["normalized"] - modulated["normalized"]).detach()
+        final_parameters = tuple(self.final_norm.parameters())
+        final_objective = torch.sum(
+            normalized * final_field) / observations
+        final_gradients = torch.autograd.grad(
+            final_objective, final_parameters)
+        self._assign_gradients(final_parameters, final_gradients)
+
+        for block, local_input, clean_output, modulated_output in zip(
+                self.blocks, modulated["block_inputs"],
+                clean["block_outputs"], modulated["block_outputs"]):
+            local_output = block(local_input.detach())
+            local_field = (clean_output - modulated_output).detach()
+            local_parameters = tuple(block.parameters())
+            local_objective = torch.sum(
+                local_output * local_field) / observations
+            local_gradients = torch.autograd.grad(
+                local_objective, local_parameters)
+            self._assign_gradients(local_parameters, local_gradients)
+
+        base_embedding = (
+            self.token_embedding(tokens)
+            + self.position_embedding[:tokens.shape[1]])
+        embedding_field = (
+            clean["embedded"] - modulated["embedded"]).detach()
+        embedding_objective = torch.sum(
+            base_embedding * embedding_field) / observations
+        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 {
+            "clean_loss": float(F.cross_entropy(
+                clean["logits"].reshape(-1, self.config.vocab_size),
+                targets.reshape(-1))),
+            "embedding_offset_rms": float(
+                embedding_offset.square().mean().sqrt()),
+            "mean_block_field_rms": float(torch.stack([
+                (first - second).square().mean().sqrt()
+                for first, second in zip(
+                    clean["block_outputs"],
+                    modulated["block_outputs"])]).mean()),
+            "training_presentations": 2,
+            "uses_task_loss_backward": False,
+            "detached_block_boundaries": len(self.blocks),
+        }
+
     def feedback_linears(self) -> Iterable[FeedbackLinear]:
         return (
             module for module in self.modules()
@@ -439,7 +550,10 @@ class LocalDecoderTransformer(nn.Module):
                 self.dfa_block_feedback, self.dfa_embedding_feedback,
                 self.dfa_final_norm_feedback)
             if tensor is not None)
-        return affine_feedback + dfa_feedback
+        pepita_feedback = (
+            self.pepita_input_feedback.numel()
+            if self.pepita_input_feedback is not None else 0)
+        return affine_feedback + dfa_feedback + pepita_feedback
 
     def teaching_statistics(self) -> Dict[str, float]:
         modules = list(self.feedback_linears())