¿Qué es una capa FFN en un transformador y por qué es necesaria la no linealidad entre los bloques de atención?

La capa FFN (Red de Propagación Adelante) en un transformador es una red neuronal completamente conectada de dos capas, que se aplica de forma independiente a cada elemento de la secuencia después de la capa de atención. Su tarea es añadir un procesamiento no lineal adicional y aumentar la capacidad expresiva del modelo.

La no linealidad entre los bloques de atención es necesaria para que el modelo pueda aproximar funciones y dependencias complejas. La atención se encarga de la suma ponderada de la información en diferentes posiciones, y la FFN con activación no lineal (por ejemplo, ReLU o GELU) ayuda al transformador a modelar patrones más complejos que simplemente combinaciones lineales.

Ejemplo de estructura de una capa FFN:

import torch
import torch.nn as nn

class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = nn.GELU()
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(self.activation(self.linear1(x)))

Aquí, d_model es la dimensión de los vectores de entrada y salida, d_ff es la dimensión de la capa interna, generalmente mayor que d_model.

De esta manera, la capa FFN con no linealidad proporciona una transformación no lineal profunda de las características, lo que mejora el entrenamiento y la calidad del modelo.