File size: 5,362 Bytes
4754707 | 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 | """v39: BitHop — binary attention + content-addressable memory bank.
v38 (pure BitMixer, no attention) tanked at 3.12 BPC vs v17's 1.68, proving
that binary attention — even though the analysis showed heads look locally
biased — is still doing data-dependent work a static mix cannot replicate.
But the analysis was not wrong about the *missing* capability: binary attention
with ±1 QK and ALiBi cannot form long-range content-addressable routing.
v39 adds it explicitly via a Hopfield-style prototype memory bank per layer:
BitHopHead: M learnable ±1 prototype (key, value) pairs.
Each token: q = BitLinear(x); scores = q @ keys.T (integer);
Gumbel-argmax over M picks one prototype; out = BitLinear(V_sel).
All signal paths ±1. Prototypes are latent-float weights whose sign() is used
at forward — identical discipline to BitLinear weights. At eval, argmax is a
pure integer compare over M options.
Each v39 block: attention(x) + hop(x) + ffn(x), all summed into residual.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from model import sign_ste, sign_ste_clipped, BitLinear, BitFFN, BinaryEmbedding
from model_v16 import _get_tau
from model_v18 import IntBinaryAttention
class BitHopHead(nn.Module):
"""Content-addressable memory: M ±1 prototype key/value pairs, top-1 routing."""
def __init__(self, d_model, n_proto=32):
super().__init__()
self.n_proto = n_proto
self.q_proj = BitLinear(d_model, d_model, binarize_input=True)
self.o_proj = BitLinear(d_model, d_model, binarize_input=True)
# Latent-float prototypes; forward uses sign(proto).
self.key_proto = nn.Parameter(torch.randn(n_proto, d_model) * 0.02)
self.val_proto = nn.Parameter(torch.randn(n_proto, d_model) * 0.02)
def forward(self, x):
B, T, D = x.shape
q = self.q_proj(x) # (B, T, D) ±1
K = sign_ste(self.key_proto) # (M, D) ±1
V = sign_ste(self.val_proto) # (M, D) ±1
scores = q @ K.t() # (B, T, M) integer popcount
tau = _get_tau(scores.device)
if scores.requires_grad:
g = -torch.log(-torch.log(torch.rand_like(scores).clamp(min=1e-9)) + 1e-9)
y_soft = F.softmax((scores + g) / tau, dim=-1)
y_hard = torch.zeros_like(y_soft).scatter_(-1, y_soft.argmax(-1, keepdim=True), 1.0)
A = y_soft + (y_hard - y_soft).detach()
else:
A = torch.zeros_like(scores).scatter_(-1, scores.argmax(-1, keepdim=True), 1.0)
out = A @ V # (B, T, D)
return self.o_proj(out)
class BitBlockV39(nn.Module):
def __init__(self, d_model, n_heads, d_ff, n_proto):
super().__init__()
self.attn = IntBinaryAttention(d_model, n_heads)
self.hop = BitHopHead(d_model, n_proto)
self.ffn = BitFFN(d_model, d_ff)
def forward(self, x):
a = self.attn(x)
h = self.hop(x)
f = self.ffn(x)
return sign_ste(x + a + h + f)
class BitLMv39(nn.Module):
def __init__(self, vocab_size=128, d_model=256, n_layers=8, n_heads=8,
d_ff=288, n_proto=32, max_seq_len=256):
super().__init__()
self.vocab_size = vocab_size
self.d_model = d_model
self.n_layers = n_layers
self.max_seq_len = max_seq_len
self.n_proto = n_proto
self.embed = BinaryEmbedding(vocab_size, d_model)
self.blocks = nn.ModuleList([
BitBlockV39(d_model, n_heads, d_ff, n_proto) for _ in range(n_layers)
])
self.out_codebook = nn.Parameter(torch.randn(vocab_size, d_model) * 0.02)
self.logit_scale = nn.Parameter(torch.tensor(1.0 / math.sqrt(d_model)))
self.out_bias = nn.Parameter(torch.zeros(vocab_size))
def forward(self, idx, targets=None):
x = self.embed(idx)
for blk in self.blocks:
x = blk(x)
W_out = sign_ste(self.out_codebook)
scores = torch.matmul(x, W_out.t())
logits = scores * self.logit_scale + self.out_bias
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new_tokens=200, temperature=1.0, top_k=None):
self.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.max_seq_len:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / max(temperature, 1e-5)
if top_k is not None:
v, _ = torch.topk(logits, top_k)
logits[logits < v[:, [-1]]] = -float('inf')
probs = F.softmax(logits, dim=-1)
nxt = torch.multinomial(probs, num_samples=1)
idx = torch.cat([idx, nxt], dim=1)
return idx
if __name__ == '__main__':
from model_v16 import set_gumbel_tau
set_gumbel_tau(0.5)
for d_ff in (256, 272, 288, 304):
m = BitLMv39(d_ff=d_ff)
n = sum(p.numel() for p in m.parameters())
print(f'd_ff={d_ff}: {n:,} ({n/1e6:.3f}M)')
m = BitLMv39()
x = torch.randint(0, 128, (2, 64))
y = torch.randint(0, 128, (2, 64))
logits, loss = m(x, y)
loss.backward()
print(f'loss={loss.item():.3f}, backward OK')
|