TikTok 手写 Transformer Decoder：因果注意力与完整前向

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 没有指定decoder-only或encoder-decoder；示例选择decoder-only，先声明合同。
• NumPy只写前向，learned position、pre-LN、tanh近似GELU、tied输出embedding；无padding/dropout/KVcache/可学LN仿射项，不声称生产完整实现。

解题路线

1. 输入ids[B,T]查token embedding并加position，得到[B,T,d]；Q/K/V线性投影拆为[B,H,T,dh]，dh=d/H。
2. score=QKᵀ/√dh，未来j>i置负无穷，按最后轴稳定softmax，再乘V，转置拼回[B,T,d]接输出投影。因果mask必须作用在softmax前。
3. 每层x←x+Attention(LN(x))，随后x←x+FFN(LN(x))；FFN逐token扩到f再缩回d。最后Norm和词表矩阵得到[B,T,V]logits。
4. 训练next-token时输入前T-1项，目标后T-1项，mask padding标签。此代码未包含反向/优化器，不把前向验证当训练验证。若需要encoder-decoder，新增cross-attention，Q来自decoder、K/V来自encoder。
5. 测试同一prefix配不同suffix，prefix logits应相同；独立逐head逐token标量参考核实维度、mask和拼接，不能只看shape。

手动走一遍

B=2,T=4,d=8,H=2时score[2,2,4,4]，每位置只能看到0..i。把第四token改掉，前三位置输出应保持；无位置表示时模型不具备同样的位置信息。

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

复杂度

L层dense前向约 O(L(BTd²+BTdf+BT²d))，词表投影 O(BTdV)；注意力概率 O(BHT²)空间。增量cache需另实现才能避免重复prefill。

易错点与追问

容易失分的边界

• mask方向反了。
• 把forward宣称完整训练。
• 省略位置/残差不解释。
• 只靠维度正确证明实现正确。

继续深挖

• encoder-decoder如何加入cross attention？
• 增量decode与全序列结果怎样对照？