antonio leandro

ia generativa

Efficient Streaming Language Models with Attention Sinks

paper · Xiao, Guangxuan, Tian, Yuandong, Chen, Beidi, Han, Song, Lewis, Mike · · ~46 min de leitura do original

a tese

os primeiros tokens de uma sequência viram ralo de atenção: o softmax precisa somar um e o modelo despeja o excedente ali — jogue esses quatro fora do cache e ele colapsa; mantenha-os e ele gera 4 milhões de tokens

o que fica

  1. O softmax obriga a atenção a somar um, então o modelo precisa despejar o excedente em algum lugar — e os tokens iniciais, visíveis a todos os seguintes, são os candidatos mais fáceis de treinar para esse papel.
  2. Descartar o KV do primeiro token leva a perplexidade do Llama-2-13B de 5,40 para 5.158,07: não é degradação gradual, é colapso.
  3. O que importa nos tokens iniciais é a posição absoluta, não o conteúdo — substituí-los por quatro quebras de linha recupera a perplexidade para 5,60.
  4. A posição atribuída a cada token é a posição dentro do cache, não a do texto original: um cache com os tokens [0, 1, 2, 3, 6, 7, 8] recebe posições 0 a 7.
  5. Aumentar o cache não reduz a perplexidade de forma consistente, o que indica que o modelo não aproveita todo o contexto que já tem.
  6. StreamingLLM não estende a janela de contexto: o que sai do cache some, e em QA de documento longo ele perde para simplesmente truncar o texto.

o problema

Um assistente de conversa precisa rodar por dias sem reiniciar. Um LLM baseado em transformer não faz isso por dois motivos independentes. Primeiro, o KV cache cresce com cada token gerado: memória sobe, latência de decodificação sobe junto. Segundo, o modelo foi treinado com uma janela finita — 4K no Llama-2 — e a qualidade despenca quando a sequência passa disso. Nenhum dos dois problemas se resolve sozinho: mesmo que você tivesse memória infinita, o modelo continuaria quebrando ao ultrapassar o comprimento de pré-treino.

A saída óbvia é window attention: guarde só o KV dos L tokens mais recentes e descarte o resto. Memória constante, velocidade constante. Só que não funciona. O modelo colapsa no instante exato em que o cache enche e o primeiro token é despejado — perplexidade de 5,40 para 5.158,07 no Llama-2-13B. A alternativa que funciona, recomputar o KV da janela a cada token gerado, tem custo O(TL²) e é lenta demais para uso real. O paper parte dessa anomalia: por que remover o KV de um token irrelevante, a milhares de posições de distância, destrói o modelo?

a ideia

Olhando os mapas de atenção do Llama-2-7B, a partir da terceira camada quase toda cabeça, em quase toda camada, joga uma fatia enorme da atenção nos primeiros tokens da sequência — independente do que eles digam. Os autores chamam isso de attention sink, o ralo.

A explicação é o softmax. Ele exige que os scores somem um. Quando a query atual não tem correspondência forte com nenhum token anterior — o caso comum, porque o embedding já carrega o que precisa —, o modelo ainda assim precisa colocar aquela massa de probabilidade em algum lugar. Ele escolhe os tokens iniciais, e a escolha é estrutural: pela natureza autorregressiva, os primeiros tokens são visíveis a todos os outros, então são os únicos que podem servir de ralo universal. Não é semântica, é topologia. Trocar os quatro primeiros tokens por quebras de linha “\n” recupera a perplexidade (5,60), o que mostra que o modelo está ancorado na posição, não no conteúdo.

Se o ralo é o problema, mantenha o ralo. É isso.

como funciona

O KV cache do StreamingLLM tem duas partes: os quatro tokens iniciais, fixos, que nunca saem; e uma janela rolante com os tokens recentes. Quando o cache enche, só a janela rola. Nada de fine-tuning, nada de kernel customizado.

O detalhe que faz a coisa funcionar é o encoding posicional. As posições são atribuídas dentro do cache, não no texto original. Se o cache contém os tokens [0, 1, 2, 3, 6, 7, 8] e o modelo está decodificando o nono, as posições usadas são [0, 1, 2, 3, 4, 5, 6, 7] — nunca [0, 1, 2, 3, 6, 7, 8, 9]. Com RoPE, isso significa guardar as keys antes da rotação e aplicar a transformação posicional a cada passo de decodificação. Com ALiBi, aplica-se o viés linear contínuo em vez de um viés com salto. Como as distâncias nunca ultrapassam o tamanho da janela de pré-treino, o modelo nunca sai da faixa que conhece.

Por que quatro e não um? Porque nenhum desses modelos foi treinado com um token inicial constante — o <s> do Llama-2 é aplicado antes do chunking, então a posição zero acaba ocupada por um token quase aleatório. O modelo então distribui o ralo entre vários tokens iniciais. Um ou dois não bastam; quatro bastam; oito não melhoram. Pré-treinar do zero com um token de sink dedicado e aprendível resolve isso: em modelos de 160 milhões de parâmetros, um único sink já estabiliza tudo, sem custo na curva de loss nem nos sete benchmarks zero-shot testados.

o que isso custou

O paper é explícito: StreamingLLM não estende a janela de contexto e não dá memória de longo prazo. O que saiu do cache não existe mais. No StreamEval, a acurácia se mantém enquanto a resposta está dentro do cache e cai para zero assim que a distância entre pergunta e resposta passa desse limite. No LongBench, a configuração 4+3496 fica abaixo do baseline que simplesmente trunca o texto pelo meio, porque o prompt inicial se perde. Sumarização e QA de documento longo estão fora do escopo.

Há também um resultado desconfortável: aumentar o cache não reduz a perplexidade de forma monótona. O Llama-2-7B piora de 9,08 (4+2044) para 9,59 (4+4092). O modelo tem contexto e não usa.

onde isso aparece hoje

O ganho prático é grande: até 22,2× de speedup por token contra a recomputação da janela, com memória equivalente, e modelagem estável de textos de 4 milhões de tokens em Llama-2, MPT, Falcon e Pythia. O paper registra adoção em TensorRT-LLM, HuggingFace Transformers, MLC LLM e Intel Extension for Transformers.

O fenômeno vale mais que o método. Os autores encontram o mesmo ralo em BERT — o [SEP] recebe atenção desproporcional na maioria das camadas — e apontam o paralelo com os “registers” de Vision Transformers, o que sugere que o sink é propriedade do softmax do transformer, não de um treino específico. A dependência de encoding relativo liga direto ao RoPE, e a observação de que cache maior não melhora perplexidade confirma, por outro caminho, o que Lost in the Middle mostrou sobre uso de contexto longo.

lido na íntegra por pipeline de llm, revisado por antonio leandro antes de publicar ·