O que é uma camada FFN num transformador e por que é necessária a não linearidade entre os blocos de atenção?

A camada FFN (Feed-Forward Network) em um transformador é uma rede neural totalmente conectada de duas camadas, que é aplicada separadamente a cada elemento da sequência após a camada de atenção. Sua tarefa é adicionar um processamento não linear adicional e aumentar a capacidade expressiva do modelo.

A não linearidade entre os blocos de atenção é necessária para que o modelo possa aproximar funções e dependências complexas. A atenção é responsável pela soma ponderada das informações em diferentes posições, e a FFN com ativação não linear (por exemplo, ReLU ou GELU) ajuda o transformador a modelar padrões mais complexos do que simples combinações lineares.

Exemplo de estrutura de uma camada 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)))

Aqui, d_model é a dimensão dos vetores de entrada e saída, d_ff é a dimensão da camada interna, geralmente maior que d_model.

Dessa forma, a camada FFN com não linearidade fornece uma transformação não linear profunda das características, o que melhora o treinamento e a qualidade do modelo.