Um Novo Método para Reduzir o Uso de Memória em Grandes Modelos de Linguagem

Modelos de linguagem grandes (LLMs) conseguem fazer uma porção de tarefas, mas geralmente precisam de muita memória. Esse uso alto de memória pode dificultar a implementação deles em situações do mundo real, onde a eficiência é fundamental. Uma grande parte do uso de memória vem de algo chamado cache de chave-valor (KV), que é essencial para o mecanismo de atenção usado nesses modelos. Quanto mais camadas no modelo, mais memória é necessária para esse Cache KV.

Nessa conversa, vamos falar sobre um novo método que ajuda a reduzir a memória necessária para os LLMs, permitindo que eles funcionem mais rápido enquanto usam menos memória. Nossa abordagem foca em cachear apenas os pares KV de um número pequeno de camadas, reduzindo significativamente o uso de memória e aumentando a velocidade de inferência.

#Requisitos de Memória dos LLMs

O alto uso de memória tem sido um problema significativo para os LLMs. O consumo de memória pode limitar quantas tarefas os modelos conseguem lidar de uma vez. Especificamente, o cache KV pode ocupar mais de 30% da memória necessária durante a implementação. As necessidades de memória crescem com base tanto no comprimento dos dados sendo processados quanto no número de camadas em um modelo.

Várias estratégias já foram introduzidas para diminuir a carga de memória do cache KV. Algumas técnicas têm como objetivo encurtar o comprimento das sequências KV cacheadas. Por exemplo, alguns métodos focam em manter apenas os KVs dos tokens mais recentes ou usar um processo para remover dinamicamente KVs mais antigos da memória quando não são mais necessários.

#Nossa Abordagem para Reduzir o Uso de Memória

Em vez de trabalhar para encurtar as sequências KV, nosso método foca em reduzir o número de camadas que precisam ter seus KVs computados e cacheados. A gente propõe um tipo de decodificador transformer onde todas as camadas usam as chaves e valores apenas da camada de cima. Essa abordagem nos permite cachear KVs de apenas uma camada em vez de várias, resultando em uma economia significativa de memória sem perder Desempenho.

A ideia também se baseia em como as camadas em um transformer se constroem umas sobre as outras. A camada de cima é frequentemente vista como a que contém as informações mais valiosas, fazendo sentido para as camadas inferiores dependerem dela em vez de precisarem dos seus próprios KVs.

#Avaliação de Desempenho

Testamos nosso método em modelos de linguagem grandes e encontramos um aumento impressionante no desempenho. Nossa abordagem mostrou melhorar a velocidade de processamento em até 26% em comparação com modelos tradicionais, enquanto ainda performa bem em tarefas de compreensão de linguagem.

Além disso, nosso método funciona bem com outras técnicas existentes projetadas para economizar memória. Isso significa que os usuários podem combinar nosso método com outros para obter um desempenho e eficiência ainda melhores.

#Desafios com o Treinamento

Embora nosso método traga ótimos resultados para inferência, ele cria algumas dificuldades durante o treinamento. Em um modelo transformer padrão, a computação pode ser feita em paralelo. No entanto, como nosso modelo requer que cada token dependa dos KVs da camada de cima do token anterior, isso cria uma sequência de dependências que limita o processamento paralelo.

Para resolver isso, desenvolvemos um novo método de treinamento que permite algum nível de paralelização. Ao permitir que os tokens computem seus valores em iterações, conseguimos reduzir o tempo de treinamento necessário mantendo a estrutura de dependência intacta.

#Resultados Experimentais

Nos nossos testes com vários modelos, descobrimos que nossa abordagem permite tamanhos de lote maiores e velocidades de processamento mais rápidas. Mesmo com modelos que têm de 1 bilhão a 30 bilhões de parâmetros, nosso método ainda mostra um desempenho competitivo em várias tarefas.

Destacamos especificamente a situação onde nosso modelo pode se integrar facilmente com outros métodos de economia de memória, melhorando ainda mais a eficiência em uma série de tarefas.

#Explicação Técnica do Modelo

Para alcançar nossos objetivos, redesenhamos como a atenção funciona em nosso modelo. Com base em nossa abordagem, cada nó representando uma camada na computação transformer só interage com os KVs da camada de cima. Essa mudança sozinha minimiza significativamente a memória necessária para os cálculos.

Um problema que encontramos foi que cada token ainda precisa prestar atenção a si mesmo na computação, criando um loop nas dependências. Resolvemos isso permitindo que o modelo ignore temporariamente a auto-at