Médias Chave-Valor

Resumo

Apresentamos Key-Value Means ("KVM"), uma nova bloco-recursão para atenção que pode acomodar tanto estados de tamanho fixo quanto crescentes. Equipar uma forte linha de base de transformador com camadas de atenção KVM de tamanho fixo resulta em uma robusta RNN fragmentada O(N), adicionando apenas um número insignificante de novos parâmetros. Treinamos um transformador com um cache KVM expansível e mostramos que seu desempenho é competitivo em testes de contexto longo, com tempo de pré-preenchimento subquadrático e crescimento de estado sublinear. O KVM é implementável com operações padrão e sem kernels personalizados, e suporta treinamento e pré-preenchimento paralelizáveis por fragmento. Ele oferece muitos dos benefícios tanto dos transformadores tradicionais (memória de contexto expansível, treinamento e pré-preenchimento paralelizáveis por fragmento) quanto das RNNs lineares em um único pacote unificado. Pode ser usado em todas as camadas, economizando memória do cache KV, e permite uma gama contínua de escolhas de complexidade de tempo de pré-preenchimento entre O(N) e O(N²). Também pode ser implementado em uma solução híbrida em conjunto com camadas LRNN no lugar da atenção tradicional, para complementar a LRNN com melhor uso do comprimento de contexto de crescimento de memória sublinear e decodificação de contexto longo. Disponibilizamos nosso código em https://github.com/recursal/KVM-paper e modelos treinados em https://huggingface.co/collections/recursal/key-value-means sob a licença Apache 2.0.

English

We present Key-Value Means ("KVM"), a novel block-recurrence for attention that can accommodate either fixed-size or growing state. Equipping a strong transformer baseline with fixed-size KVM attention layers yields a strong O(N) chunked RNN, while adding only an insignificant number of new parameters. We train a transformer with a growable KVM cache and show it performs competitively on long-context tests with only subquadratic prefill time and sublinear state growth. KVM is implementable with standard operations and without custom kernels, and supports chunk-wise parallelizable training and prefill. It provides many of the benefits of both traditional transformers (expandable context memory, chunk-wise parallelizable training and prefill) and linear RNNs in a single unified package. It can be used on every layer, saving KV-cache memory, and allowing a continuous range of choices of prefill time complexity between O(N) and O(N^2). It can also be implemented in a hybrid solution in tandem with LRNN layers in place of traditional attention, to supplement the LRNN with improved sublinear memory growth context length usage and long context decoding. We release our code at https://github.com/recursal/KVM-paper and trained models at https://huggingface.co/collections/recursal/key-value-means under the Apache 2.0 license.