Attention Sinks

Um sumidouro de atenção é um token que recebe uma grande proporção dos pesos de atenção em muitas cabeças e camadas, apesar de o seu conteúdo ser irrelevante para a consulta; na maioria dos casos, é o primeiro token da sequência. Guangxuan Xiao e colegas do MIT, da Meta AI e da Universidade Carnegie Mellon deram nome ao fenómeno em Efficient Streaming Language Models with Attention Sinks, publicado na ICLR 2024. Demonstraram que modelos de linguagem grandes (LLMs), como Llama 2, MPT, Falcon e Pythia, deixam de funcionar corretamente quando uma cache KV deslizante descarta os primeiros tokens, e que manter quatro desses tokens restaura uma geração estável.

Porque se formam os sumidouros#

A autoatenção transforma as pontuações de cada consulta em pesos com uma softmax, e a soma desses pesos é sempre um. Quando nenhum token do contexto é relevante para uma cabeça, o peso tem de ser atribuído a algum lado. Num modelo de linguagem causal, todos os tokens posteriores podem ver os primeiros, pelo que o treinamento faz deles o destino predefinido dos pesos de que uma cabeça não precisa. Na análise do StreamingLLM ao Llama 2, todas as camadas para além das duas inferiores concentravam a atenção nos tokens iniciais em todas as cabeças. Um logit de sumidouro adicionado à softmax mostra este efeito:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Sem um sumidouro, a cabeça distribui todo o seu peso por seis valores irrelevantes e produz uma saída com norma 0,83. Com um logit de sumidouro igual a 3, o sumidouro recebe 77% do peso e a saída diminui para 0,19, pelo que a cabeça quase não faz nada para esta consulta. Os modelos de pesos abertos gpt-oss da OpenAI incluem esta funcionalidade através de um logit de sumidouro aprendido para cada cabeça de atenção, dando a cada cabeça um local para o peso excedente que não acrescenta nada à saída.

Geração contínua com o StreamingLLM#

Um modelo que gera conteúdo indefinidamente ultrapassa qualquer janela de contexto e não consegue manter todos os tokens na cache KV; a atenção simples de janela, que mantém apenas os tokens mais recentes, falha assim que os primeiros tokens são descartados. O StreamingLLM mantém permanentemente as chaves e os valores de quatro tokens iniciais, juntamente com uma janela móvel dos tokens recentes, e atribui posições pela localização na cache em vez de no texto original, voltando a aplicar as incorporações posicionais rotativas (RoPE) a cada passo:

Isto manteve a modelação de linguagem estável ao longo de mais de 4 milhões de tokens e foi até 22,2 vezes mais rápido do que uma janela deslizante que recalcula a cache para cada novo token. A janela continua a limitar aquilo de que o modelo se consegue lembrar: os tokens descartados do meio desaparecem para sempre, pelo que um chatbot ou agente de IA em fluxo contínuo só se lembra das interações mais recentes. Os autores também pré-treinaram modelos com um token de substituição dedicado e treinável no início de cada amostra, que assumiu por si só o papel de sumidouro.

Sumidouros de atenção e termos relacionados#

• Atenção de janela deslizante: Limita cada token aos vizinhos recentes. Aplicada durante a inferência a um modelo treinado com atenção completa, falha quando os sumidouros saem da janela.
• Degradação do contexto: A redução da precisão à medida que um prompt aumenta, embora ainda caiba na janela de contexto.

Sumidouros em modelos de visão#

Os transformadores de visão (ViTs) apresentam um comportamento relacionado. Em Vision Transformers Need Registers, investigadores da Meta e do Inria encontraram tokens com normas invulgarmente elevadas em patches de fundo com pouca informação de ViTs supervisionados e autossupervisionados, reutilizados para computação interna. A adição de alguns tokens de registo à entrada absorveu esse papel, removeu os artefactos dos mapas de características e dos mapas de atenção, e melhorou a predição densa.

O SAM 2 da Meta, executado no pacote Ultralytics, gere a memória de vídeo com uma estratégia semelhante, que mantém uma âncora e uma janela: para cada novo fotograma, a atenção de memória consulta as entradas do banco de memória dos fotogramas indicados pelo utilizador, mantidas durante todo o vídeo, além das entradas dos 6 fotogramas mais recentes, pelo que o custo de atenção por fotograma permanece fixo, independentemente da duração do vídeo.