[論文レビュー] Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems

Antonio Ginart, Maxim Naumov+3 more|arXiv (Cornell University)|Sep 25, 2019

Recommender Systems and Techniques参考文献 77被引用数 32

ひとこと要約

本稿では、特徴の頻度に応じて埋め込みベクトルの次元数を可変的に設定する混合次元（MD）埋め込みを提案する。頻度の高いアイテムに多くのパラメータを割り当て、まれなアイテムに少ないパラメータを割り当てることで、Criteoデータセットにおいて精度に損失を生じさせることなくメモリ使用量を最大16倍まで削減し、GPU上で学習を2倍以上高速化する。

ABSTRACT

Embedding representations power machine intelligence in many applications, including recommendation systems, but they are space intensive -- potentially occupying hundreds of gigabytes in large-scale settings. To help manage this outsized memory consumption, we explore mixed dimension embeddings, an embedding layer architecture in which a particular embedding vector's dimension scales with its query frequency. Through theoretical analysis and systematic experiments, we demonstrate that using mixed dimensions can drastically reduce the memory usage, while maintaining and even improving the ML performance. Empirically, we show that the proposed mixed dimension layers improve accuracy by 0.1% using half as many parameters or maintain it using 16X fewer parameters for click-through rate prediction task on the Criteo Kaggle dataset.

研究の動機と目的

• 大規模な推薦システムにおける埋め込み層の高いメモリ消費量（モデル保存領域の99.9％以上を占める）を軽減すること。
• 低頻度の特徴に無駄に多くのパラメータを割り当て、高頻度の特徴に不足したパラメータを割り当てる、均一な埋め込み次元の非効率性を是正すること。
• 特徴のクエリ頻度に応じて動的に埋め込み次元を調整する、パラメータ効率の高い埋め込みアーキテクチャを開発すること。
• 歪んだデータ分布下でも、メモリ効率とモデルの一般化性能の両面でMD埋め込みが向上することを理論的および実験的に検証すること。
• 従来の非均一埋め込み手法と比較して、より高速な学習と少ないチューニングの複雑さを達成すること。

提案手法

• 各埋め込みベクトルの次元数を固定された均一な次元ではなく、クエリ頻度に応じて設定する混合次元（MD）埋め込み層を提案する。
• 全パラメータ予算を満たす条件下で再構成歪みを最小化するように、埋め込み次元選択を凸最適化問題として定式化する。
• 埋め込み行列を、それぞれがカテゴリカル特徴のペアに対応するブロックの和としてモデル化するための低ランク近似フレームワークを用いる。
• ラグランジュ緩和を用いて最適な次元割り当てルールを導出する。最適な次元は、その特徴ペアの頻度、スぺクトル減衰、およびパラメータ予算に依存する。
• べき乗則に従うスぺクトル減衰の場合、閉形式解を導出する：$ d^{*}_{ij} = \lambda \zeta_{ij} \Pi_{ij}^{1/(2\beta)} $、ここで $ \lambda $ はラグランジュ乗数、$ \zeta_{ij} $ は特徴サイズとスぺクトル減衰に依存する。
• 実装可能性を保証するため、連続的な最適次元を整数に切り下げ、近似誤差が小さい（通常1–10％のスぺクトル質量損失）ことを証明する。

Figure 1 : Matrix Architecture for UD and MD Embedding Layers.

実験結果

リサーチクエスチョン

• RQ1 特徴の頻度に基づいて埋め込み次元を動的に割り当てることで、モデル性能に損失を生じさせずにメモリ使用量を削減できるか？
• RQ2 メモリ制約下において、再構成歪みと一般化性能の観点から、MD埋め込みは均一次元埋め込みを上回るか？
• RQ3 各特徴の頻度と全体のパラメータ予算が与えられた場合、最適な埋め込み次元は何か？
• RQ4 従来の非均一埋め込みアーキテクチャと比較して、本手法は学習速度とチューニングの複雑さにおいてどのように差をつけるか？
• RQ5 スぺクトル減衰や頻度の偏りなどの条件下で、MD埋め込みは均一埋め込みに対して明確な改善効果を示すか？

主な発見

• Criteo Kaggleデータセットにおいて、MD埋め込みは均一埋め込みと比較して0.1％高い精度を達成しながら、パラメータ数を半分に削減した。
• MD埋め込みは、均一埋め込みと同等の精度を維持しながら、パラメータ数を最大16倍まで削減し、強力なパラメータ効率を示した。
• MD埋め込みを用いた学習は、GPU上で2倍以上高速化された。これは、メモリ帯域幅の圧力を軽減し、計算がより効率的になったためである。
• 理論的分析により、データがメモリ制限下にあり、頻度分布が十分に偏っている場合、MD埋め込みは均一埋め込みよりも低い再構成歪みを生じることが証明された。
• べき乗則に従うスぺクトル減衰の場合、最適な埋め込み次元は $ \Pi_{ij}^{1/(2\beta)} $ に比例して増加する。ここで $ \Pi_{ij} $ は特徴ペアの頻度を表す。
• 連続的な最適次元を整数に切り下げた際の近似誤差は小さく（通常1–10％のスぺクトル質量損失）、実用的な実装可能性が保証された。

Figure 2 : CTR prediction results for MD embeddings on Criteo dataset using DLRM. Implementation is available as part of an open-source project on GitHub: facebookresearch/dlrm . Fig. 2a (left): Learning curves for selected emb. arch. Fig. 2b (center): Loss vs. # param. for varying $\alpha$ . Fig 2c