神经网络中的梯度下降和反向传播：CUDA/C++

梯度下降与反向传播：使用CUDA/C++从零开始训练神经网络

本文是深入探索神经网络世界的第四部分，我们将逐步构建Transformer架构，从CUDA中最基本的组件开始。在本期中，我们将从实现前向传播转向模型训练的关键阶段。我们将深入探讨梯度下降和反向传播的基本机制，使用C++和CUDA在底层实现它们，以训练我们的神经网络识别MNIST数据集中的手写数字。我们的目标不仅仅是使用现成的库，更是要彻底理解这些复杂的计算过程是如何“在幕后”运作的，这对于深入理解现代机器学习和人工智能模型至关重要。

神经网络训练基础：梯度下降

在上一部分成功实现前向传播之后，我们的模型已经能够进行预测。然而，在初始阶段，未经训练的模型，其预测结果将是随机的。为了使模型开始产生正确的结果，必须调整其内部参数（权重和偏置），以最小化预测与真实标签之间的误差。梯度下降和反向传播等算法正是为此目的而生。

想象一下，你身处浓雾弥漫的山顶，任务是下到山谷的最低点。你无法看清整个地形，但能感受到脚下的坡度。梯度是一个向量，它指示了最陡峭的上升方向。因此，要下山，你需要沿着梯度的反方向移动。在神经网络的语境中，“山”是误差函数（或损失函数）的曲面，“山谷”则是其全局最小值。梯度下降的目标就是通过迭代地调整网络权重，使其沿着损失函数相对于这些权重的梯度的反方向移动，从而找到这个最小值。

反向传播是一种高效计算这些梯度的算法。它的工作原理是将误差从网络的输出层反向传播到输入层，将误差的责任分配给所有权重。在数学上，这通过微分的链式法则实现，它使我们能够计算每个权重的变化如何影响模型的整体误差。

损失函数：交叉熵

为了衡量模型预测的“好坏”，我们使用损失函数。在分类任务中，例如识别手写MNIST数字，交叉熵损失是一个标准选择。它量化了模型预测的概率分布与真实分布（其中只有正确类别具有1的概率，其他为0）之间的差异。这种差异越大，损失函数的值就越高。

对于单个样本，交叉熵使用以下公式计算：

![L = -\sum_{i=1}^{K} y_i \log(\hat{y}_i)](./images/image-1.svg)

其中：

![K](./images/image-2.svg) — 是类别数量（对于MNIST，为10）。
![y_i](./images/image-3.svg) — 是真实值（正确类别为1，其他为0）。
![\hat{y}_i](./images/image-4.svg) — 是模型的预测（Softmax层之后获得的概率）。

由于训练通常在数据批次上进行，而非单个样本，因此一个批次的总损失函数是其大小（![N](./images/image-5.svg) — 批次大小）的平均值：

![L = -\frac{1}{N} \sum_{j=1}^{N} \sum_{i=1}^{K} y_{j,i} \log(\hat{y}_{j,i})](./images/image-6.svg)

以下是用于在CPU上计算交叉熵的C++代码，这简化了调试和理解：

float cross_entropy_loss_cpu(
  const Tensor& predictions, const std::vector<int>& labels
) {
  int batch_size = predictions.shape()[0];
  int num_classes = predictions.shape()[1];
  float loss = 0.0f;
  for (int i = 0; i < batch_size; i++) {
    for (int j = 0; j < num_classes; j++) {
      float prob = predictions.get({i, j});
      float target = (labels[i] == j) ? 1.0f : 0.0f;
      loss += -std::log(prob) * target;
    }
  }
  return loss / batch_size;
}

反向传播与链式法则

为了最小化损失函数![L](./images/image-9.svg)，我们需要知道其值如何随网络中每个权重的微小变化而变化。这就是损失函数相对于权重的梯度。例如，对于第一层的权重![W_1](./images/image-10.svg)，我们需要找到![\frac{\partial L}{\partial W_1}](./images/image-8.svg)。

我们在之前的文章中开发的网络架构如下：

第一线性层： ![z_1 = W_1 \cdot x + b_1](./images/image-11.svg)
激活（ReLU）： ![a_1 = \sigma(z_1)](./images/image-12.svg)
第二线性层： ![z_2 = W_2 \cdot a_1 + b_2](./images/image-13.svg)
输出（Softmax + 损失）： ![L = Loss(Softmax(z_2))](./images/image-14.svg)

为了计算梯度![\frac{\partial L}{\partial W_1}](./images/image-16.svg)，我们应用链式法则：

![\frac{\partial L}{\partial W_1} = \frac{\partial L}{\partial z_2} \cdot \frac{\partial z_2}{\partial a_1} \cdot \frac{\partial a_1}{\partial z_1} \cdot \frac{\partial z_1}{\partial W_1}](./images/image-17.svg)

这个链中的每一项都代表一个局部梯度，它显示了给定层或函数的输出如何影响其输入。反向传播首先计算损失函数相对于最后一层输出的梯度，然后将此梯度依次向后传播通过网络的所有层。

Softmax和交叉熵梯度的详细推导

让我们从链的末端开始——损失函数![L](./images/image-25.svg)相对于应用Softmax后第二线性层输出![z_2](./images/image-20.svg)的导数。这是最重要且最常用的梯度之一。回想一下Softmax在向量![z_2](./images/image-20.svg)中每个元素![i](./images/image-19.svg)上的形式：

![a_{2, i} = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}](./images/image-21.svg)

其中![z_2](./images/image-22.svg)是第二线性层的输出（logits），![a_2](./images/image-23.svg)是Softmax的结果（概率）。梯度![\frac{\partial L}{\partial z_{2,i}}](./images/image-18.svg)可以使用链式法则分解为：

![\frac{\partial L}{\partial z_{2,i}} = \sum_{k=1}^{K} \frac{\partial L}{\partial a_{2,k}} \cdot \frac{\partial a_{2, k}}{\partial z_{2,i}}](./images/image-24.svg)

首先，让我们找到损失函数![L](./images/image-27.svg)相对于![a_i](./images/image-26.svg)的导数：

![\frac{\partial L}{\partial a_i} = \frac{\partial}{\partial a_i} \left( -y_i \ln(a_i) \right) = -y_i \cdot \frac{1}{a_i} = -\frac{y_i}{a_i}](./images/image-29.svg)

现在我们转向Softmax相对于![z](./images/image-30.svg)的导数。这里需要用到微分的商法则。Softmax公式：![a_i = \frac{e^{z_i}}{\sum_{k=1}^K e^{z_k}}](./images/image-32.svg)

如果![i=j](./images/image-43.svg)，导数为：

![\frac{\partial a_i}{\partial z_i} = a_i (1 - a_i)](./images/image-36.svg)

如果![i \neq j](./images/image-37.svg)，导数为：

![\frac{\partial a_i}{\partial z_j} = -a_i a_j](./images/image-41.svg)

使用克罗内克δ符号![\delta_{ij}](./images/image-42.svg)，这两种情况可以合并为：

![\frac{\partial a_i}{\partial z_j} = a_i (\delta_{ij} - a_j)](./images/image-44.svg)

结合交叉熵和Softmax的导数，我们得到：

![\frac{\partial L}{\partial z_i} = \sum_{k=1}^{K} \left( -\frac{y_k}{a_k} \right) \cdot \left( a_k(\delta_{ki} - a_i) \right) = \sum_{k=1}^{K} -y_k (\delta_{ki} - a_i)](./images/image-49.svg)

进一步简化，考虑到真实概率之和![\sum_{k=1}^{K} y_k](./images/image-51.svg)始终等于1，我们得到了一个出人意料的简洁优雅的公式：

![\frac{\partial L}{\partial z_i} = a_i - y_i](./images/image-52.svg)

这意味着损失函数相对于logits（Softmax之前的输出）的梯度，简单来说就是预测概率与真实标签之间的差异。这个结果是分类模型中实现反向传播的基石。

在C++中实现梯度

应用推导出的公式，我们可以实现神经网络最后一层的梯度计算。然后，这个梯度将用于更新权重，使模型朝着更小的误差方向移动。以下是C++代码的实现：

auto grad = Tensor::create_zeros({ batch_size, num_classes });

for (int i = 0; i < batch_size; ++i) {
  for (int j = 0; j < num_classes; ++j) {
    float prob = fc2.output->get({i, j});
    float target = (targets[i] == j) ? 1.0f : 0.0f;
    float gradient = (prob - target) / batch_size; // Division by batch_size for averaging
    grad->set({i, j}, gradient);
  }
}

这段代码计算批次中每个元素的误差梯度，然后将其平均到批次大小。生成的grad张量包含有关如何改变logits ![z_2](./images/image-54.svg)以减少误差的信息。然后，这个梯度将通过网络的前置层进行传播，从而可以使用相应的局部导数调整权重![W_2](./images/image-13.svg)、![b_2](./images/image-13.svg)、![W_1](./images/image-11.svg)和![b_1](./images/image-11.svg)。