[論文レビュー] TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models

Chunxing Yin, Bilge Acun+2 more|arXiv (Cornell University)|Jan 25, 2021

Tensor decomposition and applications参考文献 32被引用数 13

ひとこと要約

TT-Recは、テンソルトレース分解に基づく手法を提案し、深層学習推薦モデル（DLRMs）の埋め込みテーブルを圧縮する。モデルサイズを最大117倍まで削減可能であり、精度に損なわれることなく、学習時間のオーバーヘッドもたった13.9%に抑えられる。最適化されたTT-EmbeddingBagカーネル、適応的キャッシュ、および新規のガウス分布に基づく初期化を用いることで、大規模なスケーリングにおいても性能と精度を維持する。

ABSTRACT

The memory capacity of embedding tables in deep learning recommendation models (DLRMs) is increasing dramatically from tens of GBs to TBs across the industry. Given the fast growth in DLRMs, novel solutions are urgently needed, in order to enable fast and efficient DLRM innovations. At the same time, this must be done without having to exponentially increase infrastructure capacity demands. In this paper, we demonstrate the promising potential of Tensor Train decomposition for DLRMs (TT-Rec), an important yet under-investigated context. We design and implement optimized kernels (TT-EmbeddingBag) to evaluate the proposed TT-Rec design. TT-EmbeddingBag is 3 times faster than the SOTA TT implementation. The performance of TT-Rec is further optimized with the batched matrix multiplication and caching strategies for embedding vector lookup operations. In addition, we present mathematically and empirically the effect of weight initialization distribution on DLRM accuracy and propose to initialize the tensor cores of TT-Rec following the sampled Gaussian distribution. We evaluate TT-Rec across three important design space dimensions -- memory capacity, accuracy, and timing performance -- by training MLPerf-DLRM with Criteo's Kaggle and Terabyte data sets. TT-Rec achieves 117 times and 112 times model size compression, for Kaggle and Terabyte, respectively. This impressive model size reduction can come with no accuracy nor training time overhead as compared to the uncompressed baseline.

研究の動機と目的

• 深層学習推薦モデル（DLRMs）の急速に増大するメモリ要件に対処する。現在、これらのモデルはTBスケールの埋め込みテーブルを必要としている。
• 精度を損なわず、学習時間のオーバーヘッドが著しく増大するのを防ぐことにより、モデルのメモリ容量を削減する。
• 埋め込みのパラメータ化された、原理的かつ体系的な圧縮手法としてのテンソル化を検討し、モデル品質を維持する。
• 産業スケールのDLRMsに適した、計算効率的かつハードウェアに配慮した圧縮フレームワークを設計する。
• 最適化されたカーネル設計とキャッシュ戦略により、大規模な推薦モデルの実用的で低オーバーヘッドなオンライン学習を可能にする。

提案手法

• 深層学習推薦モデル（DLRMs）内の大きな埋め込みテーブルを、テンソルトレース（TT）分解を用いた低ランク行列積の系列に置き換える。
• TTベースの埋め込み参照を高速化するカスタム最適化TT-EmbeddingBagカーネルを実装し、SOTAのTT実装よりも3倍の高速化を達成する。
• 頻繁にアクセスされる埋め込みベクトルを非圧縮形式でキャッシュするメカニズムを導入し、再計算を減らして精度を回復する。
• トレーニングの安定化とモデル精度の向上を図るため、ガウス分布からのサンプリングを用いた新規な重み初期化戦略を提案する。
• GPUアクセラレータ搭載システムにおける推論および学習パフォーマンスのさらなる最適化のため、バッチ処理された行列乗算を活用する。
• メモリ、精度、速度の複数の設計次元において、TTランクおよび圧縮パラメータを最適化し、パレート最適な設定を同定する。

実験結果

リサーチクエスチョン

• RQ1 テンソルトレース分解は、精度を劣化させることなく、DLRMsにおける顕著なモデル圧縮を達成できるか？
• RQ2 重み初期化の選択が、推薦モデルにおけるTT-Recのパフォーマンスおよび収束に与える影響はいかほどか？
• RQ3 キャッシュは、学習中にTT構造の埋め込みを復元する際のパフォーマンスオーバーヘッドをどの程度軽減できるか？
• RQ4 異なるTTランクおよび埋め込み次元の設定において、メモリ削減、学習時間、モデル精度の間のトレードオフはどのようなものか？
• RQ5 TT-Recは、産業スケールの学習パイプラインおよびハードウェアアクセラレータと互換性を保ちつつ、桁違いの圧縮を達成できるか？

主な発見

• Criteo KaggleおよびTerabyteデータセットにおいて、TT-Recはそれぞれ117倍および112倍のモデルサイズ圧縮を達成し、非圧縮ベースラインと比較して精度に劣化が認められなかった。
• 提案されたTT-EmbeddingBagカーネルは、SOTAのTT実装よりも3倍高速であり、計算効率の向上が顕著に示された。
• 90％のキャッシュヒット率を達成した場合、TT-Recは標準的なEmbeddingBagを上回る学習時間パフォーマンスを示し、キャッシュ戦略の有効性が裏付けられた。
• 極端な圧縮に対しても、ベースラインと同等の精度を維持した。これは、提案された初期化および圧縮手法の強靭性を裏付けている。
• 学習時間のオーバーヘッドは、ベースラインと比較してわずか13.9％に抑えられ、TT-Recがスケールアップしたオンライン学習に実用的であることが示された。
• モデルのメモリ容量、精度、学習速度の間で、調整可能なパラメータを備えたパレート最適なトレードオフを実現しており、さまざまなハードウェアおよび展開制約に適応可能である。