[论文解读] 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），MLET利用了学习表征中的更低有效秩，使得在不损失准确率的前提下，嵌入表大小减少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开源的DLRM框架中进行MLET的原型设计与真实性能评估。

提出的方法

• MLET用两层线性变换替代单个嵌入表 W ∈ ℝⁿˣᵈ：qW₁W₂，其中 W₁ ∈ ℝⁿˣᵏ 且 W₂ ∈ ℝᵏˣᵈ。
• 模型通过反向传播进行端到端训练，推理阶段仅使用 W₂ 以保持输出为 d 维。
• 该方法利用了线性网络动态的理论洞见：更深的线性网络通常具有更低的有效秩，从而实现更好的泛化能力。
• 该方法在Facebook基于PyTorch的DLRM框架中实现并评估，数据集为Criteo-Kaggle和Avazu。
• 训练内存相比单层训练增加 k/d 倍，但推理内存仍与 d 成正比。
• 推理阶段将最终模型转换为单层形式，以保持与现有系统的兼容性。

实验结果

研究问题

• RQ1 多层线性架构能否在准确率和模型尺寸方面优于单层嵌入？
• RQ2 为何多层线性网络能生成优于单层网络的嵌入？其背后的动态机制是什么？
• RQ3 在保持或提升模型准确率的前提下，MLET最多能将嵌入表大小减少多少？
• RQ4 是否存在MLET在不增加额外训练内存成本的情况下同时提升准确率和减少模型尺寸的配置？

主要发现

• 在Criteo-Kaggle和Avazu数据集上，MLET在保持相同模型准确率水平的前提下，将嵌入表大小减少了4–8倍。
• MLET在Criteo-Kaggle上的最大LogLoss改进为0.0025，在Avazu上为0.006，且当k超过某一临界值后性能趋于饱和。
• 在某些k/d比例下，MLET在不增加额外训练内存成本（即k ≤ d）的情况下，同时实现了更高的准确率和更小的模型尺寸，该现象在Avazu中观察到，但在Criteo-Kaggle中未观察到。
• MLET训练的运行时间平均比单层训练高出25%，但推理时间与标准DLRM完全相同。
• 性能提升主要由k/d比例决定，k/d越高，准确率单调提升。
• 即使在d值较低时，该方法依然有效，此时模型尺寸的相对改进最为显著，原因在于可利用k ≤ d实现零成本收益。