[論文レビュー] Training with Multi-Layer Embeddings for Model Reduction

Benjamin Ghaemmaghami, Zihao Deng+5 more|arXiv (Cornell University)|Jun 10, 2020

Recommender Systems and Techniques参考文献 24被引用数 5

ひとこと要約

本稿では、単一の埋め込みテーブルを二層の線形アーキテクチャに置き換えることで、精度対モデルサイズのトレードオフを向上させる、マルチレイヤー埋め込み学習（MLET）を提案する。幅広い隠れ層（k）をまず学習し、その後で最終的な埋め込み次元（d）に射影することで、学習された表現の有効ランクを低く抑え、精度を落とさずに埋め込みテーブルサイズを4–8倍小さくできる。訓練時間は平均で25%増加するが、これに伴うコストは最小限である。

ABSTRACT

Modern recommendation systems rely on real-valued embeddings of categorical features. Increasing the dimension of embedding vectors improves model accuracy but comes at a high cost to model size. We introduce a multi-layer embedding training (MLET) architecture that trains embeddings via a sequence of linear layers to derive superior embedding accuracy vs. model size trade-off. Our approach is fundamentally based on the ability of factorized linear layers to produce superior embeddings to that of a single linear layer. We focus on the analysis and implementation of a two-layer scheme. Harnessing the recent results in dynamics of backpropagation in linear neural networks, we explain the ability to get superior multi-layer embeddings via their tendency to have lower effective rank. We show that substantial advantages are obtained in the regime where the width of the hidden layer is much larger than that of the final embedding (d). Crucially, at conclusion of training, we convert the two-layer solution into a single-layer one: as a result, the inference-time model size scales as d. We prototype the MLET scheme within Facebook's PyTorch-based open-source Deep Learning Recommendation Model. We show that it allows reducing d by 4-8X, with a corresponding improvement in memory footprint, at given model accuracy. The experiments are run on two publicly available click-through-rate prediction benchmarks (Criteo-Kaggle and Avazu). The runtime cost of MLET is 25%, on average.

研究の動機と目的

• 推薦システムにおける高次元埋め込みテーブルの増大するメモリコストに対処すること。
• 推論アーキテクチャを変更せずに、ディープラーニングベースの推薦モデルにおける精度対モデルサイズのトレードオフを改善すること。
• 多層線形変換が単一層のテーブルよりも優れた埋め込みを生成できるかどうかを検証すること。
• 実世界の性能を評価するために、FacebookのオープンソースDRLMフレームワーク内にMLETのプロトタイプを構築し、評価すること。

提案手法

• MLETは、単一の埋め込みテーブル W ∈ ℝⁿˣᵈ を、qW₁W₂ という二層の線形変換に置き換える。ここで W₁ ∈ ℝⁿˣᵏ および W₂ ∈ ℝᵏˣᵈ である。
• モデルはバックプロパゲーションを用いてエンドツーエンドで学習され、最終的な推論では d 次元出力を得るために W₂ のみが使用される。
• このアプローチは、線形ネットワークダイナミクスの理論的知見を活用しており、より深い線形ネットワークは有効ランクが低くなる傾向があり、一般化性能が向上する。
• 本手法は、Criteo-KaggleおよびAvazuデータセットを用いて、FacebookのPyTorchベースのDLRMフレームワーク内で実装および評価された。
• 訓練時のメモリは、単一層学習と比較して k/d の割合で増加するが、推論時のメモリは d に比例する。
• 最終的なモデルは推論時に単一層形式に変換され、既存のシステムとの互換性が保たれる。

実験結果

リサーチクエスチョン

• RQ1 多層線形アーキテクチャは、精度とモデルサイズの観点で単一層埋め込みを上回る埋め込みを生成できるか？
• RQ2 なぜ多層線形ネットワークは単一層よりも優れた埋め込みを生成するのか？その背後にあるダイナミクスは何か？
• RQ3 MLETを用いることで、モデル精度を維持または向上させながら、どの程度埋め込みテーブルサイズを削減できるか？
• RQ4 訓練メモリコストが増加しない範囲で、MLETが精度とモデルサイズの両方を改善する設定は存在するか？

主な発見

• Criteo-KaggleおよびAvazuの両データセットにおいて、MLETは単一層埋め込みと同等のモデル精度を維持しながら、埋め込みテーブルサイズを4–8倍小さくした。
• MLETの最大LogLoss改善は、Criteo-Kaggleで0.0025、Avazuで0.006であった。kが一定値を超えて増加すると性能は飽和する。
• 特定のk/d比において、MLETは追加の訓練メモリコストなしに精度向上とモデルサイズ削減の両方を達成した（すなわち、k ≤ d）。この現象はAvazuでは観察されたが、Criteo-Kaggleでは観察されなかった。
• MLETの訓練時間は、単一層学習と比較して平均で25%長くなったが、推論時間は標準的なDLRMと同一であった。
• 性能向上の主な要因はk/d比であり、k/dが高くなるほど精度が単調に向上した。
• k ≤ d を用いることで、コストフリーの利点が得られるため、dが小さい場合でも効果的であり、モデルサイズの相対的改善が最も顕著に現れた。