[論文レビュー] Extreme Classification in Log Memory using Count-Min Sketch: A Case Study of Amazon Search with 50M Products

Tharun Medini, Qixuan Huang+3 more|arXiv (Cornell University)|Jan 1, 2019

Advanced Image and Video Retrieval Techniques被引用数 23

ひとこと要約

本論文は、ユニバーサルハッシュを用いたCount-Min Sketchを採用することで、大規模な分類におけるメモリスケーリングをO(K)からO(log K)に削減する、MACHと呼ばれる画期的な極端分類フレームワークを紹介する。4946万件の製品を含むAmazonの検索データセット上で訓練されたMACHは、64億パラメータを有し、35時間未塔の単一のp3.16xインスタンス上で、最先端の精度と再現率を達成した。これは、従来手法と比較して7〜10倍速く、2〜4倍メモリ効率が良かった。

ABSTRACT

In the last decade, it has been shown that many hard AI tasks, especially in NLP, can be naturally modeled as extreme classification problems leading to improved precision. However, such models are prohibitively expensive to train due to the memory bottleneck in the last layer. For example, a reasonable softmax layer for the dataset of interest in this paper can easily reach well beyond 100 billion parameters (> 400 GB memory). To alleviate this problem, we present Merged-Average Classifiers via Hashing (MACH), a generic $K$-classification algorithm where memory provably scales at $O(\log K)$ without any assumption on the relation between classes. MACH is subtly a count-min sketch structure in disguise, which uses universal hashing to reduce classification with a large number of classes to few embarrassingly parallel and independent classification tasks with a small (constant) number of classes. MACH naturally provides a technique for zero communication model parallelism. We experiment with 6 datasets; some multiclass and some multilabel, and show consistent improvement in precision and recall metrics compared to respective baselines. In particular, we train an end-to -end deep classifier on a private product search dataset sampled from Amazon Search Engine with 70 million queries and 49.46 million documents. MACH outperforms, by a significant margin, the state-of-the-art extreme classification models deployed on commercial search engines: Parabel and dense embedding models. Our largest model has 6.4 billion parameters and trains in less than 35 hrs on a single p3.16x machine. Our training times are 7-10x faster, and our memory footprints are 2-4x smaller than the best baselines. This training time is also significantly lower than the one reported by Google’s mixture of experts (MoE) language model on a comparable model size and hardware.

研究の動機と目的

• 数百万クラスの分類において、ソフトマックス層が数百ギガバイトのメモリを要するという、極端分類におけるメモリボトルネックを解消すること。
• クラス間の関係に関する仮定を一切設けず、スケーラブルでメモリ効率の良い分類フレームワークを構築すること。
• 大規模分類を独立した小規模クラス問題に還元することで、通信を要しないモデル並列化を可能にする。
• 特に5000万クラス以上のクラスを有する製品検索を含む、実世界の産業スケールのデータセットにおいて、優れた性能と効率を示すこと。

提案手法

• MACHは、ユニバーサルハッシュを用いて、高次元の分類タスクを一連のより小さな独立した分類問題にマッピングするCount-Min Sketch構造を採用する。
• 各ハッシュ関数は、元のKクラス問題を定数サイズの部分問題に変換し、並列性の高いトレーニングと推論を可能にする。
• 最終的な予測は、すべてのハッシュベースの部分分類器からの予測の平均値を統合することで計算され、モデルの表現力は維持されつつ、メモリ使用量が削減される。
• ユニバーサルハッシュとスケッチ技術の集中性の性質を活用することで、理論的にO(log K)のスケーリングが保証される。
• MACHは、極端出力層を備えた深層ニューラルネットワークのエンドツーエンドトレーニングを可能とし、標準的なディープラーニングフレームワークとシームレスに統合できる。
• 各部分分類器がパラメータ同期なしに独立してトレーニング可能であるため、通信を要しないモデル並列化が実現可能である。

実験結果

リサーチクエスチョン

• RQ1 クラス数に比例しない、具体的にはO(log K)のメモリ複雑性で極端分類が可能になるか。この際、モデルの精度を損なわずに行えるか。
• RQ2 MACHは、Parabel や密度埋め込み手法といった最先端モデルと比較して、大規模製品検索データにおいて、精度、再現率、トレーニング効率の点でどのように差をつけるか。
• RQ3 4900万クラスを超えるデータセットに対しても、MACHは低メモリ使用量と高速トレーニングを維持できるか。
• RQ4 ユニバーサルハッシュを用いたCount-Min Sketchの使用は、極端分類における効果的な通信ゼロのモデル並列化を可能にするか。
• RQ5 同程度のハードウェアとモデルサイズにおいて、MACHはGoogleのMixture-of-Expertsモデルよりも高速なトレーニングと低いメモリ使用量を達成できるか。

主な発見

• MACHは、7000万件のクエリと4946万件のドキュメントを含むプライベートなAmazon製品検索データセットで、最先端の精度と再現率を達成した。
• 64億パラメータを有する最大のMACHモデルは、単一のp3.16xインスタンス上で35時間未塔でトレーニングが完了し、ベースラインモデルと比較して著しく高速かつ効率的であった。
• MACHは、Parabel や密度埋め込みモデルを含む最良の既存ベースラインと比較して、メモリ使用量を2〜4倍削減し、トレーニング時間を7〜10倍短縮した。
• MACHのトレーニング時間は、同程度のモデルサイズとハードウェア構成におけるGoogleのMixture-of-Experts言語モデルの報告時間よりも顕著に短かった。
• MACHは、マルチクラスおよびマルチラベル設定を含む6つの多様なデータセットにおいて一貫した改善を示し、その汎用性と頑健性を裏付けた。
• 理論的メモリスケーリングO(log K)は、実験的にも裏付けられ、クラス構造に関する仮定を一切設けず、極端なクラス数に対してもスケーラブルであることが確認された。