Memória de Chaves de Produto com Peso Rápido

Resumo

As camadas de modelagem de sequência em modelos de linguagem modernos geralmente enfrentam um compromisso entre capacidade de armazenamento e eficiência computacional. Enquanto a atenção Softmax oferece armazenamento ilimitado a custos quadráticos proibitivos, variantes lineares proporcionam eficiência, mas sofrem com armazenamento limitado e de tamanho fixo. Propomos a Memória de Chaves de Produto com Pesos Rápidos (FwPKM), uma arquitetura inovadora que resolve essa tensão ao transformar a esparsa Memória de Chaves de Produto (PKM) de um módulo estático em uma memória episódica dinâmica de "pesos rápidos". Diferente da PKM, a FwPKM atualiza seus parâmetros dinamicamente, tanto durante o treinamento quanto na inferência, por meio de gradiente descendente local a nível de blocos (chunks), permitindo que o modelo memorize e recupere rapidamente novos pares chave-valor a partir de sequências de entrada. Experimentos revelam que a FwPKM funciona como uma memória episódica eficaz que complementa a memória semântica dos módulos padrão, resultando em reduções significativas de perplexidade em conjuntos de dados de contexto longo. Notavelmente, em avaliações do tipo "Agulha no Palheiro" (Needle in a Haystack), a FwPKM generaliza para contextos de 128 mil tokens apesar de ter sido treinada apenas em sequências de 4 mil tokens.

English

Sequence modeling layers in modern language models typically face a trade-off between storage capacity and computational efficiency. While Softmax attention offers unbounded storage at prohibitive quadratic costs, linear variants provide efficiency but suffer from limited, fixed-size storage. We propose Fast-weight Product Key Memory (FwPKM), a novel architecture that resolves this tension by transforming the sparse Product Key Memory (PKM) from a static module into a dynamic, "fast-weight" episodic memory. Unlike PKM, FwPKM updates its parameters dynamically at both training and inference time via local chunk-level gradient descent, allowing the model to rapidly memorize and retrieve new key-value pairs from input sequences. Experiments reveal that FwPKM functions as an effective episodic memory that complements the semantic memory of standard modules, yielding significant perplexity reductions on long-context datasets. Notably, in Needle in a Haystack evaluations, FwPKM generalizes to 128K-token contexts despite being trained on only 4K-token sequences.