ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """对照组：标准的 GPT 式 Transformer（因果自注意力），结构上和 ZeroX 一一对应。 | |
| 位置编码两种可选（《Attention Is All You Need》§3.5 两种都试过，结果几乎一样）： | |
| pos="learned" 可学习的位置向量，只有 block_size 个，处理不了更长的序列 | |
| pos="sin" 论文用的正弦编码，任意长度都能算（论文选它就是因为"可能可以外推"） | |
| """ | |
| import numpy as np | |
| from zerocg import Tensor | |
| from zerocg import functional as F | |
| from zerocg import nn | |
| def sinusoidal_encoding(T, d): | |
| """PE(pos, 2i) = sin(pos / 10000^(2i/d))，PE(pos, 2i+1) = cos(...)。论文 §3.5。""" | |
| pos = np.arange(T)[:, None] | |
| div = 10000 ** (np.arange(0, d, 2) / d) | |
| pe = np.zeros((T, d), dtype=np.float32) | |
| pe[:, 0::2] = np.sin(pos / div) | |
| pe[:, 1::2] = np.cos(pos / div) | |
| return pe | |
| class Block(nn.Module): | |
| def __init__(self, n_embd, n_head, dropout=0.0): | |
| self.ln1 = nn.LayerNorm(n_embd) | |
| self.attn = nn.CausalSelfAttention(n_embd, n_head, dropout) | |
| self.ln2 = nn.LayerNorm(n_embd) | |
| self.fc1 = nn.Linear(n_embd, 4 * n_embd) | |
| self.fc2 = nn.Linear(4 * n_embd, n_embd) | |
| self.drop = nn.Dropout(dropout) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln1(x)) | |
| return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x))))) | |
| class TransformerLM(nn.Module): | |
| def __init__(self, vocab_size, block_size, n_embd=64, n_head=4, n_layer=2, dropout=0.0, pos="learned"): | |
| assert pos in ("learned", "sin") | |
| self.block_size = block_size | |
| self.pos_kind = pos | |
| self.tok = nn.Embedding(vocab_size, n_embd) | |
| if pos == "learned": | |
| self.pos = nn.Embedding(block_size, n_embd) | |
| self.blocks = [Block(n_embd, n_head, dropout) for _ in range(n_layer)] | |
| self.ln_f = nn.LayerNorm(n_embd) | |
| self.head = nn.Linear(n_embd, vocab_size) | |
| def forward(self, idx): | |
| T = idx.shape[1] | |
| if self.pos_kind == "learned": | |
| assert T <= self.block_size, f"位置编码只有 {self.block_size} 个，处理不了 {T} 个字" | |
| x = self.tok(idx) + self.pos(np.arange(T)) | |
| else: | |
| x = self.tok(idx) + Tensor(sinusoidal_encoding(T, self.tok.weight.shape[1])) | |
| for b in self.blocks: | |
| x = b(x) | |
| return self.head(self.ln_f(x)) | |
| def state_floats(self, T): | |
| """推理到第 T 个字时 KV 缓存的大小（浮点数个数），随 T 线性增长。""" | |
| return sum(2 * T * b.ln1.weight.shape[0] for b in self.blocks) | |