item 31 pre-registered (Nils's design): synthetic memory tokens — burst states → per-band-layer KV prefix via KVMemoryAdapter (in-place attention wrap, bit-exact disarmed, gated silent init); composed with frozen item-29 arm-1; +29 tf+fr in-flight note (30.5, FR term hurts)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-17 13:25:25 +02:00
parent 11604d0949
commit b18ffbd739
6 changed files with 231 additions and 11 deletions
@@ -0,0 +1,107 @@
+"""Item 31: synthetic memory tokens — per-layer KV prefix (Nils's design).
+
+"The one read mechanism a frozen transformer natively possesses is
+attention over the KV cache." A KVMemoryAdapter maps carried L30 states
+to (k, v) entries appended to every band layer's attention — frozen
+heads consume the state through their existing machinery. Keys live in
+post-RoPE space (prefix-tuning convention: learned keys place themselves
+where frozen queries already point).
+
+Activation is global-flag-scoped like lora_band.LoopLoRA: a custom
+attention interface "kvmem" wraps the stock eager path; when _MEM is
+armed and the module is a band layer, memory kv columns are appended
+(mask columns visible to every query).
+"""
+import torch
+import torch.nn as nn
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.gemma4.modeling_gemma4 import (
+    eager_attention_forward)
+
+_MEM = {"on": False, "kv": None, "band": (14, 30), "base": None}
+
+
+def arm_memory(kv):
+    """kv: dict layer_idx -> (k, v), each (B, kv_heads, M, head_dim)."""
+    _MEM["kv"] = kv
+    _MEM["on"] = kv is not None
+
+
+def kvmem_attention(module, query, key, value, attention_mask, **kwargs):
+    if (_MEM["on"] and _MEM["kv"] is not None
+            and getattr(module, "layer_idx", -1) in _MEM["kv"]):
+        mk, mv, gate = _MEM["kv"][module.layer_idx]
+        m = mk.shape[2]
+        key = torch.cat([key, mk.to(dtype=key.dtype, device=key.device)], 2)
+        value = torch.cat([value, mv.to(dtype=value.dtype,
+                                        device=value.device)], 2)
+        if attention_mask is not None:
+            pad = torch.zeros(*attention_mask.shape[:-1], m,
+                              dtype=attention_mask.dtype,
+                              device=attention_mask.device) + gate.to(
+                attention_mask.dtype)
+            attention_mask = torch.cat([attention_mask, pad], -1)
+    base = _MEM["base"] or eager_attention_forward
+    return base(module, query, key, value, attention_mask, **kwargs)
+
+
+def install(model):
+    """Wrap the registered implementation IN PLACE — the config name is
+    untouched, so mask construction and every other branch stays
+    bit-identical; only the (q,k,v,mask) seam gains the memory append."""
+    if _MEM.get("installed"):
+        return
+    impl = model.config._attn_implementation
+    _MEM["base"] = ALL_ATTENTION_FUNCTIONS.get_interface(
+        impl, eager_attention_forward)
+    try:
+        ALL_ATTENTION_FUNCTIONS[impl] = kvmem_attention
+    except TypeError:
+        ALL_ATTENTION_FUNCTIONS.register(impl, kvmem_attention)
+    _MEM["installed"] = True
+
+
+class KVMemoryAdapter(nn.Module):
+    """Carried states (B, M, d) -> per-band-layer post-RoPE (k, v).
+    Per-layer geometry (head_dim, kv heads) read off the live modules —
+    gemma-4 mixes sliding (256) and global (512) head dims in the band."""
+
+    def __init__(self, model, band=(14, 30), code=512):
+        super().__init__()
+        self.band = band
+        tm = model.model.language_model
+        d = model.config.get_text_config().hidden_size
+        self.trunk = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, code),
+                                   nn.GELU())
+        self.geom = {}
+        self.k_heads = nn.ModuleDict()
+        self.v_heads = nn.ModuleDict()
+        tc = model.config.get_text_config()
+        for l in range(band[0], band[1] + 1):
+            att = tm.layers[l].self_attn
+            hd = att.head_dim
+            kvh = (tc.num_global_key_value_heads
+                   if getattr(att, "use_alternative_attention", False)
+                   else tc.num_key_value_heads)
+            self.geom[l] = (kvh, hd)
+            self.k_heads[str(l)] = nn.Linear(code, kvh * hd)
+            self.v_heads[str(l)] = nn.Linear(code, kvh * hd)
+            nn.init.zeros_(self.v_heads[str(l)].weight)
+            nn.init.zeros_(self.v_heads[str(l)].bias)
+        # learnable per-layer attention gate on the memory columns:
+        # init -10 -> memory is invisible until training opens it
+        self.gate = nn.ParameterDict(
+            {str(l): nn.Parameter(torch.tensor(-10.0))
+             for l in range(band[0], band[1] + 1)})
+
+    def forward(self, states):
+        B, M, _ = states.shape
+        c = self.trunk(states.float())
+        kv = {}
+        for l in range(self.band[0], self.band[1] + 1):
+            kvh, hd = self.geom[l]
+            k = self.k_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
+            v = self.v_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
+            kv[l] = (k, v, self.gate[str(l)])
+        return kv