Memoria a Chiavi Prodotto a Pesi Rapidi

Abstract

Gli strati di modellazione sequenziale nei moderni modelli linguistici affrontano tipicamente un compromesso tra capacità di memorizzazione ed efficienza computazionale. Mentre l'attenzione Softmax offre una capacità di memorizzazione illimitata a costi quadratici proibitivi, le varianti lineari garantiscono efficienza ma soffrono di una capacità di memorizzazione limitata e di dimensioni fisse. Proponiamo Fast-weight Product Key Memory (FwPKM), un'architettura innovativa che risolve questa tensione trasformando la Product Key Memory (PKM) sparsa da un modulo statico in una memoria episodica dinamica a "pesi rapidi". A differenza della PKM, FwPKM aggiorna dinamicamente i propri parametri sia durante l'addestramento che l'inferenza attraverso una discesa del gradiente locale a livello di segmento, consentendo al modello di memorizzare e recuperare rapidamente nuove coppie chiave-valore dalle sequenze di input. Gli esperimenti rivelano che FwPKM funziona come una memoria episodica efficace che complementa la memoria semantica dei moduli standard, producendo riduzioni significative della perplexity su dataset a contesto lungo. Significativamente, nelle valutazioni "Needle in a Haystack", FwPKM generalizza a contesti di 128K token nonostante sia stato addestrato solo su sequenze di 4K token.

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.