File size: 2,961 Bytes
d997548 | 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 | """
Cell Memory: per-layer attention over cell embeddings.
Each cell gets a learned embedding that serves as both key and value.
Forward pass queries over cells, adds result residually.
New params per layer:
- cell_emb: [max_cells, d_model]
- q_proj, out_proj: [d_model, d_model]
For 28 layers, max_cells=256, d_model=1536:
cell_emb: 28 × 256 × 1536 = 11.0 M
q_proj + out_proj: 28 × 2 × 1536² = 132 M
Total: ~143 M new params (14% of R7)
Checkpoint compatibility:
- Zero-init out_proj → memory contributes 0.0 at start
- Checkpoint 3200 loads fine; new params learned from scratch
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class CellMemory(nn.Module):
def __init__(self, d_model, max_cells=256, n_heads=8):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.max_cells = max_cells
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.cell_emb = nn.Embedding(max_cells, d_model)
self.q_proj = nn.Linear(d_model, d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
# Init: small cell emb, xavier q, ZERO out (identity at start)
nn.init.normal_(self.cell_emb.weight, std=0.02)
nn.init.xavier_uniform_(self.q_proj.weight)
nn.init.zeros_(self.out_proj.weight)
def forward(self, hidden, slot_ids, slot_mask):
"""
hidden: [B, T, d]
slot_ids: [N] long tensor in [0, max_cells)
slot_mask: [N] 1 for active, 0 for padding
"""
device = hidden.device
if slot_ids.device != device:
slot_ids = slot_ids.to(device)
if slot_mask.device != device:
slot_mask = slot_mask.to(device)
B, T, D = hidden.shape
N = slot_ids.shape[0]
H, hd = self.n_heads, self.head_dim
# Multi-head query
q = self.q_proj(hidden).view(B, T, H, hd).transpose(1, 2) # [B,H,T,hd]
# Cell keys = values
k = self.cell_emb(slot_ids) # [N, D]
k = k.view(N, H, hd).transpose(0, 1) # [H, N, hd]
v = k
# Attention
attn = torch.matmul(q, k.transpose(-2, -1)) * (hd ** -0.5) # [B,H,T,N]
attn = attn.masked_fill(slot_mask.view(1, 1, 1, N) == 0, float("-inf"))
attn = F.softmax(attn, dim=-1)
attn = torch.nan_to_num(attn, nan=0.0)
out = torch.matmul(attn, v) # [B,H,T,hd]
out = out.transpose(1, 2).contiguous().view(B, T, D)
return self.out_proj(out)
if __name__ == "__main__":
import torch
cm = CellMemory(d_model=64, max_cells=16, n_heads=4)
hidden = torch.randn(1, 8, 64)
slot_ids = torch.arange(4)
slot_mask = torch.ones(4)
out = cm(hidden, slot_ids, slot_mask)
print(f"out shape: {tuple(out.shape)}")
print(f"all zeros (out_proj=0): {torch.allclose(out, torch.zeros_like(out))}")
|