Reformer

O Reformer é um Transformer eficiente em termos de memória que Nikita Kitaev, Łukasz Kaiser e Anselm Levskaya, da UC Berkeley e do Google Research, apresentaram em Reformer: The Efficient Transformer na ICLR 2020. Ele aborda os dois custos que limitam o comprimento das sequências. A atenção por hashing sensível à localidade (LSH) substitui a matriz de atenção completa n × n e custa cerca de L log L para uma sequência de comprimento L; as camadas residuais reversíveis armazenam as ativações necessárias para a retropropagação uma única vez para toda a rede, em vez de uma vez por camada. O Google informou que o Reformer processa janelas de contexto de até 1 milhão de palavras em um único acelerador com 16 GB de memória e que processou o romance inteiro Crime e Castigo como um único exemplo de treinamento.

Atenção LSH#

Após o softmax, a atenção de cada consulta é dominada pelas poucas chaves mais semelhantes a ela, por isso o Reformer procura apenas essas chaves. Ele usa uma única projeção compartilhada para consultas e chaves e aplica hashing a cada vetor com hashing sensível à localidade angular: multiplica por uma matriz aleatória R e usa como bucket a posição do maior valor em [xR, −xR], de modo que vetores que apontam em direções semelhantes geralmente ficam no mesmo bucket. Em seguida, os tokens são ordenados por bucket, a sequência ordenada é dividida em blocos iguais e cada token presta atenção apenas ao próprio bloco e ao bloco anterior. Várias rodadas de hashing são executadas em paralelo para capturar pares semelhantes que uma única rodada separaria.

LSH é um método clássico de busca vetorial aproximada de vizinhos mais próximos: encontra a maioria dos k vizinhos mais próximos de uma consulta por uma fração do custo de uma busca exata. A etapa de agrupamento em buckets é curta o bastante para ser executada isoladamente:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

O vetor próximo, com similaridade de cosseno de 0,96 em relação à consulta, compartilha o bucket em 72% das 1.000 rodadas; o vetor não relacionado, com 0,04, compartilha o bucket em 9%. Portanto, a atenção restrita a um bucket mantém a maioria dos pares relevantes.

Camadas reversíveis e divisão em blocos#

Uma rede neural padrão mantém as ativações de todas as camadas para a passagem de volta, por isso a memória aumenta tanto com a profundidade quanto com o comprimento. O Reformer adapta as redes residuais reversíveis (RevNets): divide as ativações em duas metades e atualiza cada metade com uma função da outra.

Durante a passagem de volta, cada camada recupera suas entradas a partir das saídas: x2 = y2 − FeedForward(y1) e, em seguida, x1 = y1 − Attention(x2). Assim, apenas as ativações finais são armazenadas e cada camada é recalculada uma vez. O checkpointing de gradientes faz uma troca semelhante: armazena ativações em alguns checkpoints e recalcula as camadas entre eles. O Reformer também executa suas camadas feed-forward em blocos de posições, porque a camada oculta delas é várias vezes mais larga que o modelo. A memória das ativações é o que limita o tamanho do lote ao treinar o Ultralytics YOLO26: com batch=-1, o AutoBatch mede a memória da GPU para vários tamanhos de lote, ajusta uma linha reta e escolhe o lote que, segundo a previsão, ocupará 60% da memória livre.

Resultados e implementações#

O Reformer igualou os Transformers padrão em texto ao nível de caracteres do enwik8 e na geração de imagens pixel a pixel no ImageNet-64, usando muito menos memória e executando mais rapidamente em sequências longas. A biblioteca Hugging Face Transformers implementa o modelo como ReformerModel, com checkpoints disponibilizados e treinados no enwik8 e em Crime e Castigo.

Reformer e modelos relacionados#

• Longformer: Define o padrão de atenção antecipadamente: uma janela deslizante em torno de cada token, além de alguns tokens globais.
• Atenção esparsa: Família de métodos que pontuam apenas alguns pares de tokens. A atenção LSH é um método dessa família baseado no conteúdo.
• Atenção linear: Substitui o softmax por um mapa de características de kernel para que a atenção possa ser calculada em tempo linear.
• Atenção Flash: Calcula a atenção exata em blocos na memória interna da GPU sem armazenar a matriz de pontuações, de modo que a memória cresce linearmente enquanto o número de pontuações permanece quadrático.