Implementando um Transformer do zero com TensorFlow.js: guia prático

Introdução

Os Transformers revolucionaram o campo de processamento de linguagem natural (NLP) desde que foram apresentados em 2017 com o paper “Attention Is All You Need”. Diferente de RNNs sequenciais, o Transformer permite processar todas as palavras de uma frase em paralelo, focando em diferentes partes do texto graças ao mecanismo de atenção. Neste guia, você aprenderá como construir e treinar um modelo Transformer do zero usando TensorFlow.js. Veremos em detalhe o que é o mecanismo de atenção multi-cabeça, como funcionam as estruturas de codificador e decodificador do modelo, e como implementar cada componente passo a passo com código JavaScript. Ao final, você terá compreendido as principais peças do Transformer e um exemplo de como treinar um modelo simples.

Visão geral do Transformer

O Transformer é uma arquitetura de aprendizado profundo baseada quase que inteiramente no mecanismo de atenção. Sua principal vantagem sobre modelos sequenciais (como LSTM) é a habilidade de capturar dependências de longo alcance de forma eficiente e paralela. Sua estrutura básica consiste em camadas empilhadas de codificador (encoder) e decodificador (decoder), onde cada camada usa mecanismos de atenção e redes neurais feed-forward. Entre as principais características estão:

Atenção Scaled Dot-Product: cada palavra consulta outras palavras (incluindo ela mesma) para capturar contexto.
Atenção Multi-Cabeça: múltiplas “cabeças” de atenção paralelas permitem focar em diferentes aspectos da frase simultaneamente.
Codificação Posicional: como Transformers não são recorrentes, adiciona-se informação de posição das palavras às representações.
Conexões residuais e normalização: cada subcamada (atenção ou feed-forward) tem conexão residual e normalização em camadas (LayerNorm).

De forma resumida, um Transformer típico para tradução, por exemplo, funciona assim: o codificador processa a frase de entrada e produz uma representação contextualizada. O decodificador, por sua vez, lê essa representação e gera palavra por palavra a frase de saída, usando atenção tanto consigo mesmo (atenção mascarada para não “olhar o futuro”) quanto à saída do codificador. Diagramas ilustrativos como o abaixo ajudam a entender a estrutura geral:

Camadas do Codificador: cada camada do encoder tem um bloco de atenção (self-attention) e um bloco feed-forward.
Camadas do Decodificador: cada camada do decoder tem atenção mascarada (self-attention no próprio decoder), atenção encoder-decoder, e um feed-forward.

Em suma, o Transformer combina atenção e redes feed-forward em camadas empilhadas para aprender representações do texto.

Atenção multi-cabeça

O mecanismo de atenção permite que cada posição da entrada “atenda” a outras posições para construir representações contextualizadas. A versão básica é chamada de atenção por produto escalar escalonado (Scaled Dot-Product Attention). Para cada conjunto de entradas, definimos três matrizes: Q (queries), K (keys) e V (values). Grosso modo, podemos pensar em Q como “perguntas”, K como “índices” e V como “conteúdos de respostas”. O cálculo de atenção padrão segue estes passos:

Calcular scores de similaridade entre queries e keys:
[ \text{scores} = Q \times K^T ]
Escalonar (scale): divide-se por (\sqrt{d_k}) (dimensão de K) para evitar gradientes muito grandes.
Aplicar softmax para normalizar em probabilidade:
[ \text{weights} = \text{softmax}(\text{scores}) ]
Multiplicar pelos valores (values) para obter o resultado da atenção:
[ \text{output} = \text{weights} \times V ]

Em palavras, cada posição (query) olha para todas as outras (keys) e decide quão importante cada uma é (via softmax nos scores), produzindo assim uma combinação ponderada dos valores correspondentes. Esse processo captura o contexto de forma dinâmica: cada palavra agrega informações de outras baseando-se na similaridade.

Analogias didáticas: imagine que você faz perguntas (Q) numa biblioteca. Cada livro na biblioteca tem uma etiqueta (K, key) que indica seus tópicos, e o conteúdo do livro (V, value). A atenção ajuda a encontrar quais livros são relevantes para a sua pergunta: você calcula uma pontuação de quão bem cada etiqueta (key) responde à sua pergunta (query) e, em seguida, lê trechos (value) desses livros de acordo (no output). O softmax garante que você foque nos livros mais relevantes (as maiores probabilidades).

Atenção por produto escalar (Scaled Dot-Product)

A implementação básica da atenção no TensorFlow.js pode ser feita usando operações matriciais. Por exemplo:

function scaledDotProductAttention(Q, K, V) {
  // Matriz de similaridade: Q * K^T
  const matmulQK = tf.matMul(Q, K, false, true);
  // Escalonamento pela raiz quadrada da dimensão das keys
  const dk = tf.scalar(Math.sqrt(K.shape[2]));
  const scaled = tf.div(matmulQK, dk);
  // Aplicar softmax nas últimas dimensões (normaliza scores)
  const weights = tf.softmax(scaled, -1);
  // Multiplicar pesos pelos valores
  const output = tf.matMul(weights, V);
  return output;
}

Neste código:

Q, K, V são tensores do TensorFlow.js com forma [batch, seqLen, depth].
tf.matMul(Q, K, false, true) calcula (Q K^T).
Dividimos pelo escalar com tf.div.
tf.softmax(..., -1) aplica a função softmax ao longo da última dimensão (labels).
Por fim, multiplicamos weights * V para obter o output: essencialmente, cada vetor final é uma combinação ponderada dos vetores V.

Atenção multi-cabeça (Multi-Head)

A atenção multi-cabeça amplia o conceito acima. Em vez de ter apenas uma atenção, dividimos a projeção de Q, K e V em várias partes (cabeças). Cada cabeça aprende a focar em diferentes tipos de relações no texto. Depois, concatenamos os resultados de todas as cabeças e projetamos de volta ao espaço original. Em fórmula conceitual:

Linhas de projeção para cada cabeça:
[ Q_i = Q W_i^Q,\quad K_i = K W_i^K,\quad V_i = V W_i^V ]
Atenção escalonada em cada cabeça:
[ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) ]
Concatenar todas as cabeças e projetar:
[ \text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O ]

Em código multi-cabeça simplificado (ignora muitas verificações) ficaria assim:

function multiHeadAttention(Q, K, V, numHeads) {
  const depth = Q.shape[2];
  const headDim = depth / numHeads;
  
  // Camadas lineares para projecoes de Q, K, V
  const denseQ = tf.layers.dense({units: depth}).apply(Q);
  const denseK = tf.layers.dense({units: depth}).apply(K);
  const denseV = tf.layers.dense({units: depth}).apply(V);
  
  // Dividir em múltiplas cabeças
  const reshapeForHeads = x => tf.transpose(
    tf.reshape(x, [-1, x.shape[1], numHeads, headDim]),
    [0, 2, 1, 3]
  );
  
  const Q_ = reshapeForHeads(denseQ);
  const K_ = reshapeForHeads(denseK);
  const V_ = reshapeForHeads(denseV);
  
  // Atenção escalar para cada cabeça
  const attentionHeads = scaledDotProductAttention(Q_, K_, V_);
  
  // Reconstruir tensor concatenando as cabeças
  const concatHeads = tf.reshape(
    tf.transpose(attentionHeads, [0, 2, 1, 3]),
    [-1, Q.shape[1], depth]
  );
  
  // Projeção final
  const output = tf.layers.dense({units: depth}).apply(concatHeads);
  return output;
}

Nesse exemplo:

Primeiro criamos camadas tf.layers.dense para projetar Q, K, V para um espaço de dimensão depth.
Em seguida, usamos tf.reshape e tf.transpose para dividir as projeções em numHeads cabeças e posicionar as dimensões corretamente.
Calculamos a atenção escalonada em cada uma (scaledDotProductAttention) em paralelo.
Reagrupamos os resultados com tf.reshape (concatenando as cabeças de volta) e aplicamos uma última camada dense para combinar tudo.

A grande vantagem é que cada cabeça pode “olhar” para diferentes partes da frase. Por exemplo, uma cabeça pode focar na concordância verbal enquanto outra foca no sujeito-objeto, etc. No fim, os resultados são combinados dando uma representação rica dos contextos.

Estrutura de codificador e decodificador

Com o mecanismo de atenção definido, podemos montar o Transformer completo, que consiste em pilhas de camadas de codificador (encoder) e decodificador (decoder).

Codificador (Encoder)

O codificador recebe a frase de entrada (por exemplo, um vetor de índices de palavras) e produz uma sequência de representações contextuais. Cada camada de codificador é composta por:

Camada de Embedding: mapeia cada índice de palavra para um vetor denso de dimensão dModel.
Positional Encoding: adiciona vetores de codificação posicional a cada embedding para indicar a posição de cada palavra na frase.
Atenção Multi-Cabeça (Self-Attention): cada palavra interage com todas as outras (no próprio input).
Dropout (opcional) aplicado à saída da atenção.
Conexão Residual + Normalização em Camada: adiciona a entrada da subcamada à saída (skip connection) e aplica LayerNorm.
Feed-Forward (FFN): rede neural de duas camadas (com ReLU) aplicado individualmente a cada posição.
Outra Conexão Residual + Normalização: adiciona a saída da atenção normalizada à saída do FFN e normaliza novamente.

Em pseudocódigo, uma camada do encoder pode ser assim:

function encoderLayer(x, numHeads, dff) {
  // x: [batch, seqLen, dModel]
  const selfAttn = multiHeadAttention(x, x, x, numHeads);  // Atenção self
  const attnOut = tf.layers.layerNormalization().apply(
                      tf.add(selfAttn, x));               // Residual + norm
  const ffn1 = tf.layers.dense({units: dff, activation: 'relu'}).apply(attnOut);
  const ffn2 = tf.layers.dense({units: x.shape[2]}).apply(ffn1);
  const out = tf.layers.layerNormalization().apply(
                 tf.add(ffn2, attnOut));    // Residual + norm após feed-forward
  return out;
}

numHeads: número de cabeças no multi-head.
dff: dimensão interna da camada feed-forward (geralmente maior que dModel).
Usamos tf.add para implementar a conexão residual (adicionando a entrada à saída da subcamada).
tf.layers.layerNormalization normaliza cada posição (melhora a estabilidade durante o treinamento).

Para compor o encoder completo, empilhamos numEncoderLayers dessas camadas. Antes da primeira camada, convertemos os índices em embeddings e somamos a codificação posicional:

// Parâmetros exemplo
const vocabSize = 8000;
const dModel = 64;
const maxSeqLen = 100;
const numEncoderLayers = 2;

// Input de índices de palavras
const input = tf.input({shape: [maxSeqLen], dtype: 'int32'});

// Camada de Embedding + Positional Encoding
const embeddingLayer = tf.layers.embedding({inputDim: vocabSize, outputDim: dModel});
let x = embeddingLayer.apply(input);  // shape: [batch, seqLen, dModel]

// Adiciona codificação posicional (exemplo simples)
const posEncoding = tf.tensor(getPositionalEncoding(maxSeqLen, dModel));  
x = tf.add(x, posEncoding);

// Pilha de camadas do codificador
for (let i = 0; i < numEncoderLayers; i++) {
  x = encoderLayer(x, /*numHeads=*/4, /*dff=*/128);
}

// x agora contém a saída do encoder (contextos das palavras)
 const encoderOutput = x;

Observação: a função getPositionalEncoding geraria uma matriz [maxSeqLen, dModel] usando senos e cossenos (você pode implementá-la seguindo a fórmula de Vaswani et al.). O importante é que cada posição receba um vetor único para manter a noção de ordem.

Decodificador (Decoder)

O decodificador gera a sequência de saída (por exemplo, traduzindo palavra a palavra) e também consiste em camadas empilhadas com três componentes principais:

Masked Self-Attention: a cada passo, o decoder só deve olhar para posições anteriores (incluindo atual), não conhecer futuros. Para isso, aplicamos uma máscara que impede que a atenção veja além da posição atual.
Atenção Encoder-Decoder: vogar interage com a saída do codificador. Aqui, Q vem do decoder atual, e K/V vêm do encoderOutput.
Feed-Forward: idem ao caso do codificador.
Conexões residuais e normalização após cada subcamada.

Em código, uma camada do decoder pode ficar assim:

function decoderLayer(x, encOutput, numHeads, dff) {
  // x: [batch, targetSeqLen, dModel], encOutput: [batch, inputSeqLen, dModel]
  // Masked self-attention (não olhando o futuro)
  let selfAttn = multiHeadAttention(x, x, x, numHeads);
  let selfAttnOut = tf.layers.layerNormalization().apply(
                      tf.add(selfAttn, x));

  // Atenção encoder-decoder
  let encDecAttn = multiHeadAttention(selfAttnOut, encOutput, encOutput, numHeads);
  let encDecOut = tf.layers.layerNormalization().apply(
                     tf.add(encDecAttn, selfAttnOut));

  // Feed-forward
  const ffn1 = tf.layers.dense({units: dff, activation: 'relu'}).apply(encDecOut);
  const ffn2 = tf.layers.dense({units: encDecOut.shape[2]}).apply(ffn1);
  const out = tf.layers.layerNormalization().apply(
                 tf.add(ffn2, encDecOut));  // Residual + norm
  return out;
}

Aqui, x é a entrada atual do decoder (inicialmente apenas o token start e depois os tokens gerados até então), e encOutput vem do encoder. Note o uso de uma máscara (não mostrada no pseudo-código), que eliminaria as entradas futuras na atenção mascarada. Em TensorFlow.js, isso exigiria inserir valores muito negativos nos pontos mascarados antes do softmax.

Para montar o decoder completo, também usamos uma camada de embedding + posicional para os tokens alvo, seguida por numDecoderLayers dessas camadas:

const targetVocabSize = 8000;
const targetInput = tf.input({shape: [maxSeqLen], dtype: 'int32'});
let y = embeddingLayer.apply(targetInput);  // usando a mesma layer de embedding ou outra
y = tf.add(y, posEncoding);  // soma posicional

for (let i = 0; i < numEncoderLayers; i++) {
  y = decoderLayer(y, encoderOutput, /*numHeads=*/4, /*dff=*/128);
}

const decoderOutput = tf.layers.dense({units: targetVocabSize}).apply(y);
// Este dense final converte para logits de cada palavra no vocabulário

Dessa forma, decoderOutput terá forma [batch, targetSeqLen, targetVocabSize], que podemos comparar com as sequências esperadas durante o treinamento.

Implementação com TensorFlow.js

Agora que compreendemos os blocos fundamentais, vamos resumir como montar e treinar esse modelo em TensorFlow.js. A ideia principal é usar a API Layers (semelhante ao Keras) para definir o modelo e depois compilar/treinar.

Passos principais:

Definir Inputs: tensores de entrada para as frases de origem e destino (por exemplo, input e targetInput no exemplo acima).
Construir camadas de embedding e posicional (podem ser funções auxiliares ou camadas customizadas).
Empilhar camadas do Encoder e Decoder conforme ilustrado, chamando repetidamente encoderLayer e decoderLayer.
Criar saídas: a saída final (logits) é geralmente um tf.layers.dense sobre o decoderOutput.
Criar o modelo: tf.model({inputs: [input, targetInput], outputs: decoderOutput}).
Compilar: definir o otimizador, função de perda (por exemplo, cross-entropy categorico) e métricas.
Treinar com .fit(): forneça os dados de treinamento (pares de entrada-saída).

Exemplo simplificado de compilação e treinamento:

const model = tf.model({
  inputs: [input, targetInput],
  outputs: decoderOutput
});

model.compile({
  optimizer: tf.train.adam(),
  loss: 'sparseCategoricalCrossentropy',  // apropriado para índices de classe
  metrics: ['accuracy']
});

// Exemplo de dados dummy (substitua pelos seus dados reais)
const xTrain = tf.tensor2d([[1,2,3,4,0], [5,6,7,0,0]], [2,5], 'int32');  // frases origem
const yTrainInput = tf.tensor2d([[1,2,3,4,0], [5,6,7,0,0]], [2,5], 'int32');  // shift por iniciar token
const yTrainTarget = tf.tensor2d([[2,3,4,0,0], [6,7,0,0,0]], [2,5], 'int32'); // alvo

// Treinamento
await model.fit([xTrain, yTrainInput], yTrainTarget, {
  epochs: 10,
  batchSize: 2
});

No snippet acima:

xTrain são frases de entrada codificadas (shape [2,5], batch de 2, 5 tokens cada).
yTrainInput é a frase destino deslocada à direita (em tradução, geralmente começa pelo token <start>).
yTrainTarget são as palavras-alvo que queremos prever, ignorando o token inicial.
Com .fit, o TensorFlow.js treina o modelo, ajustando os pesos das camadas (densas, etc.) pra minimizar o erro de predição.

Este exemplo é muito simplificado e serve apenas para ilustrar o processo. Em uma aplicação real você prepararia um dataset maior (por exemplo, pares de sentenças em idiomas), trataria vocabulário (tokenização), e ajustaria hiperparâmetros (tamanho das camadas, taxa de aprendizado, etc.).

Conclusão

Neste artigo, vimos como um modelo Transformer funciona e como implementá-lo do zero usando TensorFlow.js. Aprendemos que o coração do Transformer é o mecanismo de atenção multi-cabeça, que permite ao modelo aprender relações entre diferentes partes do texto de forma paralela. Também exploramos a estrutura de codificador e decodificador: o codificador lê o input e cria representações contextuais, enquanto o decodificador gera o output palavra a palavra, utilizando atenção tanto sobre si próprio quanto sobre a saída do codificador.

No caminho prático, apresentamos código em TensorFlow.js para:

Calcular a atenção (scaled dot-product).
Implementar atenção multi-cabeça.
Montar camadas do encoder e decoder com conexões residuais e normalização.
Definir e treinar um modelo Transformer simples.

Como próximos passos, você pode experimentar aumentando a quantidade de camadas, ajustando dimensões ou usando dados reais de tradução. Também é interessante explorar bibliotecas e modelos pré-treinados (como os disponibilizados pelo Hugging Face) para ver como Transformers robustos são usados em produção. O TensorFlow.js possibilita até rodar esses modelos no navegador ou em aplicações Node.js, tornando essas pesquisas acessíveis. A jornada do Transformer é incrível e está apenas começando – com a fundação que construímos aqui, você estará pronto para avançar em aplicações práticas de NLP e até mesmo visão computacional usando transformadores de forma criativa.