От MNIST к Transformer. Часть 4. Gradient Descent. Обучаем нашу модель / Хабр

Мы живем в эпоху, когда ИИ стал доступен каждому. Но за магией PyTorch скрывается колоссальная инженерная работа и сложные вычислительные процессы, которые для большинства остаются черным ящиком.

Это четвертая статья из цикла От MNIST к Transformer, цель которого пошагово пройти путь от простого CUDA ядра до создания архитектуры Transformer - фундамента современных LLM моделей. Мы не будем использовать готовые высокоуровневые библиотеки. Мы будем разбирать, как все устроено под капотом, и пересобирать их ключевые механизмы своими руками на самом низком уровне. Только так можно по настоящему понять как работают LLM и что за этим стоит. В этой статье мы разберем как работает градиентный спуск, реализуем его и обучим нашу модель для распознования mnist датасета.

Приготовьтесь, будет много кода на C++ и CUDA, работы с памятью и погружения в архитектуру GPU. И конечно же математика что за этим стоит. Поехали!

О цикле "От MNIST к Transformer"

Это будет долгий, сложный, но невероятно увлекательный путь. Мы начнем с самых азов и будем шаг за шагом усложнять задачу. Наша цель не только досконально понять как работают современные модели, но и научится программировать на GPU. Статьи для цикла будут выходить постепенно.

Самое начало. Разберемся с основами работы GPU. Напишем наш первый код на C++ и CUDA для простейшей операции сложения векторов, hello world в мире GPGPU и CUDA, чтобы понять, как CPU и видеокарта общаются друг с другом. Статья вот тут.
Основы работы с памятью. Разберемся с основами работы и устройства памяти в GPU. Напишем четыре ядра для двух операций (простой подход и оптимизированный), построение гистрограммы для вектора и транспонирование матрицы. Операция транспонирования нам понадобится в будущем что бы собирать наш Transformer. Статья вот тут .
Умножение тензоров. Пишем Linear Layer. Перейдем от матриц к тензорам, напишем умножение простое умножение тензоров и потом сделаем tiled вариант c broadcasting. Создадим полносвязный слой нейронной сети и напишем прямой проход для нашей модели которая будет распозновать mnist дата сет. Так же не забываем про CUDA: немного углубимся в способы синхронизации CPU и GPU, в реализации mnist модели мы их использовать не будем для простоты, но они пригодятся в дальнейшем. Статья вот тут .
Gradient Descent. Обучаем нашу модель. (Мы здесь) Наконец-то мы можем обучить нашу модель из прошлой статьи. В этой статье разберем как работает градиентный спуск, один из основных алгоритмов обучения нейросетей. Реализуем его на практике и обучим нашу первую модель расопзнованию рукописных чисел.

Gradient Discent и Backpropagation

В прошлой статье мы написали Forward Pass - прогнали картинку через слои и получили какой-то ответ. Скорее всего, на нетренированной модели этот ответ будет случайным. Чтобы модель начала угадывать цифры, нам нужно менять веса так, чтобы ошибка (Loss) становилась меньше. Для того что бы это сделать существует такие алгоритмы как Gradient Discent (градиентный спуск) и Backpropagation (обратное распространение ошибки).

Что же такое градиентный спуск общими словами. Давайте представим, что мы стоим на вершине горы (global maximum), и нам нужно спуститься в самую глубокую яму (global minimum). Но мы не видим где же находится эта ямы мы только можем понимать наклон поверхности. Для того что бы вычислить наклон нам как раз и нужен градиент. Он просто показывает, в какую сторону поверхность идет вверх. А раз нам нужно вниз, мы идем в противоположную сторону. А вся поверхность на которой мы стоим (и гора и яма) это поверхность функции ошибки. То есть глобальная задача такого алгоритма это найти минимум функции ошибки.

Backpropagation (обратное распространение ошибки) - это процесс как бы передачи ошибки от выхода сети к её входу. Если на этапе Forward мы получили ошибку, нам нужно понять, как каждый конкретный вес на неё повлиял. Математически это реализуется через правило цепочки (Chain Rule): мы вычисляем локальную производную текущего слоя и умножаем её на накопленный градиент, пришедший от слоев, стоящих ближе к выходу.

Функция ошибки

Кросс-энтропия (Cross-Entropy Loss) - это стандартный способ измерить, насколько «плохо» наша модель предсказывает классы. В случае с MNIST это способ понять, насколько сильно нейросеть ошиблась, назвав рукописную «пятерку» «тройкой».

Представим, что на выходе нейросети после слоя Softmax у нас есть вероятности. Например, для картинки с цифрой 3 модель выдала: Цифра 2: 10%, Цифра 3: 70% (правильный ответ) и тд.

Формула Кросс-энтропии для одного примера представлена ниже:

$L = -\sum_{i=1}^{K} y_i \log(\hat{y}_i)$

- количество классов (в случае MNIST это 10).
- истинное значение (1 для правильного класса, 0 для остальных).
$\hat{y}_i$ - предсказание модели (вероятность из Softmax).

Так как часто мы на вход для обучения мы даем не один пример, а целую пачку или batch, то для батча мы должны поделить значение на размера батча ( - batch size):

$L = -\frac{1}{N} \sum_{j=1}^{N} \sum_{i=1}^{K} y_{j,i} \log(\hat{y}_{j,i})$

Код для вычисления кросс-энтропии, опять же будем считать на CPU для упрощения:

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;
}

Chain rule

Так мы определились с функцией ошибки и для того что бы найти минимум функции нам нужно найти наклон функции в точке в которой мы находимся для этого нам нужно найти градиент функции ошибки по отношению к весам сети.

$\frac{\partial L}{\partial W_1}$

Для того что бы найти градиент нам нужно взять производную от функции ошибки по весам. В данной статье я не буду подробно рассматривать что такое производная функции. Если совсем просто, производная - это скорость изменения чего-либо (функции в нашем случае). Для того что бы найти производную от функции ошибки по отношению к давайте вспомним архитекутру сети из прошлой статьи:

Первый слой:
Активация (ReLU):
Второй слой:
Выход (Softmax + Loss):

Исходя из архитектуры сети производную от функции ошибки по отношению к можно найти по правилу цепочки (Chain rule)

$\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}$

Для начала давайте рассмотрим первую часть производной $\partial L / \partial z_2$ , для этого так же вспомним как выглядит функция softmax, для каждого элемента в векторе :

$a_{2, i} = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$

: выход второго линейного слоя (вектор логитов).
: результат Softmax (вектор вероятностей).

$\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}}$

Давайте по очереди возьмем производные, сначала возьмем первую часть по :

$L = -\sum_{k=1}^{K} y_k \ln(a_k)$ $\frac{\partial L}{\partial a_i} = \frac{\partial}{\partial a_i} \left( -y_i \ln(a_i) \right)$ $\frac{\partial L}{\partial a_i} = -y_i \cdot \frac{1}{a_i} = -\frac{y_i}{a_i}$

Теперь давайте возьмем производную второй части, а именно Softmax по , тут чуть сложнее. Для того чтобы найти производную Softmax, нам понадобится правило производной частного:

$\frac{d}{dx} \left( \frac{u(x)}{v(x)} \right) = \frac{u'(x)v(x) - u(x)v'(x)}{[v(x)]^2}$

Формула softmax:

$a_i = \frac{e^{z_i}}{\sum_{k=1}^K e^{z_k}}$ $\frac{\partial u}{\partial z_i} = e^{z_i}$ $\frac{\partial v}{\partial z_i} = \frac{\partial}{\partial z_i} (e^{z_1} + e^{z_2} + \dots + e^{z_i} + \dots) = e^{z_i}$ $\frac{\partial a_i}{\partial z_i} = \frac{e^{z_i} (\sum e^{z_k}) - e^{z_i} (e^{z_i})}{(\sum e^{z_k})^2}$

Упростим, разбив на две дроби

$\frac{\partial a_i}{\partial z_i} = \frac{e^{z_i}}{\sum e^{z_k}} \cdot \frac{\sum e^{z_k} - e^{z_i}}{\sum e^{z_k}} = a_i (1 - a_i)$

Рассмотрим так же случай когда $i \neq j$ (внедиагональные элементы):

$\frac{\partial u}{\partial z_j} = 0 \quad \text{(так как } e^{z_i} \text{ не зависит от } z_j)$ $\frac{\partial v}{\partial z_j} = e^{z_j}$

Итого для второго случая:

$\frac{\partial a_i}{\partial z_j} = \frac{0 \cdot (\sum e^{z_k}) - e^{z_i} (e^{z_j})}{(\sum e^{z_k})^2}$ $\frac{\partial a_i}{\partial z_j} = - \frac{e^{z_i}}{\sum e^{z_k}} \cdot \frac{e^{z_j}}{\sum e^{z_k}} = -a_i a_j$

Чтобы не писать два случая, математики используют символ Кронекера $δ_{ij}$ (он равен 1, если , и 0 в остальных случаях):

$\frac{\partial a_i}{\partial z_j} = a_i (\delta_{ij} - a_j)$

Давайте теперь обьеденим 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)$ $\frac{\partial L}{\partial z_i} = \sum_{k=1}^{K} -y_k (\delta_{ki} - a_i)$ $\frac{\partial L}{\partial z_i} = -y_i + a_i \left( \sum_{k=1}^{K} y_k \right)$

Так как сумма всех (истинных вероятностей) всегда равна 1, получаем финальную, максимально простую формулу:

$\frac{\partial L}{\partial z_i} = a_i - y_i$

Вот так это считается в коде:

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;
    grad->set({i, j}, gradient);
  }
}

Все это было только для градиентов результата (Softmax + Loss), давайте пойдем дальше но дальше будет чуть легче. Теперь мы подходим к моменту, когда ошибку нужно пробросить от выхода второго слоя () назад ко входу (), который прошел через активацию ReLU.

$z_2=W_2⋅a_1+b_2 ; \frac{\partial z_2}{\partial a_1} = W_2$

Далее берем $\partial a_1 / \partial z_1$

$\text{ReLU}'(z_1) = \begin{cases} 1 & \text{if } z_1 > 0 \\ 0 & \text{if } z_1 \leq 0 \end{cases}$

И последний шаг $\partial z_1 / \partial W_1$ по аналогии $\partial a_1 / \partial z_1$

$z_1=W_1⋅x+b_1 ; \frac{\partial z_1}{\partial W_1} = x$

Итого давайте соединим все вместе:

$\frac{\partial L}{\partial W_1} = \left( \left( W_2^T (a_2 - y) \right) \odot \mathbb{1}(z_1 > 0) \right) x^T$

Теперь давайте для каждого линейного слоя напишем код для вычисления градиента при обратном прохождении.

void LinearLayer::backward(Tensor& input, Tensor& grad, float learning_rate) {
    // Calculate grad_weights = input ^ T * grad;
    MatmulOperation::matmul_inplace(input, grad, *grad_weights, true, false);

    // Calculate grad_bias = sum(grad) over batch
    TensorOperation::sum_rows_inplace(grad, *grad_bias);

    // Calculate grad_input = grad * weights^T (for the next layer back)
    MatmulOperation::matmul_inplace(grad, *weights, *grad_input, false, true);

    // Update parameters (SGD step)
    TensorOperation::add_weighted_inplace(*weights, *grad_weights, learning_rate);
    TensorOperation::add_weighted_inplace(*bias, *grad_bias, learning_rate);
}

Давайте еще немного разберем вот эту строчку кода

  // Calculate grad_bias = sum(grad) over batch
  TensorOperation::sum_rows_inplace(grad, *grad_bias);

Для этого вспомним формулу для одного нейрона: Чему равна производнаяпо ? Правильно она равна 1. Итого мы имеем:

$\frac{\partial L}{\partial b} = \frac{\partial L}{\partial z} \cdot \frac{\partial z}{\partial b} = \frac{\partial L}{\partial z} \cdot 1$

Или для одних входных данных градиент по смещению - это и есть сам градиент ошибки на этом слое. Ядро CUDA для вычисления sum_rows_inplace выглядит так:

__global__ void sumrow_kernel(const float* __restrict__ A, float* __restrict__ B, int size, int batch_size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < size) {
        B[idx] = 0.0f;
    }

    for (int i = 0; i < batch_size; i ++) {
        if (idx < size) {
            B[idx] += A[idx + i * size];
        }
    }
}

Gradient Discent

Давайте еще посмотрим на эти две строчки ниже в функции обратного подхода для одного слоя.

TensorOperation::add_weighted_inplace(*weights, *grad_weights, learning_rate);
TensorOperation::add_weighted_inplace(*bias, *grad_bias, learning_rate);

Эти строчки по сути финал всей работы алгоритма градиентного спуска. Когда мы вычислили градиенты (нашли направление, куда нужно двигаться), нам нужно обновить веса. learning_rate) — коэффициент скорости обучения. Общая форму GCD выглядит так:

$W_{t+1} = W_t - \eta \nabla L(W_t)$

Давайте теперь посмотрим на код обратного прохода для всей сети.

void MNISTModel::backward(Tensor& input, std::vector<int>& targets, float lr) {
    int batch_size = fc2.output->shape()[0];
    int num_classes = fc2.output->shape()[1];

    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;
            grad->set({i, j}, gradient);
        }
        std::cout << std::endl;
    }

    grad->prepare_for_gpu_work();

    fc2.backward(*fc1.output, *grad, lr);

    // Backprop through ReLU
    // The gradient of ReLU is 1 if x > 0, else 0.
    // fc2.grad_input now contains the gradient w.r.t. fc1's output
    TensorOperation::relu_backward_inplace(*fc2.grad_input, *fc1.output);
    fc1.backward(input, *fc2.grad_input, lr);
}

И один этап обучения выглядит таким образом, сначала делаем прямой проход, считаем loss потом делаем обратный проход и обновляем веса.

float MNISTModel::sgd_step(Tensor& input, std::vector<int>& targets, float lr) {
    forward(input);
    backward(input, targets, lr);
    float loss = TensorOperation::cross_entropy_loss_cpu(*fc2.output, targets);
    return loss;
}

Обучение сети

После того как мы разобрались с градиентами сети обратным и прямым проходом при обучении сети, давайте теперь напишем код для всего обучения.

void mnist_train() {
    std::cout << "MNIST Train" << std::endl;
    int batch_size = 60;
    auto train_data = MnistReader::load_train(batch_size);
    train_data.prepare_for_gpu_work();
    std::cout << "Train data: " << train_data.images.size() << std::endl;
    std::cout << "Train data labels: " << train_data.labels.size() << std::endl;
    print_vector("Train data images", train_data.images[0]->shape());
    MNISTModel model(batch_size);
    model.prepare_for_gpu_work();

    int epochs = 10;
    int steps = train_data.images.size();
    for (int epoch = 0; epoch < epochs; epoch++) {
        std::cout << "Epoch " << epoch << std::endl;
        for (size_t i = 0; i < steps; i++) {
            auto& input = train_data.images[i];
            std::vector<int> targets;
            for (size_t j = 0; j < input->shape()[0]; j++) {
                targets.push_back(train_data.labels[i * batch_size + j]);
            }
            float loss = model.sgd_step(*input, targets, 0.001f);
            std::cout << "Epoch: " << epoch << " Loss: " << loss << std::endl;
        }
    }

    auto test_data = MnistReader::load_test(batch_size);
    test_data.prepare_for_gpu_work();
    std::cout << "Test data: " << test_data.images.size() << std::endl;
    std::cout << "Test data labels: " << test_data.labels.size() << std::endl;
    print_vector("Test data images", test_data.images[0]->shape());
    auto predictions = model.predict(*test_data.images[0]);
    print_vector("Predictions", predictions);
    int label_count = test_data.labels.size();
    int correct_count = 0;
    for (int i = 0; i < batch_size; i++) {
        if (predictions[i] == static_cast<int>(test_data.labels[i])) {
            correct_count++;
        }
    }
    std::cout << "Correct count: " << correct_count << std::endl;
    std::cout << "Accuracy: " << (float)correct_count / batch_size << std::endl;
    int label = static_cast<int>(test_data.labels[0]);
    std::cout << "Label: " << label << std::endl;
}

Для обучения мы возьмем батч размером в 60 изображений и пройдем 10 эпох обучения. В начале обучения Loss у нас будет принимать значение около 2.3 что объясняется тем что в начала обучения модель просто угадывает число с вероятностью 10% или 0.1 если мы подставим это число в функцию ошибки то как раз получим примерно 2.3. В конце обучения значение loss опустится до значений примерно 0.3 - 0.4.

Loss ~ 2.3: Рандомное угадывание (10% точности).
Loss ~ 1.0: Модель начала что-то понимать, точность около 50–60%.
~0.3: Хороший результат для простой полносвязной сети (90%+ точности).

Если мы потом возьмем примеры из тестовой выборки и проверим Accuracy то получим примерно 93%. Что может означать что наша модель научилась распознавать числа и MNIST датасета.

Adam Optimizer

Еще пару слов стоит сказать про другие оптимизаторы. Ранее мы рассмотрели просто Gradient Discent и это просто шаг в сторону спуска, Adam Optimizer (Adaptive Moment Estimation) - это улучшенная версия обычного градиентного спуска, который учитывает инерцию и подстраивает скорость для каждого веса отдельно. Современные LLM и Трансформеры почти никогда не обучаются на обычном SGD. В данной статье мы рассмотрели обычный оптимизатор для понимания основ. Давайте немного опишем теорию Adam оптимизатора.

Нам понадобятся две дополнительные переменные для каждого веса (еще два вспомогательных тензора такого же размера, как веса слоя):

(первый момент) — среднее значение градиента (инерция).
(второй момент) — среднее значение квадрата градиента (адаптивность).

Эти два тензора так же обновляются в процессе обучения.

$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$ $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$

Где - это текущий градиент, а (обычно 0.9) и (обычно 0.999) - коэффициенты затухания.

Обновление же весов сети происходит по формуле:

$w_{t+1} = w_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

Где - крошечное число (например $10^{-8}$ ), чтобы не делить на ноль.

На этом пока остановимся и не будем реализовывать Adam Optimizer в текущей статье. В целом доработки к обычному оптимизатору не сложные.

Заключение

На этом пока все. Мы разобрали очень важную тему, а именно как работает алгоритм обучения нашей нейронной сети. Так же полноценно обучили нашу полносвязную двухслойную нейронную сеть. И cамое главное заложили большой фундамент к построению трансформера. В следующих статьях мы начнем уже по блокам строить наш трансформер что является основой всех современных LLM.

Надеюсь было понятно и самое главное интересно. Продолжение следует ...

Подробный код для статьи можно посмотреть на github. Так же если хотите следить и данная тема вам интересна можете подписаться на канал в телеграмме созданный специально для этого цикла статей, там будут анонсы статьей по циклу, обсуждения, ваши запросы и многое другое.