前置知識：前面的 micrograd、makemore 系列課程（可選），熟悉 Python，微積分和統計學的基本概念
目標：理解和欣賞 GPT 的工作原理
你可能需要的資料：
Colab Notebook 地址
Twitter 上看到的一份很細緻的筆記，比我寫得好
ChatGPT#
在 2022 年底問世的 ChatGPT 到如今的 GPT4、Claude3，這些 LLM（大語言模型）已經融入了很多人的日常生活。它們都是概率系統，對於同樣的 prompt，它們的答案是多樣的。相較於我們之前實現的語言模型，GPT 等可以能夠模擬單詞、字符或更一般的符號序列，並知道英語中某些單詞是如何相互跟隨的。從這些模型的視角來看，我們的 prompt 就是一段序列的開始，模型要做的就是補全這個序列。
那麼，為這些單詞序列建模的神經網絡是什麼呢？
Transformer#
2017 年，AI 領域裡程碑式的論文《Attention is All You Need》提出了 Transformer 架構。
我們所熟知的 GPT，全稱是 Generative Pre-trained Transformer （生成式預訓練 Transformer）。儘管原文章對應的領域是機器翻譯，但其深遠影響了整個 AI 領域，稍微改動這個架構就可以應用到大量 AI 應用程序中，也是 ChatGPT 的核心。
當然，本節的目標並不是訓練一個 ChatGPT，畢竟那可是個超級工業級項目，涉及大量數據的訓練、預訓練和微調過程，我們要做的只是訓練一個基於 Transformer 的語言模型，跟前面一樣，這也會是一個字符級的語言模型。
搭建模型#
數據集#
使用 toy 級別的小規模數據集 “Tiny Shakespeare”，它深得 Andrej 喜愛。這個數據集基本就是莎士比亞所有作品的大雜燴，文件大小約 1MB。和 ChatGPT 的一個區別是，ChatGPT 的輸出單位是 token，類似 “單詞塊” 的概念，我們後面也會提到。
# 我們總是從一個數據集開始訓練, 下載小莎士比亞數據集
!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
# 閱讀以檢查
with open('input.txt', 'r', encoding='utf-8') as f:
text = f.read()
Tokenize#
chars = sorted(list(set(text))) # set得到序列中不重複的字符的無序序列，轉為list來得到排序功能
vocab_size = len(chars)
print(''.join(chars)) # 合併為一個字符串
print(vocab_size)
# 輸出(按ASCII碼排序)：
# !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz
# 65
這裡和我們在前面幾節中做的字符表功能是一樣的，然後我們需要開發一個能 tokenize (標記化) 輸入序列的功能，這個名字的意義是將原始文本作為一個字符串轉換為一些整數序列。對於我們要做的字符級模型來說，那就只是簡單的把單個字符映射為數字。
如果看過前面幾節內容的話，這部分代碼給你的感覺應該相當熟悉，如 Bigram 中 “創建 Lookup Table 和字符組合的映射關係” 就和這裡很相像。
# 創建一個從字符到整數的映射
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: 接收一個字符串，輸出一個整數列表（編碼器）
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: 接收一個整數列表，輸出一個字符串（解碼器）
print(encode("hii there"))
print(decode(encode("hii there")))
這裡我們同時構建了編碼器和解碼器，這裡的作用就是在字符級別，將字符串和整數相互轉換。這只是個非常簡單的 tokenize 算法，人們設想了很多方法，如谷歌的SentencePiece ，它可以將文本分割成子詞（subwords），這也是實踐中常採用的方法；OpenAI 也有TikToken這個用字節對（byte pairs）來 tokenize。
使用 tiktoken 來編碼：gpt2 的詞彙表中包含了 50,257 個 token，面對同樣的字符串，相較於我們簡單的算法，只用了 3 個整數就完成了編碼。
# 現在對整個文本數據集進行編碼，並將其存儲到一個 torch.tensor 中。
import torch
data = torch.tensor(encode(text), dtype=torch.long)
print(data.shape, data.dtype)
print(data[:1000])
之前看到的 1000 個字符在 GPT 中會呈現為這樣：
可以看到 0 就是空格，1 就是換行符
目前為止，整個數據集就被重新表示為一個龐大的整數序列。
訓練、驗證集劃分#
# 將數據分成訓練集和驗證集，檢查模型的過擬合程度
n = int(0.9*len(data)) # 前90%將用於訓練，其餘部分為驗證
train_data = data[:n]
val_data = data[n:]
我們不希望模型完美記住莎翁的作品，而是讓它創造模仿莎翁風格的文本。
chunks & batches#
要注意的是，我們不會把整個文本一次輸入 Transformer 中，而是使用數據集的 chunks，也就是從訓練集中隨機抽取小塊樣本。
分塊處理#
Block Size 被用來指定模型訓練時每個輸入數據塊（如文本片段）的固定長度。
block_size = 8
train_data[:block_size+1]
x = train_data[:block_size]
y = train_data[1:block_size+1]
for t in range(block_size):
context = x[:t+1]
target = y[t]
print(f"when input is {context} the target: {target}")
這實際上是一種逐步揭露上下文信息給模型的策略
這個方法能強迫模型學會基於先前的字符（或 token）來預測序列中的下一個字符（或 token），提高推理能力。
分批處理#
為了提高 GPU 擅長的並行運算的效率，我們還要考慮分批訓練，將多個批次的文本塊堆疊在一個 tensor 中，同時處理多個之間相互獨立的數據塊。
batch size 的含義就是我們的 Transformer 每一次 forward & backward pass 需要處理多少個獨立序列。
torch.manual_seed(1337) # 提供採樣和可復現能力
batch_size = 4 # 並行處理獨立序列的個數
block_size = 8 # 預測的最大上下文長度
# 這裡的作用類似torch中的dataloader
def get_batch(split):
# 生成一小批輸入數據 x 和目標數據 y
data = train_data if split == 'train' else val_data
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1:i+block_size+1] for i in ix])
return x, y
xb, yb = get_batch('train')
torch.stack用於沿著新的維度對一系列張量進行堆疊（stack），所有的張量需要具有相同的形狀。
可以看到 inputs 的形狀是 4x8，每一列都是訓練集中的一部分；而 targets 的作用是在模型最後計算損失函數。
for b in range(batch_size): # 批次維度
for t in range(block_size): # 時間維度
context = xb[b, :t+1]
target = yb[b,t]
這樣可以更清楚了解 inputs & outpus 兩個數組的關係
Bigram#
在 Makemore 系列中，我們深入了解並實現了 bigram 語言模型，現在改用 PyTorch Module 來快速重新實現。
模型搭建#
import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)
class BigramLanguageModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
# 每個token直接從lookup table中讀取下一個token的對數
self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
這個 embedding layer 想必也很熟悉，比如輸入 24 那就是到 embedding table 中取出第 24 行。
def forward(self, idx, targets=None):
# idx和targets都是整數的（B，T）張量
logits = self.token_embedding_table(idx) # (Batch=4,Time=8,Channel=65)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
在 Makemore 系列中，我們知道了衡量損失的一個好方法：負對數似然損失，在 PyTorch 中對應的實現是 “cross-entropy”（交叉熵）。直觀來講，就是模型對於 logit 對應的正確分類應該有一個很高的概率（信心高），同時其他所有維度都是很低的概率（信心極低）。此時的損失是可以估計的，大概是 -log (1/65)，約等於 4.17，但因為一些熵的存在實際結果會更大一些。
# 從模型中生成
def generate(self, idx, max_new_tokens):
# idx 是當前上下文中索引的 (B, T) 陣列
for _ in range(max_new_tokens):
# 獲取預測結果
logits, loss = self(idx)
# 僅關注最後一個time step
logits = logits[:, -1, :] # becomes (B, C)
# 應用softmax以獲取概率
probs = F.softmax(logits, dim=-1) # (B, C)
# 從分布中抽樣
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# 將採樣索引添加到運行序列中
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
print(loss)
print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist())) # 序列的第一個字符為換行符(0)
generate 的任務就是把大小為 BxT 的代表上下文信息的 idx 擴展為 $B\times T + 1, + 2 ,+\ldots$，也就是在時間維度中的所有批次維度上繼續生成 。
模型未訓練時的生成結果，完全是隨機的
模型訓練#
下面開始訓練這個模型。相較於 Makemore 系列使用隨機梯度下降 (SGD)，這裡我們使用的是更先進和流行的 AdamW 優化器，
# 創建一個 PyTorch 優化器
optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)
優化器基本作用就是獲取梯度，並根據梯度更新參數。
batch_size = 32 # 選擇更大的 batch-size
for steps in range(100): # 增加步數以獲得良好結果
# 采樣一個 batch 的數據
xb, yb = get_batch('train')
# 評估損失
logits, loss = m(xb, yb)
optimizer.zero_grad(set_to_none=True) # 將上一步的梯度清零
loss.backward() # 反向傳播
optimizer.step() # 相當於“新參數=舊參數-學習率*梯度”，跟我們以前手動實現的梯度下降循環基本一樣
print(loss.item())
可以看到我們的優化起作用了，損失在減小
增加訓練輪數，最終達到了約 2.48，我們複製之前的採樣代碼片段再次生成，應該能得到有進步的結果。
初具人形，但又沒那麼具。
再怎麼訓練也很難得到理想的結構，因為這個模型假設十分簡單（只根據前一個預測後一個 token），各個 token 之間沒有聯繫，這就是使用 transformer 的原因。
Transformer#
如果有 Nvidia 顯卡的話可以加速訓練：
device = 'cuda' if torch.cuda.is_available() else 'cpu'
這麼設定後需要對代碼進行一些改動，大概就是要讓數據加載、計算和採樣生成都在 device (GPU) 上進行，具體細節請見 Andrej 的 lecture 倉庫，其中的 bigram.py 是我們的起點。
除此之外，我們的模型分為了訓練階段和評估階段，不過目前模型中只有一個nn.Embedding 層而兩個階段表現一樣，並沒有引入 dropout layer、batch norm layer 等。這樣做是訓練模型中的最佳實踐 ( best practice )，因為一些層在訓練和推理階段的行為不同。
Self-Attention#
自注意力機制
在上手 Transformer 之前，我們要做的第一件事是通過簡單例子來習慣一個實現 Transformer 內部自注意力核心的數學技巧。
torch.manual_seed(1337)
B,T,C = 4,8,2 # batch, time, channels
x = torch.randn(B,T,C)
# x.shape = torch.Size([4, 8, 2])
我們希望通過一個特定的方式將這些本來獨立的 token 結合起來，比如說，第五個 token 不應該能與第 6、7 和 8 號 token 交互通信，因為在這個序列中，這三個屬於未來的future tokens。因此，5 號 token 只能與第 4、3、2 和 1 號 token 交互通信，也就是說信息只能從之前的上下文流向當前的 time step，根據這些信息來預測未來信息。
那 token 之間相互通信的最簡單的方式是什麼呢？
答案很令人意外：對前面的 token 取平均，變為當前背景下的歷史特徵向量。當然，也很容易猜到這種交互方式太弱了，丟失了大量關於這些 token 空間排列的信息。
v1. 循環#
# 我們希望有 x[b,t] = mean_{i<=t} x[b,i]
xbow = torch.zeros((B,T,C))
for b in range(B):
for t in range(T):
xprev = x[b,:t+1] # (t,c)
xbow[b,t] = torch.mean(xprev, 0)
打印出來就能理解這段代碼的作用了
這個方法運算效率較低，我們可以使用矩陣乘法來更高效地完成這個工作。
v2. 矩陣乘法#
# 一個簡化示例，說明矩陣乘法如何用於“加權聚合”。
torch.manual_seed(42)
a = torch.ones(3, 3)
b = torch.randint(0,10,(3,2)).float()
c = a @ b
這裡就是基本的矩陣乘法運算，c 中的每個數是 a 和 b 中分別對應的行與列的點積。如 c 的 (1, 1) 元素的就是 。
為了實現一樣的效果，可以將 a 換為下三角矩陣：
這樣可以實現 “分別抽取第一、二行加和” 的效果，通過torch.trill實現
a = torch.tril(torch.ones(3, 3))
現在的作用是加和，因為 a 中的元素都為 1，為了實現加權聚合，可以對 A 的每行進行歸一化，使每行元素和為 1：
a = torch.tril(torch.ones(3, 3))
a = a / torch.sum(a, 1, keepdim=True) # keepdim保證廣播語義可行
現在 a 的每行和為 1，c 就是 b 中的對應前幾行的平均了
回到前面來應用這個更有效的方法：
# 版本2：使用矩陣乘法進行加權聚合
wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)
weight 權重矩陣對應上面的 a 矩陣
# 這裡torch會給weight創造一個batch維度
xbow2 = wei @ x # (T, T) @ (B, T, C) ----> (B, T, C)
torch.allclose(xbow, xbow2) # 比較兩個張量是否在一定的數值容差範圍內相等
# 輸出 True,兩種方法效果一樣
總結一下這個竅門：我們可以用批矩陣乘法來實現加權聚合，權重在這個 T✖️T 矩陣中指定，而加權和是根據維度和權重呈倒三角分布，使得 t 維上的 token 只能從先前的 token 中獲取信息。
v3. Softmax#
除此之外，可以使用 Softmax 來實現第三個版本。
其中一個重要的 api 是 torch.masked_fill()，根據指定的掩碼 ( mask ) tensor 對輸入 tensor 進行填充，如下圖所示：
那對於每一行都取 Softmax 會發生什麼呢？在前面的章節中提過，Softmax 是一個歸一化操作，這裡的作用就是對每行中 “過去的” 元素使用下三角矩陣乘法進行加權聚合：
# 版本3：使用Softmax
tril = torch.tril(torch.ones(T, T))
wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
xbow3 = wei @ x
除了編碼 token 的身份，還會對 token 的位置進行編碼：
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每個token直接從查找表中讀取下一個token的logits
self.token_embedding_table = nn.Embedding(vocab_size, n_embd) # token編碼
self.position_embedding_table = nn.Embedding(block_size, n_embd) # 位置編碼
def forward(self, idx, targets=None):
B, T = idx.shape
# idx和targets都是整數的(B,T)張量
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
這裡用 x 保存 token 嵌入和位置嵌入的加和，但由於目前只是個簡單的 bigram 模型，不同位置是具有平移不變性的，但當我們使用注意力機制時就不一樣了。
v4. self-attention#
我們當前所做的只是簡單的平均，但現實中，每個 token 的意義都不是不同的，也就是數據相關的。比如對於一個元音，它就想知道在自己前面傳遞信息的輔音們是什麼。這就是自注意力機制解決的問題。
每個 token 都將發出兩個向量，一個 query（查詢：我要找什麼，對啥感興趣）和一個key（鍵：我包含什麼信息，和誰相似）。
我們得到序列中這些 token 之間親和力（權重）的方法，基本上就是在鍵和查詢之間做點積。如果一個鍵與查詢非常匹配或對齊，這個鍵對應的權重會很高。因此，模型的注意力會集中在這個鍵對應的值上 —— 即模型會更加關注這個特定的信息（或 token ），而不是序列中的其他任何信息。
也就是說，注意力機制通過計算查詢和鍵之間的匹配度並據此分配權重，使模型能夠聚焦於最相關的信息，從而提高了處理和理解序列數據的能力。
我們還需要一個 Value（值：你對我感興趣的話，我會貢獻給你的信息），最終聚合的是經過一個線性層傳播的 x 而不是直接聚合 x。
現在我們來實現這個單頭注意力機制:
# 版本 4: self-attention!
torch.manual_seed(1337)
B,T,C = 4,8,32 # batch, time, channels
x = torch.randn(B,T,C)
# 讓我們看看單頭自注意力是如何運作的
head_size = 16
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)
# 現在key和query對x進行前向傳播
k = key(x) # (B, T, 16)
q = query(x) # (B, T, 16)
# 點積得到親和力（權重）
wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)
tril = torch.tril(torch.ones(T, T))
# wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
v = value(x)
out = wei @ v
# out = wei @ x
可以看到權重矩陣 wei 中，並非像以前一樣完全平均了，而是數據依賴的：親和力高的 token 會在加權聚合中為當前 token 提供更多信息
注意力總結#
-
注意力是一種通信機制。可以看作有向圖中的節點相互觀察，並通過指向它們的所有節點的加權和來聚合信息，權重是數據相關的。
-
本身並沒有空間概念。注意力只是作用在一組向量上。這就是為什麼我們需要對標記進行位置編碼。
-
每個批次維度上的示例完全獨立處理，永遠不會互相通信。
-
注意力機制並不在乎你是否只關心過去的信息。我們這裡的實現中，當前 token 對於未來的信息是屏蔽的，但只要在 “編碼器” 注意力模塊中刪除使用
masked_fill進行掩碼的代碼就能允許所有 token 互相通信。這裡稱為 “解碼器” 模塊，因為它具有三角矩陣掩碼。注意力機制支持任意節點之間的連接。 -
"自注意力" 意味著鍵和值是從與查詢相同的源（x）產生的。在 " 交叉注意力（cross-attention）" 中，查詢仍然由 x 生成，但鍵和值來自其他某些外部來源（例如一個編碼器模塊）
-
"Scaled" Attention：額外將
wei除以 $\sqrt {head_size}$ 。
原因是，我們當前只有標準高斯分布 （均值為 0，方差為 1）輸入時，會發現簡單的加權得到的 wei 的方差其實是 head size 的數量級（我們的實現中正是 16），
加上這個 normalization 後，權重的方差就是 1 了：
為啥這一步很重要呢？
# 我們的 wei 會經過 softmax
wei = F.softmax(wei, dim=-1)
對於 softmax 來說，一個性質是經過 softmax 後分布中絕對值較大的元素會向 1 趨近：
可以看到，同樣的分布在乘上 8 後，變得趨於 one-hot 分布了（one-hot：一個元素是 1 其它全 0），即在初始階段導致分布過於尖銳，基本上只是從單個節點獲取信息。
因此，“縮放” 注意力機制說白了就是控制初始化時的方差，通過將 wei 除以 1/√(頭大小) 來進行額外調整。這樣，當輸入 Q（查詢）和 K（鍵）的方差為 1 時，wei也將保持單位方差，這意味著 Softmax 將保持分散，而不會過度飽和。使得在應用 Softmax 函數之前，權重分布不會因為過大的值而導致梯度消失或梯度爆炸，進而保持了模型的穩定性和效果。
代碼實現#
class Head(nn.Module):
""" one head of self-attention """
def __init__(self, head_size):
super().__init__() # 人們一般不在這裡用 bias
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B,T,C = x.shape
k = self.key(x) # (B,T,C)
q = self.query(x) # (B,T,C)
# 計算注意力分數 ("affinities")，這裡應用了上面提到的 scaled
wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
# 執行加權聚合
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
構造函數中， tril 並非 nn.Module 的參數，這在 PyTorch 命名約定中被稱為緩衝區（buffer），要調用的話必須使用 register_buffer來將它賦值給 nn.Module，
Multi-Head Attention#
論文中還有這部分沒有復現，多頭注意力就是並行地應用多個注意力，並將它們的結果連接起來：
代碼實現#
在 PyTorch 中，我們可以簡單地通過創建多個 head 來做到這一點。
class MultiHeadAttention(nn.Module):
""" 自注意力中的多頭並行。 """
def __init__(self, num_heads, head_size):
super().__init__()
# 在一個列表中並行運行，然後將輸出連接起來
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
def forward(self, x):
return torch.cat([h(x) for h in self.heads], dim=-1) # 在 chanmel 維度上連接
現在我們有了 4 個並行的通信 channel 而不是一個，單個通道會相應地變小。嵌入維度是 32，對應的就有 8 維度的 self-attention，連接起來又得到 32 個，這就是原始嵌入。這裡有點類似於 群卷積 ( group convolution )，不進行一個大卷積而是分組卷積。
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每個 token 直接從查找表中讀取下一個 token 的 logit
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.sa_head = nn.MultiHeadAttention(4, n_embd//4)
self.lm_head = nn.Linear(n_embd, vocab_size)
Blocks#
對於下圖，也就是論文所展示的網絡結構，我們不會實現對編碼器的交叉注意力。但圖中還有一個前饋部分，它被分組成一個 block，一個不斷重複（$N\times$）的塊。
前饋網絡#
這個前饋部分只是個簡單的 MLP：
注意論文這裡說輸入和輸出的維數是 512，前饋的內層維數是 2048，所以前饋的內層通道大小應該乘以 4
class FeedFoward(nn.Module):
""" 一個簡單的線性層，後面跟著一個非線性函數 """
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
""" Transformer block: 分散了通信和計算
通信：多頭注意力
計算：前饋網絡在所有 token 上獨立完成
"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我們想要的 head 個數，這裡是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
#
self.ffwd = FeedFoward(n_embd) # 計算
def forward(self, x):
x = self.sa(x)
x = self.ffwd(x)
return x
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每個 token 直接從查找表中讀取下一個 token 的 logit
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
)
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx 和 targets 都是 (B,T) 大小的整數 tensor
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
...
嘗試解碼（decode），發現結果並沒有改善多少。原因是我們現在得到了相當深的神經網絡，會受到優化問題影響，我們還需要從 Transformer 的論文中借鑒一個解決這個問題的方法。
現在有兩種能極大地提高網絡深度的同時，確保網絡保持可優化狀態的方法：
殘差連接#
block 中紅色圈起的部分（箭頭和 Add）就是殘差連接（residual connection），這個概念由 Deep Residual Learning for Image Recognition論文提出。
Andrej 的原話是 “你轉換數據，然後與先前特徵進行跳連接並相加”，我們這裡詳細解釋一下（可以配合圖片食用）：
-
你轉換數據（You transform data）：深度神經網絡中的每一層，輸入數據會通過權重矩陣的乘法和激活函數的非線性變換等操作進行 “轉換”，從而學習到數據的抽象表示。
-
但然後你有一個跳躍連接（But then you have a skip connection）：在傳統的深度網絡中，數據的這種轉換是連續且線性的。殘餘連接通過引入 “跳過連接” 打破了這種模式。跳過連接直接將某一層的輸入連接到後面的某一層（通常是相隔一層或幾層的輸出），這樣做的目的是為了將前面層的特徵直接傳遞給後面的層。
-
與之前的特徵連接並相加（With addition from the previous features）：跳躍連接的實現通常是通過將跳過的輸入和目標層的輸出進行元素級的加法操作來完成的。這種加法操作保證了原始特徵信息能夠在網絡中直接傳遞，而不會被後續層的變換所 “稀釋”。
殘餘連接的引入使得網絡能夠更容易地學習到恆等映射（identity mapping），這對於訓練非常深的網絡是非常有用的。實際上，殘餘連接允許梯度直接流過網絡，有助於緩解梯度消失或梯度爆炸的問題，從而使得訓練深層網絡變得更加可行和高效。
我們在 micrograd 中提到過，神經網絡中加法節點的作用就是將梯度均勻分配給所有輸入（因為加法操作對於每個輸入來說是線性的，並且每個輸入對輸出的貢獻是獨立的） 看到這裡的時候不禁感慨：鮮活的知識點居然又串起來了。
class MultiHeadAttention(nn.Module):
""" 自注意力中的多頭並行。 """
def __init__(self, num_heads, head_size):
super().__init__()
# 在一個列表中並行運行，然後將輸出連接起來
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(n_embd, n_embd) # 引入投射
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1) # 在 chanmel 維度上連接
out = self.proj(out) # 殘差路徑的投影就是 out 的線性變換
return out
class FeedFoward(nn.Module):
""" 一個簡單的線性層，後面跟著一個非線性函數 """
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(n_embd, 4 * n_embd), # 投射回殘差通路
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
""" Transformer block: 分散了通信和計算
通信：多頭注意力
計算：前饋網絡在所有 token 上獨立完成
"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我們想要的 head 個數，這裡是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
self.ffwd = FeedFoward(n_embd) # 計算
def forward(self, x):
x = x + self.sa(x)
x = x + self.ffwd(x)
return x
Layer Norm#
第二個優化方法是這裡的 Norm，指的是一種叫做層歸一化（Layer Norm）的東西：
Layer Norm 和我們之前實現過的 Batch Norm （確保在 batch 維度上，任意神經元都是標準正態分布）非常相似，唯一與 Batch Norm 不同的是，Layer Norm 不是在批次維度上進行歸一化，而是在特徵維度上進行歸一化。這意味著，對於網絡中的每個樣本，Layer Norm 會計算該樣本所有特徵的均值和標準差，並用這些統計量來歸一化該樣本的所有特徵。
這裡需要注意的是，Transformer 論文發布到現在內部的細節並沒有太多變動，但我們這裡的實現和原論文略有出入，可以看到論文中 Add&Norm 是在 Transform 後添加的，但現在更普遍的做法是在 Transform 之前應用 Layer Norm，這被稱為 Pre-norm 公式。
現在我們已經擁有了一個相當完整的 Transformer 了（只有解碼器 decoder）
class Block(nn.Module):
""" Transformer block: 分散了通信和計算
通信：多頭注意力
計算：前饋網絡在所有 token 上獨立完成
"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我們想要的 head 個數，這裡是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
self.ffwd = FeedFoward(n_embd) # 計算
self.ln1 = nn.LayerNorm(embd)
self.ln2 = nn.LayerNorm(embd)
def forward(self, x):
x = self.sa(self.ln1(x))
x = self.ffwd(self.ln2(x))
return x
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每個 token 直接從查找表中讀取下一個 token 的 logit
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
''' 這部分等價於下面的兩行
self.blocks = nn.Sequential(
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
nn.LayerNorm(n_embd),
)
'''
self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)]) # 擴大模型，指定 block 層數
self.ln_f = nn.LayerNorm(n_embd) # 最終的 Layer Norm
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx 和 targets 都是 (B,T) 大小的整數 tensor
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
Dropout#
Dropout 是可以在殘差連接回到殘差路徑之前添加的東西，由Dropout: A Simple Way to Prevent Neural Networks from Overfitting論文提出，它基本上讓你的神經網絡在每次 forward-backward pass 隨機關閉一些神經元（也就是降為 0，不參加後面的訓練）
這裡只需要知道它是一種正則化技術就可以了
# 超參數（可以用 Colab V100 跑，CPU的話太慢了，或者調低超參數）
batch_size = 64
block_size = 256 # 增加上下文長度，預測第 257 個 token
max_iters = 5000
eval_interval = 100
learning_rate = 1e-3 # 網絡變大，降低學習率
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6 # 384 / 6
n_layer = 4 # 4 層 Block
dropout = 0.2 # 每次 forward-backward pass 都有 20% 機率丟棄神經元
class Head(nn.Module):
""" one head of self-attention """
def __init__(self, head_size):
super().__init__() # 人們一般不在這裡用 bias
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x):
B,T,C = x.shape
k = self.key(x) # (B,T,C)
q = self.query(x) # (B,T,C)
# 計算注意力分數 ("affinities")，這裡應用了上面提到的 scaled
wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
wei = self.dropout(wei) # Dropout
# 執行加權聚合
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
class MultiHeadAttention(nn.Module):
def __init__(self, num_heads, head_size):
super().__init__()
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(n_embd, n_embd)
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
class FeedFoward(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout), # Dropout
)
def forward(self, x):
return self.net(x)
現在的結果就像是莎士比亞風格的胡言亂語，不過顯然很可觀了。
圖中左半邊所示的編碼器（encoder）和右邊紅圈中的交叉注意力並沒有在這裡實現。
我們只用了解碼器的原因是，我們只是生成不受任何條件限制的文本，就像最後的結果一樣，只是根據給定的莎士比亞數據集胡言亂語，我們通過三角掩碼來實現注意力機制，使其具有自回歸的性質便於抽樣進行語言模型建模。
而原始論文採用 encoder-decoder 的結構是因為它是在機器翻譯領域的，模型期待的輸入是外文編碼 token（如法語），然後解碼翻譯為英語，如下圖所示：
這裡編碼器就是取感興趣的法語句子來創建 token ， 在上面使用 Transformer 結構但不用三角掩碼，讓所有的 tokens 都盡可能多地相互通信。負責語言建模的解碼器會連接編碼完成後的輸出（上面架構圖中的左半部分的最上面），這是通過交叉注意力完成的。實際所做的就是對解碼進行制約，不僅僅是在當前解碼過去信息，更是在完整的、完全編碼的法語 tokens 上進行。而我們實現的是 decoder-only 的版本。
本節知識對應的項目：karpathy/nanoGPT，同樣是只關注預訓練部分的實現。
重回 ChatGPT#
要訓練 ChatGPT 大致有兩個階段：預訓練和微調階段。
預訓練#
在大量互聯網語料上訓練，試圖得到一個 encoder-only 的 Transformer。我們現在已經完成了一個很小的預訓練步驟。還有一點不同是， OpenAI 的訓練使用的是 tokenizer，也就是詞彙表不是單個字符而是字符塊，我們使用的莎士比亞數據集對應的大概會有 30 萬 token，我們在上面訓練了約 1000 萬個參數，而 GPT3 的 Transformer 最多擁有 1750 億個參數，在 3000 億個 token 上進行了訓練。
在完成這一步後，你並不能向模型提問，因為它目前只會創建互聯網上的新聞等信息，也就是只有補全序列的作用。
微調#
這個階段就是把它調教為一個語言模型助手。
第一步是收集成千上萬的文檔，其格式為 “問題：答案”，對模型進行微調對齊（align），在對大模型中微調過程中的樣本效率是很高的。
第二步，評分者根據模型的回應進行排名，用此來訓練一個獎勵模型（reward model）。
第三步，運行 PPO（ Proximal Policy Optimization，一種策略梯度強化學習優化器）來微調抽樣策略，把模型從一個文檔補全器變成了一個問題回答器。
當然，這些部分個人是基本無法復刻的，大公司才行。
關於 GPT 的詳細講述，Andrej 在 2023 年 3 月份的 Microsoft Build 演講中進行了全面的講述，詳見GPT 的現狀。