Transformer-XL

O Transformer-XL é um Transformer que funciona como modelo de linguagem e mantém o contexto entre segmentos de texto de comprimento fixo, armazenando em cache os estados ocultos do segmento anterior e permitindo que o segmento seguinte preste atenção a eles. Zihang Dai, Zhilin Yang e colegas da Carnegie Mellon University e do Google Brain o apresentaram em Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context, na ACL 2019. Um Transformer convencional treinado em segmentos de comprimento fixo, cada um com sua própria janela de contexto, não consegue ver nada anterior ao início do segmento atual, um problema que os autores chamaram de fragmentação de contexto. O Transformer-XL aprendeu dependências 80% mais longas do que redes neurais recorrentes (RNNs) e 450% mais longas do que Transformers convencionais.

Recorrência no nível dos segmentos#

Quando o modelo processa um segmento, cada camada concatena os estados ocultos em cache do segmento anterior aos estados do segmento atual. As consultas vêm apenas do segmento atual, enquanto as chaves e os valores abrangem ambos; os gradientes param no cache, mantendo fixo o custo de treinamento por segmento. Uma única camada de atenção em PyTorch demonstra esse padrão:

import torch
from torch import nn

torch.manual_seed(0)
seg_len, dim = 4, 16
attn = nn.MultiheadAttention(dim, num_heads=2, batch_first=True)
hidden = torch.randn(1, 3 * seg_len, dim)  # hidden states for three consecutive segments

memory = torch.empty(1, 0, dim)  # nothing is cached before the first segment
for i, segment in enumerate(hidden.split(seg_len, dim=1), start=1):
    # Queries come from the current segment; keys and values also cover the cached one
    context = torch.cat([memory, segment], dim=1)
    out, weights = attn(segment, context, context)
    print(f"Segment {i}: {weights.shape[1]} queries attend over {weights.shape[2]} positions")
    # Cache this segment for the next one; detach() stops gradients as in Transformer-XL
    memory = segment.detach()

As 4 consultas do primeiro segmento veem apenas os 4 tokens desse segmento, e as consultas de cada segmento seguinte veem 8: seus próprios 4 mais os 4 armazenados em cache. Como a camada n de um segmento lê os estados em cache da camada n − 1 do segmento anterior, o alcance aumenta em mais um segmento a cada camada, chegando a cerca de N × L tokens para segmentos de comprimento N e L camadas:

O cache também acelera a avaliação. Um Transformer convencional que prevê um token após o outro precisa recalcular uma janela inteira para cada novo token, enquanto o Transformer-XL reutiliza os estados em cache, tornando a avaliação até 1.800 vezes mais rápida. O rastreamento do Ultralytics YOLO26 enfrenta o mesmo problema de fragmentação quando um script envia vídeo para model.track() um quadro por vez: persist=True mantém o estado do rastreador entre chamadas, para que cada objeto mantenha o mesmo ID de um quadro para o seguinte.

Codificação posicional relativa#

A reutilização dos estados em cache compromete as posições absolutas: o primeiro token do segmento em cache e o primeiro token do segmento atual teriam o mesmo embedding posicional. Em vez disso, o Transformer-XL calcula cada pontuação de atenção a partir do conteúdo da consulta e da chave, além de uma codificação senoidal da distância entre elas, com dois vetores de viés aprendidos e compartilhados entre posições, um para o conteúdo e outro para a posição. Como as pontuações dependem apenas das distâncias, durante a avaliação o modelo pode prestar atenção a um cache mais longo do que aquele visto no treinamento. Os embeddings posicionais rotativos (RoPE), usados na maioria dos grandes modelos de linguagem (LLM) atuais, também fazem as pontuações de atenção dependerem da posição relativa.

Resultados e influência#

O Transformer-XL estabeleceu novos resultados de ponta em cinco benchmarks: 0,99 bits por caractere no enwik8, 1,08 no text8 e perplexidade de 18,3 no WikiText-103, 21,8 no One Billion Word e 54,5 no Penn Treebank, sem ajuste fino. Treinado apenas com o WikiText-103, ele gerou artigos novos razoavelmente coerentes com milhares de tokens. O XLNet, desenvolvido por vários dos mesmos autores, baseou seu pré-treinamento de compreensão de linguagem natural na arquitetura Transformer-XL, e o Compressive Transformer da DeepMind ampliou o cache comprimindo memórias antigas em vez de descartá-las.

Transformer-XL vs. modelos relacionados#

• Memória de longo e curto prazo (LSTM): atualiza um estado oculto, token por token, portanto as informações distantes precisam passar por cada etapa intermediária.
• Longformer: lê de uma só vez um documento inteiro de 4.096 tokens, usando atenção em janelas e global, sem recorrência.
• Cache KV: armazena chaves e valores de tokens anteriores durante a geração para que não precisem ser recalculados. A memória do Transformer-XL também é usada no treinamento, com os gradientes interrompidos no limite do segmento.
• Modelos de espaço de estados (SSM): comprimem todo o histórico em um estado de tamanho fixo.