11อ่านโค้ดจริงทั้งชั้น

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

torch.manual_seed(1337)          # ตัวเลขทุกตัวในบทนี้มาจาก seed นี้

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.n_head, self.n_embd = n_head, n_embd
        self.attn = nn.Linear(n_embd, 3 * n_embd, bias=False)   # Q, K, V ในก้อนเดียว
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.attn(x).split(C, dim=2)                  # แยกสามทางที่นี่
        hs = C // self.n_head
        q = q.view(B, T, self.n_head, hs).transpose(1, 2)       # แตกเป็นหลายหัว
        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
        v = v.view(B, T, self.n_head, hs).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * hs ** -0.5            # คะแนนความเข้ากัน
        att = att.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = att @ v                                             # ถ่วงน้ำหนักแล้วบวก
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y), att

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.sa = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),                      # ขยาย 4 เท่า
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),                      # บีบกลับ
        )

    def forward(self, x):
        delta, att = self.sa(self.ln1(x))
        x = x + delta                                           # residual เส้นที่ 1 — สื่อสาร
        x = x + self.mlp(self.ln2(x))                           # residual เส้นที่ 2 — คิด
        return x, att

B, T, n_embd, n_head = 1, 5, 8, 2
block = Block(n_embd, n_head, block_size=T)
x = torch.randn(B, T, n_embd)
y, att = block(x)

print("params in one block =", sum(p.numel() for p in block.parameters()))
print("x.shape =", tuple(x.shape), " -> y.shape =", tuple(y.shape))
print("\nattention of head 0 (rows = who is asking):")
print(att[0, 0].detach().numpy().round(2))

with torch.no_grad():                                  # ปิดส่วนแก้ทั้งสองเส้น
    for p in block.sa.proj.parameters():
        p.zero_()
    for p in block.mlp[-1].parameters():
        p.zero_()
    y_zero, _ = block(x)
print("\nblock becomes identity?", torch.allclose(y_zero, x, atol=1e-6))

params in one block = 840
x.shape = (1, 5, 8)  -> y.shape = (1, 5, 8)

attention of head 0 (rows = who is asking):
[[1.   0.   0.   0.   0.  ]
 [0.41 0.59 0.   0.   0.  ]
 [0.31 0.35 0.34 0.   0.  ]
 [0.22 0.28 0.22 0.28 0.  ]
 [0.34 0.19 0.21 0.15 0.11]]

block becomes identity? True