File size: 5,062 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 | """v41: BitConv — depthwise causal ±1 convolution replaces attention.
Findings so far:
- v38 (global static T×T mix, shared-across-channels) failed catastrophically (3.12 BPC).
Too few degrees of freedom: one pattern for 256 channels can't capture diverse
per-channel routing.
- v39/v40 (attention + content-addressable memory) failed to beat v17. Adding
more routing options didn't help; the underlying local routing was already OK.
- Analysis said 102/120 heads are local-recent (mean argmax-distance < 3 tokens).
v41 takes the "attention is local" finding literally and replaces attention with a
CHANNEL-WISE causal ±1 convolution. Each of d_model channels gets its own learnable
K-length ±1 kernel over the past K tokens. Per-channel (not shared) is the key
distinction from v38. Parameter cost is trivial (D·K), so most of the 5M budget
goes to FFN.
y[b, t, d] = sign( sum_{k=0}^{K-1} sign(kernel[d, k]) * x[b, t-k, d] / sqrt(K) )
Everything stays ±1. No Q/K/V, no softmax, no ALiBi, no multi-head.
Config (5M): d_model=256, n_layers=8, K=16, d_ff=800.
"""
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
class BitCausalConv(nn.Module):
"""Depth-wise causal ±1 1D convolution. One K-length kernel per channel."""
def __init__(self, d_model, kernel_size=16):
super().__init__()
self.d_model = d_model
self.K = kernel_size
# Latent float kernel; forward uses sign(kernel).
self.kernel = nn.Parameter(torch.randn(d_model, kernel_size) * 0.02)
# Scale so sum of K ±1 values has unit variance on average.
self.scale = 1.0 / math.sqrt(kernel_size)
def forward(self, x):
# x: (B, T, D), ±1
B, T, D = x.shape
W = sign_ste(self.kernel) # (D, K) ±1
# F.conv1d expects (B, C, T_in) and weight (C_out, C_in/groups, K).
# Depthwise: C_out = C_in = D, groups = D, weight shape (D, 1, K).
W_conv = W.unsqueeze(1) # (D, 1, K)
x_t = x.transpose(1, 2) # (B, D, T)
# Causal: pad K-1 on the left. Output length = T.
x_pad = F.pad(x_t, (self.K - 1, 0))
y = F.conv1d(x_pad, W_conv, groups=D) # (B, D, T)
y = y.transpose(1, 2) * self.scale # (B, T, D), ~unit scale
return sign_ste_clipped(y)
class BitBlockV41(nn.Module):
def __init__(self, d_model, d_ff, kernel_size):
super().__init__()
self.conv = BitCausalConv(d_model, kernel_size)
self.ffn = BitFFN(d_model, d_ff)
def forward(self, x):
c = self.conv(x)
f = self.ffn(x)
return sign_ste(x + c + f)
class BitLMv41(nn.Module):
def __init__(self, vocab_size=128, d_model=256, n_layers=8, d_ff=800,
kernel_size=16, 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.kernel_size = kernel_size
self.embed = BinaryEmbedding(vocab_size, d_model)
self.blocks = nn.ModuleList([
BitBlockV41(d_model, d_ff, kernel_size) 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__':
for d_ff in (768, 800, 816):
for K in (8, 16, 32):
m = BitLMv41(d_model=256, n_layers=8, d_ff=d_ff, kernel_size=K)
n = sum(p.numel() for p in m.parameters())
print(f'd_ff={d_ff}, K={K}: {n:,} ({n/1e6:.3f}M)')
m = BitLMv41()
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')
|