[論文レビュー] Large Memory Layers with Product Keys

Guillaume Lample, Alexandre Sablayrolles+3 more|HAL (Le Centre pour la Communication Scientifique Directe)|Jul 10, 2019

Topic Modeling被引用数 50

ひとこと要約

スケーラブルなメモリ拡張レイヤーをプロダクトキーに基づいて導入し、正確な最近傍探索を可能にする。大規模言語モデル用のトランスフォーマーに統合され、深いベースラインよりも推論が速く、精度が改善される。

ABSTRACT

This paper introduces a structured memory which can be easily integrated into a neural network. The memory is very large by design and significantly increases the capacity of the architecture, by up to a billion parameters with a negligible computational overhead. Its design and access pattern is based on product keys, which enable fast and exact nearest neighbor search. The ability to increase the number of parameters while keeping the same computational budget lets the overall system strike a better trade-off between prediction accuracy and computation efficiency both at training and test time. This memory layer allows us to tackle very large scale language modeling tasks. In our experiments we consider a dataset with up to 30 billion words, and we plug our memory layer in a state-of-the-art transformer-based architecture. In particular, we found that a memory augmented model with only 12 layers outperforms a baseline transformer model with 24 layers, while being twice faster at inference time. We release our code for reproducibility purposes.

研究の動機と目的

• 大きな学習可能なメモリ層を導入することにより、比例的な計算コストを伴わずにモデル容量を増加させる。
• プロダクトキーを用いたスケーラブルなキー-バリューメモリ内で厳密な最近傍探索を提供する。
• 大規模言語モデルタスクでのトランスフォーマー内におけるメモリ層の有効性を示す。
• メモリ拡張モデルが、計算量が同等または少なくても、より深いベースラインを上回ることを示す。
• メモリサイズ、ネットワーク内の配置、正規化が性能と使用状況に与える影響を分析する。

提案手法

• クエリネットワーク、2部構成のサブキーを持つプロダクトキー空間、値のルックアップテーブルを備えたメモリを定義する。
• 2つのコードブックの外積によって形成されるプロダクトキーを用いて、大規模なキー集合を暗黙的に定義する。
• クエリとの内積類似度で上位k個のキーを選択し、対応する値の重み付き和を計算する（スパースリード）。
• 各ヘッドが独自のクエリネットワークとサブキーを持ち、値を共有するマルチヘッドメモリアテンションを採用する。
• クエリネットワークに対してスパースメモリ更新とバッチ正規化を用いて、キーのカバレッジと使用を改善する。
• 選択されたトランスフォーマーレイヤーのFFNをメモリ層に置換し、加法的な残差接続（x → x + PKM(x)）を可能にする。

実験結果

リサーチクエスチョン

• RQ1 プロダクトキー・メモリ（PKM）は、厳密な探索を維持しつつ、ほとんどオーバーヘッドなしでスケーラブルな容量を提供しますか？
• RQ2 メモリサイズ、ヘッド数、k-NN、メモリ配置は、大規模トランスフォーマーにおけるパープレキシティ、メモリ使用量、速度にどのような影響を与えますか？
• RQ3 パットキー・メモリとフラットキー・メモリを、性能・使用状況・効率の観点でどのように比較しますか？
• RQ4 大規模言語モデルデータセットで、メモリ使用とモデル性能を最適化するトレーニング設定（BatchNorm、スパース更新など）は何ですか？

主な発見

• 同じ次元の12層メモリとプロダクトキーを用いたモデルは、CC-Newsで同じ次元の24層トランスフォーマよりも性能が高く、推論時間は約半分。
• メモリサイズを増やすと一貫してパープレキシティが低下し、メモリサイズが100万スロットの場合、バッチ正規化がキーの利用を改善すると顕著な向上を達成する。
• プロダクトキーは、同等サイズでフラットキーよりもメモリ使用量が高く、パープレキシティが低く、推論が速い。
• メモリ配置は重要で、トランスフォーマー中間層（およそ4–5付近）にメモリを挿入すると最良の性能を得られる。
• 4つのメモリヘッド、各ヘッド32キー、総スロット512^2は、速度とパープレキシティの間で有利なトレードオフをもたらす。プロダクトキーはフラットキーに比べてパラメータ数を削減し、推論時のスケーリングもより効率的。
• 大規模メモリにおいて、クエリネットワークのバッチ正規化はキーの利用とパープレキシティを著しく改善する。