Memoria de Claves de Producto de Peso Rápido

Resumen

Las capas de modelado de secuencias en los modelos de lenguaje modernos generalmente enfrentan una disyuntiva entre capacidad de almacenamiento y eficiencia computacional. Mientras que la atención Softmax ofrece almacenamiento ilimitado con costos cuadráticos prohibitivos, las variantes lineales proporcionan eficiencia pero adolecen de un almacenamiento limitado y de tamaño fijo. Proponemos Fast-weight Product Key Memory (FwPKM), una arquitectura novedosa que resuelve esta tensión transformando la memoria de claves de producto (PKM) de un módulo estático en una memoria episódica dinámica de "pesos rápidos". A diferencia de PKM, FwPKM actualiza sus parámetros dinámicamente tanto durante el entrenamiento como en la inferencia mediante descenso de gradiente a nivel de fragmentos locales, permitiendo que el modelo memorice y recupere rápidamente nuevos pares clave-valor de las secuencias de entrada. Los experimentos revelan que FwPKM funciona como una memoria episódica efectiva que complementa la memoria semántica de los módulos estándar, produciendo reducciones significativas de perplejidad en conjuntos de datos de contexto largo. Notablemente, en evaluaciones de Aguja en un Pajar, FwPKM generaliza a contextos de 128K tokens a pesar de haber sido entrenado únicamente con secuencias de 4K 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.