Быстрообучаемая продуктовая ключевая память

Аннотация

Слои моделирования последовательностей в современных языковых моделях обычно сталкиваются с компромиссом между объемом памяти и вычислительной эффективностью. В то время как Softmax-внимание предлагает неограниченную память, но за запретительно высокую квадратичную стоимость, линейные варианты обеспечивают эффективность, но страдают от ограниченной памяти фиксированного размера. Мы предлагаем Fast-weight Product Key Memory (FwPKM) — новую архитектуру, которая разрешает это противоречие, преобразуя разреженную Продуктовую Память Ключей (PKM) из статичного модуля в динамическую, «быструю» эпизодическую память. В отличие от PKM, FwPKM динамически обновляет свои параметры как во время обучения, так и во время вывода с помощью локального градиентного спуска на уровне чанков, что позволяет модели быстро запоминать и извлекать новые пары «ключ-значение» из входных последовательностей. Эксперименты показывают, что FwPKM функционирует как эффективная эпизодическая память, дополняющая семантическую память стандартных модулей, и приводит к значительному снижению перплексии на наборах данных с длинным контекстом. Примечательно, что в оценках «Иголка в стоге сена» FwPKM обобщается на контексты длиной 128 тыс. токенов, несмотря на обучение на последовательностях длиной всего 4 тыс. токенов.

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.