Descenso de gradiente y retropropagación en redes neuronales: CUDA/C++

Descenso de Gradiente y Retropropagación: Entrenando Redes Neuronales desde Cero con CUDA/C++

Este artículo es la cuarta parte de una inmersión profunda en el mundo de las redes neuronales, donde estamos construyendo una arquitectura Transformer paso a paso, comenzando con los componentes más básicos en CUDA. En esta entrega, pasamos de la implementación del paso hacia adelante a la etapa crucial del entrenamiento del modelo. Exploraremos los mecanismos fundamentales del Descenso de Gradiente y la Retropropagación, implementándolos a bajo nivel usando C++ y CUDA para entrenar nuestra red neuronal para reconocer dígitos escritos a mano del conjunto de datos MNIST. El objetivo no es solo usar bibliotecas preexistentes, sino comprender a fondo cómo funcionan estos complejos procesos computacionales "en las entrañas", lo cual es fundamental para una comprensión profunda de los modelos modernos de aprendizaje automático e inteligencia artificial.

Fundamentos del Entrenamiento de Redes Neuronales: Descenso de Gradiente

Después de haber implementado con éxito el paso hacia adelante en la parte anterior, nuestro modelo es capaz de hacer predicciones. Sin embargo, en la etapa inicial, sin entrenamiento, estas predicciones serán aleatorias. Para que el modelo comience a producir resultados correctos, sus parámetros internos (pesos y sesgos) deben ajustarse para minimizar el error entre las predicciones y las etiquetas verdaderas. Algoritmos como el Descenso de Gradiente y la Retropropagación se utilizan precisamente para este fin.

Imagina que estás en la cima de una montaña en medio de una densa niebla, y tu tarea es descender al punto más bajo del valle. No puedes ver todo el paisaje, pero puedes sentir la pendiente bajo tus pies. El gradiente es un vector que indica la dirección del ascenso más pronunciado. En consecuencia, para descender, necesitas moverte en la dirección opuesta al gradiente. En el contexto de las redes neuronales, la "montaña" es la superficie de la función de error (o función de pérdida), y el "valle" es su mínimo global. El objetivo del descenso de gradiente es encontrar este mínimo ajustando iterativamente los pesos de la red en la dirección opuesta al gradiente de la función de pérdida con respecto a estos pesos.

La Retropropagación es un algoritmo que calcula eficientemente estos gradientes. Funciona propagando el error desde la capa de salida de la red hacia la capa de entrada, distribuyendo la responsabilidad del error entre todos los pesos. Matemáticamente, esto se implementa a través de la Regla de la Cadena de diferenciación, que nos permite calcular cómo un cambio en cada peso afecta el error general del modelo.

Función de Pérdida: Entropía Cruzada

Para medir la "discrepancia" de las predicciones del modelo, se utiliza una función de pérdida. En tareas de clasificación, como el reconocimiento de dígitos escritos a mano de MNIST, la pérdida de entropía cruzada es una elección estándar. Cuantifica la divergencia entre la distribución de probabilidad predicha por el modelo y la distribución verdadera (donde solo la clase correcta tiene una probabilidad de 1, y las demás son 0). Cuanto mayor sea esta divergencia, mayor será el valor de la función de pérdida.

Para un solo ejemplo, la entropía cruzada se calcula utilizando la fórmula:

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

Donde:

![K](./images/image-2.svg) — es el número de clases (para MNIST, esto es 10).
![y_i](./images/image-3.svg) — es el valor verdadero (1 para la clase correcta, 0 para las demás).
![\hat{y}_i](./images/image-4.svg) — es la predicción del modelo (probabilidad obtenida después de la capa Softmax).

Dado que el entrenamiento generalmente ocurre en lotes de datos, no en ejemplos individuales, la función de pérdida total para un lote se promedia sobre su tamaño (![N](./images/image-5.svg) — tamaño del lote):

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

A continuación se muestra el código C++ para calcular la entropía cruzada en la CPU, lo que simplifica la depuración y la comprensión:

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

Retropropagación y la Regla de la Cadena

Para minimizar la función de pérdida ![L](./images/image-9.svg), necesitamos saber cómo su valor cambia con un pequeño cambio en cada peso en la red. Este es el gradiente de la función de pérdida con respecto a los pesos. Por ejemplo, para el peso ![W_1](./images/image-10.svg) de la primera capa, necesitamos encontrar ![\frac{\partial L}{\partial W_1}](./images/image-8.svg).

La arquitectura de nuestra red, desarrollada en artículos anteriores, se ve de la siguiente manera:

Primera Capa Lineal: ![z_1 = W_1 \cdot x + b_1](./images/image-11.svg)
Activación (ReLU): ![a_1 = \sigma(z_1)](./images/image-12.svg)
Segunda Capa Lineal: ![z_2 = W_2 \cdot a_1 + b_2](./images/image-13.svg)
Salida (Softmax + Pérdida): ![L = Loss(Softmax(z_2))](./images/image-14.svg)

Para calcular el gradiente ![\frac{\partial L}{\partial W_1}](./images/image-16.svg), aplicamos la Regla de la Cadena:

![\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)

Cada término en esta cadena representa un gradiente local, que muestra cómo la salida de una capa o función dada afecta su entrada. La retropropagación comienza calculando el gradiente de la función de pérdida con respecto a las salidas de la última capa, y luego propaga secuencialmente este gradiente hacia atrás a través de todas las capas de la red.

Derivación Detallada de Gradientes para Softmax y Entropía Cruzada

Comencemos desde el final de la cadena – la derivada de la función de pérdida ![L](./images/image-25.svg) con respecto a las salidas de la segunda capa lineal ![z_2](./images/image-20.svg) después de aplicar Softmax. Este es uno de los gradientes más importantes y frecuentemente utilizados. Recordemos cómo se ve Softmax para cada elemento ![i](./images/image-19.svg) en el vector ![z_2](./images/image-20.svg):

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

Donde ![z_2](./images/image-22.svg) es la salida de la segunda capa lineal (logits), y ![a_2](./images/image-23.svg) es el resultado de Softmax (probabilidades). El gradiente ![\frac{\partial L}{\partial z_{2,i}}](./images/image-18.svg) puede descomponerse usando la Regla de la Cadena:

![\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)

Primero, encontremos la derivada de la función de pérdida ![L](./images/image-27.svg) con respecto a ![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)

Ahora pasemos a la derivada de Softmax con respecto a ![z](./images/image-30.svg). Aquí, se necesitará la regla del cociente para la diferenciación. La fórmula de Softmax: ![a_i = \frac{e^{z_i}}{\sum_{k=1}^K e^{z_k}}](./images/image-32.svg)

Si ![i=j](./images/image-43.svg), la derivada es:

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

Si ![i \neq j](./images/image-37.svg), la derivada es:

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

Usando el símbolo delta de Kronecker ![\delta_{ij}](./images/image-42.svg), estos dos casos pueden combinarse:

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

Combinando las derivadas de la entropía cruzada y Softmax, obtenemos:

![\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)

Una simplificación adicional, considerando que la suma de las probabilidades verdaderas ![\sum_{k=1}^{K} y_k](./images/image-51.svg) siempre es igual a 1, conduce a una fórmula sorprendentemente simple y elegante:

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

Esto significa que el gradiente de la función de pérdida con respecto a los logits (salidas antes de Softmax) es simplemente la diferencia entre las probabilidades predichas y las etiquetas verdaderas. Este resultado es una piedra angular para implementar la retropropagación en modelos de clasificación.

Implementando el Gradiente en C++

Aplicando la fórmula derivada, podemos implementar el cálculo del gradiente para la última capa de nuestra red neuronal. Este gradiente se utilizará luego para actualizar los pesos, moviendo el modelo hacia un error menor. Así es como se ve en el código 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);
  }
}

Este código calcula el gradiente de error para cada elemento en el lote y luego lo promedia sobre el tamaño del lote. El tensor grad resultante contiene información sobre cuánto cambiar los logits ![z_2](./images/image-54.svg) para reducir el error. Este gradiente se propagará luego a través de las capas precedentes de la red, permitiendo que los pesos ![W_2](./images/image-13.svg), ![b_2](./images/image-13.svg), ![W_1](./images/image-11.svg), y ![b_1](./images/image-11.svg) se ajusten utilizando las derivadas locales correspondientes.

Puntos Clave:

El Descenso de Gradiente es un método de optimización iterativo que ajusta los parámetros del modelo en la dirección opuesta al gradiente de la función de pérdida para encontrar su mínimo.
La Retropropagación es un algoritmo que utiliza la regla de la cadena para calcular eficientemente los gradientes de la función de pérdida con respecto a todos los pesos en una red neuronal.
La Entropía Cruzada es una función de pérdida estándar para tareas de clasificación, que mide la divergencia entre las distribuciones de probabilidad predichas y verdaderas.
La derivada de Softmax + Entropía Cruzada con respecto a los logits se simplifica a la diferencia entre las probabilidades predichas y las etiquetas verdaderas, simplificando significativamente la implementación.
Implementar desde cero en C++/CUDA permite una comprensión profunda de los mecanismos internos de modelos como Transformer, sin depender de abstracciones de alto nivel.

— Editorial Team