Chapter4.3:Implementing a feed forward network with GELU activations

4 Implementing a GPT model from Scratch To Generate Text

4.3 Implementing a feed forward network with GELU activations

• 本节即将实现子模块，用于transformer block（变换器块）的一部分。为此，我们需要从激活函数开始。

  深度学习中，ReLU因其简单和有效而被广泛使用。但在大语言模型中，还使用了GELU和SwiGLU这两种更复杂、平滑的激活函数，它们结合了高斯和sigmoid门控，提升了模型性能，与ReLU的简单分段线性不同。

• GELU（Hendrycks 和 Gimpel，2016）可以通过多种方式实现；其精确版本定义为 GELU(x)=x⋅Φ(x)，其中 Φ(x) 是标准高斯分布的累积分布函数。

  在实践中，通常会实现一种计算成本更低的近似版本：
  $$\text{GELU}(x)\approx 0.5\cdot x\cdot \left(1+\tanh \left[\sqrt{\frac{2}{\pi }}\cdot \left(x+0.044715\cdot x^3\right)\right]\right)$$
  （原始的 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))
          ))
      
  import matplotlib.pyplot as plt
  
  gelu, relu = GELU(), nn.ReLU()
  
  # Some sample data
  x = torch.linspace(-3, 3, 100)
  y_gelu, y_relu = gelu(x), relu(x)
  
  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} activation function")
      plt.xlabel("x")
      plt.ylabel(f"{label}(x)")
      plt.grid(True)
  
  plt.tight_layout()
  plt.show()
  

  ​

  如上图所示

  1. ReLU：分段线性函数，正输入直接输出，负输入输出零。

     ReLU的局限性： 零处有尖锐拐角，可能增加优化难度；对负输入输出零，限制了负输入神经元的作用。

  2. GELU：平滑非线性函数，近似 ReLU，对负值具有非零梯度（除了在约-0.75处之外）

     GELU的优势：平滑性带来更好的优化特性；对负输入输出较小的非零值，使负输入神经元仍能贡献学习过程。

• 接下来让我们使用 GELU 函数来实现小型神经网络模块 FeedForward，稍后我们将在 LLM 的转换器块中使用它：

  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)
  

  上述的前馈模块是包含两个线性层和一个GELU激活函数的小神经网络，在1.24亿参数的GPT模型中，用于处理嵌入大小为768的令牌批次。

  print(GPT_CONFIG_124M["emb_dim"])
  
  """输出"""
  768
  

  ​

  初始化一个新的前馈网络（FeedForward）模块，其 token 嵌入大小为 768，并向其输入一个包含 2 个样本且每个样本有 3 个 token 的批次输入。我们可以看到，输出张量的形状与输入张量的形状相同：

  ffn = FeedForward(GPT_CONFIG_124M)
  
  # input shape: [batch_size, num_token, emb_size]
  x = torch.rand(2, 3, 768) 
  out = ffn(x)
  print(out.shape)
  
  """输出"""
  torch.Size([2, 3, 768])
  

  本节的前馈模块对模型学习和泛化至关重要。它通过内部扩展嵌入维度到更高空间，如下图所示，然后应用GELU激活，最后收缩回原维度，以探索更丰富的表示空间。

  前馈神经网络中层输出的扩展和收缩的图示。首先，输入值从 768 个值扩大到 4 倍，达到 3072 个值。然后，第二层将 3072 个值压缩回 768 维表示。

• 本节至此，现在已经实现了 下图中LLM 的大部分构建块（打勾的部分）

  ​