GPT模型層標準化實作詳解

層標準化是 Transformer 架構中不可或缺的技術，用於穩定訓練過程並提升模型效能。不同於批次標準化，層標準化針對每個樣本的特徵維度進行標準化，使其均值接近零，方差接近一，有效解決了內部共變數偏移問題。此特性使其在處理變動批次大小或長序列資料時更具優勢，尤其適用於 GPT 等大語言模型的訓練。本文將進一步探討如何將層標準化整合至前饋網路，並說明捷徑連線如何最佳化深度網路訓練。

實作GPT模型：層標準化（Layer Normalization）詳解

在深度學習模型中，尤其是像GPT這樣的Transformer架構，層標準化是一種重要的技術，用於穩定和加速模型的訓練過程。本章節將探討層標準化的原理，並透過實作範例展示如何在PyTorch中實作層標準化。

層標準化的作用

層標準化是一種正規化技術，主要目的是將神經網路層的輸出進行標準化，使其具有零均值和單位方差。這種技術可以有效減少內部共變數偏移（Internal Covariate Shift），進而提高模型的訓練速度和穩定性。

層標準化的實作範例

首先，我們來建立一個簡單的神經網路層，並觀察其輸出：

import torch
import torch.nn as nn

torch.manual_seed(123)
batch_example = torch.randn(2, 5)
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)

輸出結果為：

tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
        [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
       grad_fn=<ReluBackward0>)

內容解密：

1. 輸入資料：我們首先建立了一個隨機的輸入張量 batch_example，其形狀為 (2, 5)，代表兩個具有五個特徵的訓練樣本。
2. 神經網路層：我們定義了一個簡單的神經網路層 layer，包含一個線性層 nn.Linear(5, 6) 和一個ReLU啟用函式 nn.ReLU()。這個層將輸入的五個特徵對映到六個輸出。
3. 輸出結果：經過神經網路層處理後，我們得到了輸出張量 out，其形狀為 (2, 6)。由於ReLU啟用函式的特性，輸出值均為非負數。

計算均值和方差

接下來，我們計算輸出張量的均值和方差：

mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

輸出結果為：

Mean:
 tensor([[0.1324],
        [0.2170]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.0231],
        [0.0398]], grad_fn=<VarBackward0>)

內容解密：

1. 均值計算：我們使用 out.mean(dim=-1, keepdim=True) 計算輸出張量在最後一個維度（即特徵維度）上的均值。keepdim=True 保證了輸出的均值張量保持與原始輸出張量相同的維度數。
2. 方差計算：同樣地，我們使用 out.var(dim=-1, keepdim=True) 計算輸出張量在最後一個維度上的方差。

應用層標準化

現在，我們將層標準化應用於輸出張量：

out_norm = (out - mean) / torch.sqrt(var + 1e-5) # 加上1e-5避免除以零
mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Normalized layer outputs:\n", out_norm)
print("Mean:\n", mean)
print("Variance:\n", var)

輸出結果為：

Normalized layer outputs:
 tensor([[ 0.6159,  1.4126, -0.8719,  0.5872, -0.8719, -0.8719],
        [-0.0189,  0.1121, -1.0876,  1.5173,  0.5647, -1.0876]],
       grad_fn=<DivBackward0>)
Mean:
 tensor([[0.0000],
        [0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

內容解密：

1. 層標準化公式：我們使用公式 (out - mean) / torch.sqrt(var + 1e-5) 對輸出張量進行標準化。這裡加上的 1e-5 是為了避免除以零的情況。
2. 標準化結果：經過層標準化後，輸出的均值接近零，方差接近一，達到了標準化的效果。

4.2 正規化啟用函式：層正規化（Layer Normalization）

到目前為止，我們已經逐步實作並應用了層正規化。現在，讓我們將這個過程封裝在一個 PyTorch 模組中，以便稍後在 GPT 模型中使用。

層正規化的實作

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

內容解密：

1. 層正規化的作用維度：此實作的層正規化操作於輸入張量 x 的最後一個維度，即嵌入維度（emb_dim）。
2. eps 的作用：eps 是一個小的常數（epsilon），新增到變異數中以防止在正規化過程中除以零。
3. scale 和 shift：這兩個引數是可訓練的，它們的維度與輸入相同。在訓練過程中，如果模型認為調整它們能夠提升在訓練任務上的表現，那麼它們就會被自動調整，以學習到最適合處理資料的縮放和偏移。

測試層正規化

ln = LayerNorm(emb_dim=5)
out_ln = ln(batch_example)
mean = out_ln.mean(dim=-1, keepdim=True)
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)
print("平均值:\n", mean)
print("變異數:\n", var)

輸出結果表明，層正規化程式碼按預期工作，將每個輸入的值正規化，使其平均值為 0，變異數為 1。

層正規化 vs. 批次正規化

層正規化與批次正規化不同，後者是沿著批次維度進行正規化，而層正規化是沿著特徵維度進行正規化。大語言模型（LLMs）通常需要大量的計算資源，並且在訓練或推斷期間，批次大小可能會受到硬體或特定使用案例的限制。由於層正規化獨立於批次大小對每個輸入進行正規化，因此它在這些場景中提供了更大的靈活性和穩定性。這對於分散式訓練或在資源有限的環境中佈署模型尤其有益。

4.3 實作帶有 GELU 啟用函式的前饋網路

接下來，我們將實作一個小型神經網路子模組，該子模組用作 LLMs 中轉換器區塊的一部分。首先，我們將實作 GELU 啟用函式，它在此神經網路子模組中扮演著至關重要的角色。

GELU 啟用函式

歷史上，ReLU 啟用函式因其簡單性和在各種神經網路架構中的有效性而被廣泛使用。然而，在 LLMs 中，除了傳統的 ReLU 外，還採用了其他幾種啟用函式。兩個值得注意的例子是 GELU（高斯誤差線性單元）和 SwiGLU（Swish-gated 線性單元）。

GELU 和 SwiGLU 是更複雜和平滑的啟用函式，分別結合了高斯和 sigmoid-gated 線性單元。它們為深度學習模型提供了更好的效能，與更簡單的 ReLU 不同。

GELU 啟用函式有多種實作方式；確切的版本定義為 GELU(x) = x⋅Φ(x)，其中 Φ(x) 是標準高斯分佈的累積分佈函式。然而，在實踐中，通常會實作一個計算成本較低的近似值（原始的 GPT-2 模型也是使用這種近似值進行訓練的，這是透過曲線擬合找到的）。

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))

內容解密：

1. GELU 的近似實作：此程式碼實作了 GELU 啟用函式的近似版本，使用 torch.tanh 和其他運算來近似累積分佈函式 Φ(x)。
2. 運算過程：首先計算 x 的三次方並乘以一個常數，然後加上 x 本身，接著乘以另一個根據 π 的常數，最後透過 torch.tanh 函式處理，並根據結果計算最終輸出。
3. 為何使用近似：使用近似版本是因為它計算成本較低，並且已被證明在實際應用中足夠準確，例如在 GPT-2 模型的訓練中。

實作帶有GELU啟用函式的前饋網路

在前一節中，我們探討了GELU（Gaussian Error Linear Unit）啟用函式的特性。現在，我們將利用GELU來實作一個小型神經網路模組——FeedForward，這個模組將在稍後的LLM（大語言模型）變壓器區塊中使用。

GELU與ReLU的比較

首先，讓我們透過繪製GELU和ReLU函式的圖形來直觀地比較這兩種啟用函式：

import matplotlib.pyplot as plt
import torch
import torch.nn as nn

# 定義GELU和ReLU啟用函式
gelu, relu = GELU(), nn.ReLU()

# 生成-3到3之間的100個資料點
x = torch.linspace(-3, 3, 100)

# 計算GELU和ReLU的輸出
y_gelu, y_relu = gelu(x), relu(x)

# 繪製GELU和ReLU的圖形
plt.figure(figsize=(8, 3))
for i, (y, label) in enumerate(zip([y_gelu, y_relu], ["GELU", "ReLU"]), 1):
    plt.subplot(1, 2, i)
    plt.plot(x, y)
    plt.title(f"{label}啟用函式")
    plt.xlabel("x")
    plt.ylabel(f"{label}(x)")
    plt.grid(True)
plt.tight_layout()
plt.show()

內容解密：

1. 匯入必要的函式庫：使用matplotlib進行繪圖，torch和torch.nn進行神經網路相關的操作。
2. 定義啟用函式：例項化GELU和ReLU啟用函式。
3. 生成資料點：使用torch.linspace生成-3到3之間的100個資料點，用於評估啟用函式。
4. 計算輸出：將資料點輸入到GELU和ReLU中，計算相應的輸出。
5. 繪圖：使用matplotlib繪製GELU和ReLU的圖形，直觀比較兩者的差異。

從圖4.8中可以看到，ReLU是一個分段線性函式，當輸入為正時直接輸出輸入值，否則輸出零。相比之下，GELU是一個平滑的非線性函式，它近似於ReLU，但對於大多數負值都有非零梯度。GELU的平滑性質使得在訓練過程中能夠進行更細微的引數調整，從而可能帶來更好的最佳化效果。

實作FeedForward模組

接下來，我們將使用GELU啟用函式來實作FeedForward模組：

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

    def forward(self, x):
        return self.layers(x)

內容解密：

1. 定義FeedForward類別：繼承自nn.Module，這是PyTorch中定義神經網路模組的基本類別。
2. __init__方法：初始化FeedForward模組，接受一個組態字典cfg作為引數。
   • 使用nn.Sequential定義一個順序容器，按照新增的順序執行其中的模組。
   • 第一個線性層(nn.Linear)將輸入維度從emb_dim擴充套件到4 * emb_dim。
   • 使用GELU啟用函式進行非線性轉換。
   • 第二個線性層將維度從4 * emb_dim壓縮回emb_dim。
3. forward方法：定義前向傳播的過程，直接呼叫self.layers對輸入x進行處理。

測試FeedForward模組

讓我們初始化一個FeedForward模組，並測試其對隨機輸入的處理：

ffn = FeedForward(GPT_CONFIG_124M)
x = torch.rand(2, 3, 768)
out = ffn(x)
print(out.shape)

內容解密：

1. 初始化FeedForward模組：使用GPT_CONFIG_124M組態初始化一個FeedForward例項。
2. 生成隨機輸入：建立一個形狀為(2, 3, 768)的隨機張量，模擬批次大小為2、序列長度為3、嵌入維度為768的輸入。
3. 前向傳播：將輸入傳遞給ffn，執行前向傳播。
4. 輸出形狀：列印輸出張量的形狀。

輸出結果表明，輸入和輸出的形狀保持一致，都是torch.Size([2, 3, 768])。這是因為FeedForward模組內部雖然對嵌入維度進行了擴充套件和壓縮，但最終輸出的維度與輸入保持一致，這樣的設計便於在模型中堆積疊多個這樣的層，而無需擔心維度的比對問題。

4.4 加入捷徑連線（Shortcut Connections）

如同圖 4.11 所示，我們已經實作了大語言模型（LLM）的大部分基本元件。接下來，我們將討論在神經網路的不同層之間加入捷徑連線（shortcut connections）的概念，這對於改善深度神經網路架構的訓練效能至關重要。

捷徑連線的概念

讓我們來探討捷徑連線（也稱為跳躍或殘差連線）的背後概念。捷徑連線最初是在電腦視覺領域（特別是在殘差網路中）為瞭解決梯度消失（vanishing gradients）的問題而提出的。梯度消失是指在訓練過程中，梯度在向後傳播時變得越來越小，使得早期層的訓練變得困難。

圖 4.12 顯示，捷徑連線透過將一層的輸出新增到後續層的輸出中，建立了一條替代的、更短的路徑，讓梯度能夠流經網路。這就是為什麼這些連線也被稱為跳躍連線的原因。它們在訓練過程中對於保持梯度的流動起著至關重要的作用。

實作捷徑連線

在下面的列表中，我們實作了圖 4.12 中的神經網路，以展示如何在 forward 方法中新增捷徑連線。

class ExampleDeepNeuralNetwork(nn.Module):
    def __init__(self, layer_sizes, use_shortcut):
        super().__init__()
        self.use_shortcut = use_shortcut
        self.layers = nn.ModuleList([
            nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]),
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]),
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]),
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]),
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]),
                          GELU())
        ])

    def forward(self, x):
        for layer in self.layers:
            layer_output = layer(x)
            if self.use_shortcut and x.shape == layer_output.shape:
                x = x + layer_output
            else:
                x = layer_output
        return x

內容解密：

1. 類別定義：ExampleDeepNeuralNetwork 繼承自 nn.Module，定義了一個具有五層的神經網路，每層由一個線性層和一個 GELU 啟用函式組成。
2. __init__ 方法：初始化神經網路的層，並根據 use_shortcut 引數決定是否使用捷徑連線。
3. forward 方法：定義了前向傳播的過程。如果啟用了捷徑連線，並且輸入和輸出的形狀相同，則將輸入新增到輸出中，否則直接使用輸出作為下一層的輸入。
4. 梯度計算：透過 print_gradients 函式計算模型的梯度，以展示捷徑連線對梯度消失問題的影響。

初始化模型並計算梯度

layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])
torch.manual_seed(123)
model_without_shortcut = ExampleDeepNeuralNetwork(layer_sizes, use_shortcut=False)

def print_gradients(model, x):
    output = model(x)
    target = torch.tensor([[0.]])
    loss = nn.MSELoss()
    loss = loss(output, target)
    loss.backward()
    # 列印梯度

內容解密：

1. 模型初始化：建立了一個不使用捷徑連線的神經網路模型。
2. print_gradients 函式：計算模型的輸出、損失，並進行反向傳播以計算梯度。
3. 梯度消失問題：透過比較有無捷徑連線的模型的梯度，可以觀察到捷徑連線如何幫助維持較大的梯度值，從而改善訓練效果。