高速重み付けプロダクトキーメモリ

要旨

現代の言語モデルにおける系列モデリング層は、一般的に記憶容量と計算効率の間でトレードオフに直面しています。Softmaxアテンションは二次コストが膨大になる代わりに無制限の記憶容量を提供する一方、線形変種は効率性を提供するものの、限定的で固定サイズの記憶容量に悩まされています。本論文では、この対立を解決する新しいアーキテクションであるFast-weight Product Key Memory（FwPKM）を提案します。FwPKMは、静的なモジュールであるProduct Key Memory（PKM）を動的な「高速重み」エピソード記憶へと変換します。PKMとは異なり、FwPKMはトレーニング時と推論時の両方で、局所的なチャンクレベルの勾配降下法を通じてそのパラメータを動的に更新し、モデルが入力系列から新しいキーと値のペアを迅速に記憶し、検索することを可能にします。実験結果から、FwPKMは標準モジュールの意味記憶を補完する効果的なエピソード記憶として機能し、長文脈データセットにおいて大幅なパープレキシティの低減をもたらすことが明らかになりました。特に、「干し草の山の中の針」評価において、FwPKMは4Kトークンの系列でトレーニングされたのみにも関わらず、128Kトークンの文脈へと一般化することが示されました。

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.