Week 8 - Transformer：把 Attention 组装成可训练的网络

Prompt A 最终位置在第二个 LayerNorm 后：
z = [1.2476, 0.2644, -1.5404, 0.0284]

教学 W1 的前四个 Hidden Pre-activations：
[1.2476, 0.2644, -1.5404, 0.0284]

GELU 后：
[1.1152, 0.1596, -0.0952, 0.0144]

再经教学 W2=0.25I：
FFN update = [0.2788, 0.0399, -0.0238, 0.0036]

Prompt A 的最终“喜欢”位置：
residual_0      = [ 0.5000, 0.3000, -0.1500,  0.2000]
attention_update = [ 0.9944, 0.1092, -1.4326, -0.0512]

逐元素相加：
residual_after_attention
= [1.4944, 0.4092, -1.5826, 0.1488]

residual_0                                      [2,2,4]
Attention(LN1(residual_0))                    [2,2,4]
residual_after_attention                      [2,2,4]

FFN(LN2(residual_after_attention))            [2,2,4]
residual_after_ffn                            [2,2,4]

先令 gamma=[1,1,1,1]、beta=[0,0,0,0]，忽略极小 epsilon 的显示误差：

x        = [1, 2, 3, 4]
mean     = (1+2+3+4) / 4 = 2.5
deviation = [-1.5, -0.5, 0.5, 1.5]
variance = (2.25+0.25+0.25+2.25) / 4 = 1.25
std      = sqrt(1.25) ≈ 1.118

normalized
= deviation / std
≈ [-1.342, -0.447, 0.447, 1.342]

Prompt A、Head 1、最终 Query 的两次打分：
q = [1.2206, 0.3715]
k_我 = [-0.2156, -1.2939]
k_喜欢 = [1.2206, 0.3715]

score_我
= (1.2206×-0.2156 + 0.3715×-1.2939) / sqrt(2)
≈ -0.5260

score_喜欢
= (1.2206×1.2206 + 0.3715×0.3715) / sqrt(2)
≈ 1.1511

Softmax([-0.5260, 1.1511])
≈ [0.1575, 0.8425]

Prompt A 的 Values 就是各 Head 选出的 LN1 Channels：

Head 1 output
= 0.1575×[-0.2156,-1.2939] + 0.8425×[1.2206,0.3715]
≈ [0.9944, 0.1092]

Head 2 output
= 0.0349×[1.5095,0.0000] + 0.9651×[-1.5390,-0.0531]
≈ [-1.4326, -0.0512]

Concat 两个 Heads；教学 W_O=I_4：
attention_update ≈ [0.9944,0.1092,-1.4326,-0.0512]

Prompt A final row
attention_update         = [ 0.9944, 0.1092, -1.4326, -0.0512]
residual_0               = [ 0.5000, 0.3000, -0.1500,  0.2000]
residual_after_attention = [ 1.4944, 0.4092, -1.5826,  0.1488]
LN2                      = [ 1.2476, 0.2644, -1.5404,  0.0284]
FFN update               = [ 0.2788, 0.0399, -0.0238,  0.0036]
residual_after_ffn       = [ 1.7732, 0.4492, -1.6064,  0.1524]
final LayerNorm h        = [ 1.3128, 0.2134, -1.4933, -0.0330]

Prompt B final row
attention_update         = [ 0.9847, 0.3949, -1.3141,  0.0686]
residual_0               = [ 0.5000, 0.3000, -0.1500,  0.2000]
residual_after_attention = [ 1.4847, 0.6949, -1.4641,  0.2686]
LN2                      = [ 1.1475, 0.4158, -1.5843,  0.0209]
FFN update               = [ 0.2508, 0.0687, -0.0224,  0.0027]
residual_after_ffn       = [ 1.7356, 0.7637, -1.4865,  0.2713]
final LayerNorm h        = [ 1.2100, 0.3787, -1.5462, -0.0425]

import math

import torch
import torch.nn.functional as F


torch.set_printoptions(precision=4, sci_mode=False)

ids = torch.tensor([
    [0, 1],  # 我 喜欢
    [4, 1],  # 猫 喜欢
])

token_table = torch.tensor([
    [ 0.20, -0.10,  0.70,  0.30],  # 我
    [ 0.60,  0.30, -0.20,  0.10],  # 喜欢
    [-0.40,  0.80,  0.50, -0.30],  # AI
    [ 0.10,  0.20,  0.90,  0.40],  # 学习
    [-0.70,  0.40,  0.30,  0.60],  # 猫
])

position_table = torch.tensor([
    [ 0.05, 0.10, -0.05, 0.00],
    [-0.10, 0.00,  0.05, 0.10],
])

B, T = ids.shape
C, num_heads = 4, 2
head_size = C // num_heads

residual_0 = F.embedding(ids, token_table) + position_table[:T]

# Teaching LayerNorm: gamma=1, beta=0, epsilon=1e-5.
normalized_for_attention = F.layer_norm(
    residual_0,
    normalized_shape=(C,),
    eps=1e-5,
)

# Teaching Q/K/V projections:
# Head 1 selects channels 0:2; Head 2 selects channels 2:4.
heads = normalized_for_attention.view(B, T, num_heads, head_size)
heads = heads.transpose(1, 2)  # [B,H,T,D]
q = heads
k = heads
value_states = heads

scores = (q @ k.transpose(-2, -1)) / math.sqrt(head_size)
causal_mask = torch.tril(torch.ones(T, T, dtype=torch.bool))
masked_scores = scores.masked_fill(~causal_mask, float("-inf"))
attention_probs = F.softmax(masked_scores, dim=-1)

head_outputs = attention_probs @ value_states
attention_update = head_outputs.transpose(1, 2).contiguous().view(B, T, C)
# Teaching W_O is I_4, so the projection leaves attention_update unchanged.

residual_after_attention = residual_0 + attention_update
normalized_for_ffn = F.layer_norm(
    residual_after_attention,
    normalized_shape=(C,),
    eps=1e-5,
)

# Sparse teaching FFN:
# W1 copies four channels into the first four of sixteen hidden units.
hidden = torch.zeros(B, T, 4 * C)
hidden[..., :C] = normalized_for_ffn
hidden = F.gelu(hidden, approximate="none")

# W2 selects those four units and multiplies them by 0.25.
ffn_update = 0.25 * hidden[..., :C]
residual_after_ffn = residual_after_attention + ffn_update

final_hidden = F.layer_norm(
    residual_after_ffn,
    normalized_shape=(C,),
    eps=1e-5,
)

# Reuse Week 6's five candidate scoring rules.
W_out = torch.tensor([
    [1.0, 0.0,  0.0, 0.0],
    [0.0, 1.0,  2.0, 2.0],
    [0.0, 0.0,  1.0, 1.0],
    [0.0, 3.0, -1.0, 0.0],
    [1.0, 2.0,  0.0, 0.0],
])
bias = torch.tensor([-0.2, 0.1, 0.0, 0.0, 0.0])

logits = final_hidden @ W_out.T + bias
vocabulary_probs = F.softmax(logits, dim=-1)

print("final attention probabilities:", attention_probs[:, :, -1, :])
print("attention update:", attention_update[:, -1, :])
print("residual after attention:", residual_after_attention[:, -1, :])
print("FFN update:", ffn_update[:, -1, :])
print("final hidden:", final_hidden[:, -1, :])
print("final logits:", logits[:, -1, :])
print("final vocabulary probabilities:", vocabulary_probs[:, -1, :])

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalSelfAttention(nn.Module):
    def __init__(self, model_dim: int, num_heads: int, dropout: float):
        super().__init__()
        if model_dim % num_heads != 0:
            raise ValueError("model_dim must be evenly divisible by num_heads")
        self.num_heads = num_heads
        self.head_size = model_dim // num_heads
        self.qkv = nn.Linear(model_dim, 3 * model_dim, bias=False)
        self.output = nn.Linear(model_dim, model_dim, bias=False)
        self.attention_dropout = nn.Dropout(dropout)
        self.output_dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, time_steps, model_dim = x.shape
        q, k, value_states = self.qkv(x).chunk(3, dim=-1)
        q = q.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)
        k = k.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)
        value_states = value_states.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_size)
        causal = torch.ones(time_steps, time_steps, device=x.device, dtype=torch.bool).tril()
        scores = scores.masked_fill(~causal, float("-inf"))

        attention_probs = F.softmax(scores, dim=-1)
        dropped_probs = self.attention_dropout(attention_probs)
        heads = dropped_probs @ value_states

        merged = heads.transpose(1, 2).contiguous().view(batch_size, time_steps, model_dim)
        return self.output_dropout(self.output(merged))


class TransformerBlock(nn.Module):
    def __init__(self, model_dim: int, num_heads: int, dropout: float):
        super().__init__()
        self.ln1 = nn.LayerNorm(model_dim)
        self.attention = CausalSelfAttention(model_dim, num_heads, dropout)
        self.ln2 = nn.LayerNorm(model_dim)
        self.ffn = nn.Sequential(
            nn.Linear(model_dim, 4 * model_dim),
            nn.GELU(),
            nn.Linear(4 * model_dim, model_dim),
            nn.Dropout(dropout),
        )

    def forward(self, residual: torch.Tensor) -> torch.Tensor:
        attention_update = self.attention(self.ln1(residual))
        residual_after_attention = residual + attention_update

        ffn_update = self.ffn(self.ln2(residual_after_attention))
        residual_after_ffn = residual_after_attention + ffn_update
        return residual_after_ffn


sample = torch.randn(2, 2, 4)
block = TransformerBlock(model_dim=4, num_heads=2, dropout=0.1)
block.eval()
output = block(sample)
print(output.shape)  # torch.Size([2, 2, 4])

Prompt A final hidden
h_A = [1.3128, 0.2134, -1.4933, -0.0330]

例如“学习”的 Logit：
z_学习 = [0,3,-1,0] · h_A + 0
       = 3(0.2134) - (-1.4933)
       ≈ 2.1336

全部 Logits [我,喜欢,AI,学习,猫]
z_A = [1.1128, -2.7390, -1.5262, 2.1336, 1.7397]

Prompt B final hidden
h_B = [1.2100, 0.3787, -1.5462, -0.0425]
z_B = [1.0100, -2.6987, -1.5887, 2.6821, 1.9674]

ids A/B                                      [2,2]
token embeddings                             [2,2,4]
position embeddings                          [2,4] (broadcast across B)
residual_0 = token + position                 [2,2,4]
one or more Transformer blocks               [2,2,4]
final LayerNorm                              [2,2,4]
LM head                                      [2,2,5]

training: logits.reshape(B*T,V_vocab)        [4,5]
          targets.reshape(B*T)                [4]
          cross entropy                       scalar

generation: logits[:, -1, :]                 [2,5]
            sampled/argmax next_id            [2,1]
            appended ids                       [2,3]

class MiniGPT(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        model_dim: int,
        num_heads: int,
        num_layers: int,
        dropout: float,
    ):
        super().__init__()
        self.context_length = context_length
        self.token_embedding = nn.Embedding(vocab_size, model_dim)
        self.position_embedding = nn.Embedding(context_length, model_dim)
        self.blocks = nn.ModuleList([
            TransformerBlock(model_dim, num_heads, dropout)
            for _ in range(num_layers)
        ])
        self.final_norm = nn.LayerNorm(model_dim)
        self.lm_head = nn.Linear(model_dim, vocab_size)

    def forward(
        self,
        token_ids: torch.Tensor,
        targets: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, T = token_ids.shape
        if T > self.context_length:
            raise ValueError("Sequence length exceeds context_length")

        positions = torch.arange(T, device=token_ids.device)
        residual = (
            self.token_embedding(token_ids)
            + self.position_embedding(positions)
        )

        for block in self.blocks:
            residual = block(residual)

        final_hidden = self.final_norm(residual)
        logits = self.lm_head(final_hidden)

        if targets is None:
            return logits, None

        loss = F.cross_entropy(
            logits.reshape(B * T, logits.size(-1)),
            targets.reshape(B * T),
        )
        return logits, loss


model = MiniGPT(
    vocab_size=5,
    context_length=8,
    model_dim=4,
    num_heads=2,
    num_layers=2,
    dropout=0.1,
)

inputs = torch.tensor([
    [0, 1],  # 我 喜欢
    [4, 1],  # 猫 喜欢
])
targets = torch.tensor([
    [1, 2],  # 喜欢 AI
    [1, 3],  # 喜欢 学习
])

logits, loss = model(inputs, targets)
print(logits.shape)  # torch.Size([2, 2, 5])
print(loss.shape)    # torch.Size([])