File size: 7,123 Bytes
f930dac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """PNS-Bind (E3): exact identity addressing, learned semantic content.
Seven prior models kept a learned recurrent state that was causally inert.
E3 tests one narrow hypothesis: learned semantic memory becomes durable when
exact identity/binding structure protects it from unrelated updates.
mode="bind" : an event writes ONLY the slot its exact address names;
every other slot is carried through BITWISE unchanged.
mode="unbound" : identical capacity (192 slots), identical everything else,
but the original unstructured mechanism - every slot is
written every event by attention. This is the E3-A3
capacity-matched control that separates "binding isolation"
from "more learned-state capacity".
The exact side supplies only identity: slot -> (entity id, attribute type) and
"which slot is authoritative for this query". It never supplies a value; the
payload s_i in R^d is learned, and the answer must be decoded from it.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
from ..data.binding import N_BIND_SLOTS
from ..world.schema import ET, ENUM_VOCAB
from .modules import SDPA, Block, HeadSystem, RecordEncoder, add_mask
@dataclass
class BindConfig:
vocab: int = 8192
d: int = 192
heads: int = 6
enc_layers: int = 3
n_slots: int = N_BIND_SLOTS
T_ev: int = 48
mode: str = "bind" # "bind" | "unbound"
gate_bias: float = 0.0
full_heads: bool = False # Stage B: the complete factorized output system
class PNSBind(nn.Module):
def __init__(self, cfg: BindConfig):
super().__init__()
self.cfg = cfg
d = cfg.d
self.tok_emb = nn.Embedding(cfg.vocab, d)
self.pos_emb = nn.Parameter(torch.zeros(cfg.T_ev, d))
self.enc = nn.ModuleList(Block(d, cfg.heads) for _ in range(cfg.enc_layers))
self.enc_ln = nn.LayerNorm(d)
self.etype_emb = nn.Embedding(len(ET), d)
self.dt_proj = nn.Linear(8, d)
self.slot0 = nn.Parameter(torch.randn(cfg.n_slots, d) * 0.02)
# exact identity features of a slot: entity and attribute type only
self.key_ent = nn.Embedding(4097, d)
self.key_attr = nn.Embedding(32, d)
self.wq = nn.Linear(d, d, bias=False)
self.write_attn = SDPA(d, cfg.heads)
self.gate = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, d))
nn.init.constant_(self.gate[-1].bias, cfg.gate_bias)
self.read_q = nn.Parameter(torch.randn(1, d) * 0.02)
self.hop = SDPA(d, cfg.heads) # learned hop 2 over slot keys
self.mix = nn.Sequential(nn.LayerNorm(3 * d), nn.Linear(3 * d, d), nn.GELU(),
nn.Linear(d, d))
self.enum = nn.Linear(d, len(ENUM_VOCAB))
if cfg.full_heads:
# identical head system to every prior model, so Stage B differs
# from the failed runs ONLY in the state mechanism
self.recenc = RecordEncoder(d, self.tok_emb)
self.heads = HeadSystem(d)
def initial_state(self, B, device):
return self.slot0.detach().float().unsqueeze(0).repeat(B, 1, 1).to(device)
def encode_event(self, tok, etype, dt):
pad = tok == 0
x = self.tok_emb(tok.long()) + self.pos_emb[: tok.shape[1]]
for blk in self.enc:
x = blk(x, self_padding_mask=pad)
x = self.enc_ln(x)
den = (~pad).sum(1, keepdim=True).clamp(min=1)
v = x.masked_fill(pad.unsqueeze(-1), 0).sum(1) / den
v = v + self.etype_emb(etype.long()) + self.dt_proj(RecordEncoder.age_feats(dt))
return x, v, pad
def step(self, state, tok, etype, dt, bind_write, bind_read,
slot_ent, slot_attr, freeze_writes=False, payload_scale=1.0,
rec_bank=None, live_mask=None):
"""state [B,S,d] fp32. bind_write/bind_read: [B] int, -1 = none."""
cfg = self.cfg
B, S, d = state.shape
x, v, pad = self.encode_event(tok, etype, dt)
cdt = x.dtype
keys = self.key_ent((slot_ent.long() + 1).clamp(0, 4096)) \
+ self.key_attr(slot_attr.long().clamp(0, 31))
if not freeze_writes:
st = state.to(cdt)
if cfg.mode == "bind":
# ---- addressed write: gather ONE slot, update it, scatter back.
# Slots other than bind_write are never touched, so they are
# carried through bitwise (asserted in tests).
w = bind_write.long()
has = (w >= 0)
idx = w.clamp(min=0).view(B, 1, 1).expand(-1, 1, d)
cur = torch.gather(st, 1, idx) # [B,1,d]
kcur = torch.gather(keys.to(cdt), 1, idx)
q = self.wq(cur + kcur)
kv = torch.cat([x, v.unsqueeze(1)], 1)
kvm = torch.cat([pad, torch.zeros(B, 1, dtype=torch.bool,
device=x.device)], 1)
delta = self.write_attn(q, kv, kv, mask=add_mask(kvm, None, q.dtype))
g = torch.sigmoid(self.gate(torch.cat([cur, delta,
v.unsqueeze(1)], -1)))
new = ((1 - g) * cur + g * delta).float()
upd = state.scatter(1, w.clamp(min=0).view(B, 1, 1).expand(-1, 1, d),
new)
state = torch.where(has.view(B, 1, 1), upd, state)
else:
# ---- unstructured control: every slot written every event
q = self.wq(st + keys.to(cdt))
kv = torch.cat([x, v.unsqueeze(1)], 1)
kvm = torch.cat([pad, torch.zeros(B, 1, dtype=torch.bool,
device=x.device)], 1)
delta = self.write_attn(q, kv, kv, mask=add_mask(kvm, None, q.dtype))
g = torch.sigmoid(self.gate(torch.cat(
[st, delta, v.unsqueeze(1).expand(-1, S, -1)], -1)))
state = ((1 - g) * st + g * delta).float()
# ---------------- read
st = (state * payload_scale).to(cdt)
kslots = st + keys.to(cdt)
r = bind_read.long()
addressed = torch.gather(
st, 1, r.clamp(min=0).view(B, 1, 1).expand(-1, 1, d))
addressed = addressed * (r >= 0).view(B, 1, 1).to(cdt) # [B,1,d]
# learned second hop: query the slot bank with what hop 1 produced
hop2 = self.hop(addressed + v.unsqueeze(1), kslots, kslots)
dense = (st * 0 + kslots).mean(1, keepdim=True) # dense read
h = self.mix(torch.cat([addressed.squeeze(1), hop2.squeeze(1),
v], -1))
if self.cfg.full_heads and rec_bank is not None:
out = self.heads(h, v, kslots.mean(1), rec_bank, live_mask)
out["enum"] = out["enum"] + self.enum(h) # binding-read pathway
return state, out
return state, {"enum": self.enum(h)}
|