StreamingLLM — implementando modelos de linguagem de streaming eficientes com attention sink
(github.com/mit-han-lab)- O StreamingLLM é um framework para implantar LLMs em aplicações de streaming com entradas de comprimento infinito sem sacrificar eficiência nem desempenho
- Aborda o problema de que, em conversas de múltiplas rodadas que exigem interações longas, o cache dos estados Key/Value de tokens anteriores usa muita memória, e LLMs comuns não conseguem generalizar para textos mais longos que o comprimento de sequência usado no treinamento
- A window attention que armazena em cache apenas o KV recente falha quando o comprimento do texto excede o tamanho do cache, enquanto o StreamingLLM recupera significativamente o desempenho da window attention com um attention sink que mantém o KV dos tokens iniciais
- Faz com que LLMs treinados com uma janela de attention de comprimento finito generalizem para comprimentos de sequência infinitos sem fine-tuning, realizando modelagem de linguagem estável e eficiente com mais de até 4 milhões de tokens no Llama-2, MPT, Falcon e Pythia
- Em configuração de streaming, mostra até 22,2× de speedup em relação ao baseline de recomputação com sliding window
- Não expande a janela de contexto; mantém apenas os tokens recentes e o attention sink, descartando os tokens intermediários
- Se o Llama-2 foi pré-treinado com uma janela de contexto de 4096 tokens, o tamanho máximo do cache do Llama-2 no StreamingLLM também é 4096
- Ao inserir um livro longo, o modelo reconhece apenas os tokens mais recentes, portanto só consegue resumir a parte da conclusão
- Os casos de uso adequados são aplicações de streaming que precisam operar continuamente e evitar dependência de dados passados ou grande uso de memória; os exemplos apresentados incluem conversas de múltiplas rodadas e assistentes diários baseados em LLM
- É ortogonal aos métodos recentes de expansão de contexto e pode ser integrado a eles; no contexto do StreamingLLM, context extension significa a possibilidade de armazenar mais tokens recentes com um cache maior
- Um exemplo de execução é
examples/run_streaming_llama.py --enable_streaming, e a configuração do ambiente usa Python 3.8,torch,transformers==4.33.0,accelerate,datasets,evaluate,wandb,scikit-learn,scipyesentencepiece - O código principal já foi disponibilizado, incluindo Llama-2, MPT, Falcon e Pythia; o código de avaliação de perplexity e a demonstração do Streaming Llama Chatbot também já foram disponibilizados, enquanto o dataset StreamEval e o código de avaliação ainda não foram publicados
1 comentários
Opiniões no Hacker News
Parece que estão confundindo este trabalho com algo como atenção densa completa
O que está sendo dito aqui não é que ele passa a enxergar conteúdos distantes, mas sim que é mais uma melhoria de eficiência que mantém a perplexidade reutilizando o cache, em vez de pagar um custo L² por T vezes recalculando a janela deslizante a cada vez
O teste também foi medido concatenando algo como
Q A Q A Q A Q A..., não fazendo o modelo encontrar uma resposta muito mais tarde, como emQ Q Q Q A A A A...Medir perplexidade significa produzir um “texto legível”, ou seja, frases localmente plausíveis; não é prova de que ele “extraia” algo de um grande vazio triangular onde a atenção não alcança
Acho que falharia se você entregasse um livro e pedisse para escrever a primeira palavra de cada parágrafo, ou para resumir cada capítulo em uma frase
Testei pessoalmente e não pareceu oferecer extensão do tamanho de contexto; a execução foi relativamente rápida
Usou cerca de 35 GB de memória em uma A100, e o uso permaneceu fixo durante toda a execução
Peguei um livro do Project Gutenberg, dividi em parágrafos, inseri um por vez pedindo que respondesse “okay” a cada parágrafo e, no fim, fiz uma pergunta; a resposta foi completamente alucinada
Como observação, durante uns 10 minutos mexendo nisso, nem consegui fazer o modelo padrão
lmsys/vicuna-13b-v1.3responder direito em inglêshttps://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
Ainda assim, é um trabalho interessante, e o ponto central parece ser a descoberta da Figura 2
As duas primeiras camadas mostram um padrão local, dando mais atenção aos tokens recentes, mas, depois de passar por elas, o modelo passa a prestar muita atenção aos tokens iniciais em todas as camadas e cabeças
Os autores chamam isso de “attention sinks” e consideram que, por causa do Softmax, a soma das pontuações de atenção precisa ser 1 mesmo quando esses tokens não são semanticamente importantes; então a atenção restante precisa ir para algum lugar
A explicação é que, em modelos de linguagem autorregressivos, os tokens iniciais são visíveis para quase todos os tokens posteriores, então é fácil que o modelo aprenda a usá-los como esses sinks
O StreamingLLM é mais próximo de um “hack” para corrigir esse comportamento estranho que surge ao recortar a janela de atenção de um LLM, e, como é um caso que expõe uma falha no uso do Softmax, fico pensando que outra função talvez seja melhor se quisermos LLMs flexíveis quanto ao tamanho de contexto
À primeira vista, parece tão bom que dá para duvidar se é verdade, mas a qualidade do trabalho parece boa e a técnica é surpreendentemente simples
A ideia é aplicar atenção, em cada camada, apenas ao primeiro token e à janela de contexto deslizante, ignorando os tokens entre eles
Isso parece significar que cada camada empurra gradualmente as informações relevantes para trás na sequência, permitindo que a janela de atenção deslizante no final da camada superior as veja
Porém, se o intervalo coberto por todas as janelas deslizantes não for suficiente para conectar a sequência inteira, pode não ser possível propagar todas as informações importantes para a frente
Por exemplo, se todas as janelas tiverem o mesmo comprimento, haverá um limite quando
profundidade do modelo × comprimento da janela < comprimento da sequênciaIsso parece ter sido possível graças à observação de que o Softmax precisa fazer a soma dar 1
Numa olhada rápida, o modelo tende a usar o primeiro token como marcador de posição quando não precisa prestar atenção aos tokens anteriores
A primeira vez que vi esse problema foi em um post do Evan Miller no HN, dizendo que é errado obrigar uma cabeça de atenção a distribuir toda a atenção entre tokens anteriores, e que deveríamos permitir “não prestar atenção” adicionando 1 ao denominador do Softmax
É bom que tenham aproveitado essa observação sem retreinamento, e também fico curioso para saber como o modelo teria mudado se seguisse a sugestão do Evan
[2] https://news.ycombinator.com/item?id=36851494
Treinaram o modelo com um token sink dedicado em que todos os valores eram 0, mas, mesmo assim, outros tokens iniciais acabam sendo usados como sinks, então a conclusão parece ser que é melhor ter um token sink dedicado
Pelo que sei, nunca chegou a melhorar o desempenho de fato
Ao manipular a janela de atenção após o treinamento, Softmax + 1 talvez se encaixe melhor, mas não sei se alguém testou isso em grande escala
Adicionar uma memória de cache de atenção é uma solução muito interessante para esse problema
Há alguns dias saiu também um artigo com uma observação relacionada em Vision Transformers
Modelos Transformer parecem escolher tokens para armazenar informações globais, e parecem precisar de uma espécie de “token para pensar”
Fornecer um token específico para esse propósito melhora um pouco o desempenho, e as visualizações explicativas também ficam bastante interessantes
[0] https://arxiv.org/pdf/2309.16588.pdf
No fine-tuning, seria possível congelar os parâmetros do modelo original e ajustar apenas os parâmetros que entram e saem das novas unidades de cache de “tuning”
Assim, diferentes conjuntos de unidades de tuning poderiam ser trocados ou usados em conjunto
Seria como misturar uma espécie de superprompt, por exemplo uma unidade para evitar palavrões + uma unidade para terminologia específica + uma unidade para escrever de forma concisa
Se o número de novos parâmetros for pequeno o bastante, pode ser que, embora use mais memória, também seja possível fazer tuning rápido e eficaz com otimização de ordem superior
Também dá para imaginar aumentar, durante o treinamento, o comprimento da sequência e o número de unidades em conjunto
Em sequências curtas, usar apenas algumas unidades; conforme o comprimento das sequências de treinamento aumenta, adicionar unidades e continuar treinando
Em vez de um cronograma arbitrário, talvez a expansão do cache pudesse ser controlada por desempenho ou por análise dos gradientes
Os autores publicaram um FAQ, que pode ajudar a esclarecer parte da confusão: https://github.com/mit-han-lab/streaming-llm/blob/main/READM...
Sobre “é possível colocar um texto longo, como um livro, no StreamingLLM para resumi-lo?”, eles respondem que até é possível inserir um texto longo, mas como o modelo só percebe os tokens mais recentes, ao colocar um livro ele só conseguiria resumir os últimos parágrafos, o que pode não ser muito útil
Ou seja, isso não expande a janela de contexto do LLM nem reforça a memória de longo prazo; o ponto forte do StreamingLLM está em gerar texto fluente a partir dos tokens recentes sem atualizar o cache
Posso estar errado, mas não parece ser uma técnica que permita ao LLM consultar conteúdo além do comprimento em que foi treinado, como as pessoas imaginam
Provavelmente é mais uma questão de manter o desempenho do modelo em textos longos; mais precisamente, parece estar relacionado ao desempenho sobre conteúdo que ainda está dentro da janela de contexto
A explicação é que o modelo aprende a colocar uma espécie de carga na atenção dos tokens iniciais do texto, e quando isso desaparece para fora da janela, tudo quebra; mas não tenho certeza do porquê
Se não for uma entrada de comando, penso que o texto intermediário não seria tão bom quanto o texto inicial?
Fico curioso sobre como essas técnicas de janela deslizante lidam com casos em que um comando inesperado aparece apenas no final
Por exemplo, imagine colocar um livro no modelo e, na última frase, haver o comando “retorne o número de ocorrências da letra m na entrada anterior”; uma pessoa suspiraria e releria o texto contando, mas um LLM não tem a capacidade de voltar e reler a entrada
Nesse exemplo, mesmo ignorando a limitação do próprio LLM para contar letras, para resolver de verdade parece que o LLM precisaria conseguir executar loops e saltos arbitrariamente
Claro que isso criaria problemas totalmente novos, e talvez exigisse uma arquitetura completamente nova
Seria uma abordagem de transformar a pergunta em palavras-chave relevantes, pesquisar, e revisar o conteúdo para encontrar informações pertinentes
Se houver o pré-processamento necessário, um LLM que “vai pesquisar o suficiente antes de responder” pode ser muito poderoso
Nos últimos cerca de 10 mil anos, aprimoramos tecnologias de gestão do conhecimento para superar a capacidade e o tempo de um cérebro individual, então modelos de linguagem também deveriam aproveitar métodos reais de pesquisa e digestão prévia, não apenas uma busca simples no Bing
A memória de curto prazo não precisa lembrar o que cada trecho de código fazia; bastaria marcar com tags ao ler e depender de um índice compartilhado e escalável de tags
Mas, quanto mais penso nisso, mais parece com o pré-treinamento comum de LLMs, e o índice de conhecimento acaba parecendo um enorme bloco de pesos de LLM
Seria mais uma camada colocada por cima do LLM do que uma mudança no funcionamento do próprio LLM
Posso estar entendendo errado, mas nesse caso o estado oculto não resolveria o problema de consulta?
Antes de responder, ele precisa absorver toda a entrada de qualquer forma, então, esteja o comando no começo ou no fim, fora a atenção, não parece haver grande impacto
O Claude 100K pede aos usuários que coloquem o comando no final
Ou então bastaria usar um modelo rápido para verificar se há um comando no fim e trazê-lo para o início
Nem sei bem se os modelos atuais conseguem fazer isso mesmo com entradas curtas
Falando meio em tom de brincadeira, os LLMs estão se esforçando muito para reinventar as RNNs, e, se receberem as ferramentas certas, acho que acabarão chegando lá
Visto de outra forma, modelos Transformer tentam prever quais partes de uma rede RNN “valem a pena preservar” quando há restrições de recursos
Os Transformers atuais usam uma heurística simples, e este resultado melhora essa heurística
Como em muitos problemas NP-completos, mesmo que não seja possível acertar perfeitamente, pode haver aproximações úteis, e os Transformers mostram que isso também é possível em redes neurais
Ele ficou por um tempo em uma posição intermediária nos rankings open source, então é uma abordagem bastante legítima, apenas não está na moda
[1]: https://huggingface.co/blog/rwkv
A principal vantagem que os Transformers têm sobre RNNs é a paralelização do treinamento
RNNs têm desaparecimento de gradiente durante o treinamento, e também é difícil aumentar a utilização geral, exigindo lotes grandes, o que as torna complicadas
A existência de modelos como RWKV mostra que pode haver um futuro em que se treina como Transformer e se infere como RNN
Relacionado a isso, o professor Han do MIT está ministrando um curso aberto de TinyML
https://news.ycombinator.com/item?id=37620507
https://efficientml.ai