7.2 Attention Is All You Need：Transformer 架構全解（2017）#

import torch, torch.nn as nn, math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.h = n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)   # Q,K,V 合一投影
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)       # (B, T, D) each
        # 拆頭：(B, h, T, d_k)
        q = q.view(B, T, self.h, self.d_k).transpose(1, 2)
        k = k.view(B, T, self.h, self.d_k).transpose(1, 2)
        v = v.view(B, T, self.h, self.d_k).transpose(1, 2)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)   # (B,h,T,T)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = torch.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, D)
        return self.out(out)                          # (B, T, D)

class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, p=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(),
                                 nn.Linear(d_ff, d_model))
        self.n1, self.n2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.drop = nn.Dropout(p)

    def forward(self, x, mask=None):
        x = self.n1(x + self.drop(self.attn(x, mask)))   # 殘差 + LN
        x = self.n2(x + self.drop(self.ffn(x)))
        return x

# 測試：並行性 —— 全序列一次前向，無時間迴圈
layer = EncoderLayer()
x = torch.randn(4, 50, 512)               # batch=4, 句長=50
print(layer(x).shape)                     # (4, 50, 512)，50 步同時處理！

import torch, math

def positional_encoding(T, d_model):
    pos = torch.arange(T).unsqueeze(1)
    i   = torch.arange(0, d_model, 2)
    div = torch.exp(-i * math.log(10000) / d_model)
    pe = torch.zeros(T, d_model)
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe

pe = positional_encoding(100, 64)
# 性質 1：相對位置的線性變換性（內積只依賴相對距離）
dot = pe @ pe.T
print("任意兩行內積僅依賴 pos_i - pos_j:",
      torch.allclose(dot[0, 5], dot[10, 15], atol=1e-4))   # True
# 性質 2：置換等變性 —— 沒有位置編碼的注意力不知道順序
x = torch.randn(1, 5, 8)
att = nn.MultiheadAttention(8, 2, batch_first=True)
out1, _ = att(x, x, x)
perm = torch.randperm(5)
out2, _ = att(x[:, perm], x[:, perm], x[:, perm])
print("打亂順序後輸出同樣打亂（置換等變）:",
      torch.allclose(out1[:, perm], out2, atol=1e-5))      # True

import torch, torch.nn as nn, math

class TinyTransformer(nn.Module):
    def __init__(self, src_v, tgt_v, d=128, h=4, layers=2):
        super().__init__()
        self.d = d
        self.se, self.te = nn.Embedding(src_v, d), nn.Embedding(tgt_v, d)
        enc = [EncoderLayer(d, h, 4*d) for _ in range(layers)]
        self.enc = nn.ModuleList(enc)
        self.dec_attn = nn.ModuleList(MultiHeadAttention(d, h) for _ in range(layers))
        self.dec_ffn  = nn.ModuleList(nn.Sequential(nn.Linear(d, 4*d), nn.ReLU(),
                                                    nn.Linear(4*d, d)) for _ in range(layers))
        self.n = nn.ModuleList(nn.LayerNorm(d) for _ in range(4*layers))
        self.out = nn.Linear(d, tgt_v)

    def causal_mask(self, T):
        return torch.tril(torch.ones(T, T)).bool()

    def forward(self, src, tgt):
        x = self.se(src) * math.sqrt(self.d) + positional_encoding(src.size(1), self.d)
        for i, l in enumerate(self.enc):                 # 編碼器（無掩碼）
            x = self.n[4*i](x + l.attn(x)); x = self.n[4*i+1](x + l.ffn(x))
        y = self.te(tgt) * math.sqrt(self.d) + positional_encoding(tgt.size(1), self.d)
        m = self.causal_mask(tgt.size(1)).to(y.device)
        for i in range(len(self.enc)):                   # 解碼器（掩碼 + 交叉）
            y = self.n[4*i+2](y + self.dec_attn[i](y, m.unsqueeze(0).unsqueeze(0)))
            cross = MultiHeadAttention.attend(self.dec_attn[i], x, x, y)  # Q=y, KV=x
            y = y + cross
            y = self.n[4*i+3](y + self.dec_ffn[i](y))
        return self.out(y)

# MultiHeadAttention 需支援 attend(Q,K,V)；完整實作見本書配套代碼庫
# 訓練流程與 6.2 的 Seq2Seq 相同（teacher forcing + 交叉熵），
# 但現在整條序列一個前向 —— GPU 平行度拉滿