题目重点：MultiheadAttention、Feedforward、DecoderLayer、GPT：从形状和因果掩码到残差、输出logits与独立验证。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 输入整数token ID[B,T]，无padding，T在已学习位置表范围；d可被head数h整除，每head维d/h。
- 采用pre-LN、learned absolute position、tanh近似GELU与tied输出embedding；线性无bias、LN固定scale1/bias0，均是教学选择。
- 仅前向计算：无autograd、优化器、dropout、KV cache或生成采样。不声称完成训练框架；题目这些要求未知。
解题路线
- 先写形状合同：embedding从[B,T]得[B,T,d]，加position；Q/K/V各投影d，再reshape[B,T,h,dh]转[B,h,T,dh]。score为[B,h,T,T]，乘1/sqrt(dh)。
- 因果mask把列位置s>行位置t设负无穷，softmax只沿最后的key轴。保留自身t，避免首行全被mask造成NaN；减每行最大值后exp，禁止把mask放到softmax之后而不重归一化。
- attention乘V得[B,h,T,dh]，先transpose回[B,T,h,dh]再合并heads，最后输出投影。直接reshape没换轴会错混token与head，维度对不代表语义对。
- Feedforward逐token做d→hidden→d，不跨token混合；使用明确tanh近似GELU。DecoderLayer两段分别x+attention(LN(x))和x+FFN(LN(x))，第二段使用更新后的x。
- GPT堆叠层后final-LN，乘token embedding转置得到[B,T,V] logits，不在类内强加softmax。next-token训练应以logits[:, :-1]预测ids[:, 1:]，并屏蔽padding；本文未实现反传，不能把NumPy前向当可训练GPT。
- 独立验证用逐batch、head、query和历史key的标量循环计算参考，包括projection/normalize/FFN，不调用生产attention。改未来token不能影响前缀；截短输入的同一前缀也应一致；零输出投影的层应成为残差identity。
手动走一遍
B=2,T=3,d=4,h=2时，Q为[2,2,3,2]、score[2,2,3,3]；第一行仅看token0，第二行看0/1，第三行看0/1/2。输出[2,3,V]最后一位置的logits可用于采样下个token。48组不同d/head/T/depth模型与标量实现对照，168组前缀不变性、48组batch独立性、1残差identity及6非法输入均通过。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
import math
import numpy as np
def layer_norm(x, eps=1e-5):
# Teaching forward pass: scale=1, bias=0, normalize last axis only.
mean = x.mean(axis=-1, keepdims=True)
return (x - mean) / np.sqrt(((x - mean) ** 2).mean(axis=-1, keepdims=True) + eps)
def weight(rng, rows, cols):
return rng.normal(0, 1 / math.sqrt(rows), (rows, cols))
class MultiheadAttention:
def __init__(self, width, heads, rng):
if width <= 0 or heads <= 0 or width % heads:
raise ValueError("positive width must be divisible by heads")
self.width, self.heads, self.head_width = width, heads, width // heads
self.wq, self.wk, self.wv, self.wo = [weight(rng, width, width) for _ in range(4)]
def __call__(self, x):
batch, length, width = x.shape
if width != self.width or length == 0:
raise ValueError("expected nonempty [batch, length, width]")
def split(w):
return (x @ w).reshape(batch, length, self.heads, self.head_width).transpose(0, 2, 1, 3)
q, k, v = split(self.wq), split(self.wk), split(self.wv)
scores = (q @ k.swapaxes(-1, -2)) / math.sqrt(self.head_width)
future = np.triu(np.ones((length, length), dtype=bool), k=1)
scores = np.where(future, -np.inf, scores)
exp = np.exp(scores - scores.max(axis=-1, keepdims=True))
probabilities = exp / exp.sum(axis=-1, keepdims=True)
context = (probabilities @ v).transpose(0, 2, 1, 3).reshape(batch, length, width)
return context @ self.wo
class Feedforward:
def __init__(self, width, hidden, rng):
if hidden <= 0:
raise ValueError("hidden must be positive")
self.w1, self.w2 = weight(rng, width, hidden), weight(rng, hidden, width)
def __call__(self, x):
z = x @ self.w1
# Explicit tanh approximation to GELU, not exact erf GELU.
gelu = 0.5 * z * (1 + np.tanh(math.sqrt(2 / math.pi) * (z + 0.044715 * z ** 3)))
return gelu @ self.w2
class DecoderLayer:
def __init__(self, width, heads, hidden, rng):
self.attention = MultiheadAttention(width, heads, rng)
self.feedforward = Feedforward(width, hidden, rng)
def __call__(self, x):
x = x + self.attention(layer_norm(x))
return x + self.feedforward(layer_norm(x))
class GPT:
def __init__(self, vocabulary, max_length, width, heads, hidden, layers, seed=0):
if min(vocabulary, max_length, width, heads, hidden, layers) <= 0:
raise ValueError("positive dimensions required")
rng = np.random.default_rng(seed)
self.tokens = weight(rng, vocabulary, width)
self.positions = weight(rng, max_length, width)
self.layers = [DecoderLayer(width, heads, hidden, rng) for _ in range(layers)]
def __call__(self, ids):
ids = np.asarray(ids)
if ids.ndim != 2 or not np.issubdtype(ids.dtype, np.integer):
raise ValueError("integer token IDs shaped [batch, length]")
batch, length = ids.shape
if batch == 0 or not 0 < length <= len(self.positions):
raise ValueError("nonempty batch and supported sequence length required")
if np.any(ids < 0) or np.any(ids >= len(self.tokens)):
raise ValueError("token ID outside vocabulary")
x = self.tokens[ids] + self.positions[np.arange(length)]
for layer in self.layers:
x = layer(x)
# Tied output embedding, logits for next token at each position.
return layer_norm(x) @ self.tokens.T
复杂度
每层QKV/输出与FFN约O(BT·d²+BT·d·hidden)，attention O(BT²d)；此未融合实现保存O(BhT²)score/probability。词表输出O(BTdV)，位置embedding O(Tmax·d)，参数随层数增加。不能只报attention复杂度忽略大词表投影。
易错点与追问
容易失分的边界
- softmax轴错或mask反了。
- 混heads不transpose。
- 声称LN只要均值不需要variance/eps。
- 训练target不shift导致标签泄露。
- 以纯NumPy前向声称实现完整训练/生产推理。
继续深挖
- padding mask与causal mask怎样组合，避免全masked query？
- 有KV cache时当前位置是offset+t，因果mask如何对齐？
- 如何加入可学习LN参数、bias和训练模式dropout？
- tied embedding对参数量和梯度累积有何影响？
