[论文解读] Differentiable Product Quantization for End-to-End Embedding Compression

Ting Chen, Lala Li+1 more|arXiv (Cornell University)|Aug 26, 2019

Advanced Data Compression Techniques参考文献 24被引用 14

一句话总结

本文提出可微分乘积量化（DPQ），一种新颖的端到端可训练框架，通过学习离散乘积量化码来压缩神经网络嵌入层。通过使乘积量化可微，DPQ 实现了在 10 个语言数据集上 14–238× 压缩比下近乎零性能损失的单阶段训练，相较于先前方法在效率和压缩性方面表现更优，且无需蒸馏。

ABSTRACT

Embedding layers are commonly used to map discrete symbols into continuous embedding vectors that reflect their semantic meanings. Despite their effectiveness, the number of parameters in an embedding layer increases linearly with the number of symbols and poses a critical challenge on memory and storage constraints. In this work, we propose a generic and end-to-end learnable compression framework termed differentiable product quantization (DPQ). We present two instantiations of DPQ that leverage different approximation techniques to enable differentiability in end-to-end learning. Our method can readily serve as a drop-in alternative for any existing embedding layer. Empirically, DPQ offers significant compression ratios (14-238$ imes$) at negligible or no performance cost on 10 datasets across three different language tasks.

研究动机与目标

• 解决神经网络中大型嵌入表带来的高内存和存储开销，尤其是在自然语言处理和推荐系统中的问题。
• 克服先前方法中使用固定或不可微离散码的局限性，这些方法需要两阶段训练与蒸馏。
• 开发一种通用的、端到端可微的框架，学习紧凑的离散码以表示嵌入向量，且不造成性能下降。
• 通过用量化码替换完整嵌入表，实现在保持模型精度的同时实现高效推理与训练。

提出的方法

• 提出一种可微分乘积量化（DPQ）框架，通过可微离散化函数 φ(·) 和反向离散化函数 ρ(·)，将连续嵌入向量映射为离散码。
• 引入两种近似技术——DPQ-SX 和 DPQ-VQ——分别采用直通梯度估计和带软分配的向量量化，以实现通过离散码的反向传播。
• 使用乘积量化（PQ）将嵌入空间划分为 D 个正交子空间，每个子空间包含 K 个中心点，从而以 D 维索引向量形式实现紧凑的码表示。
• 通过反向传播穿过量化过程，实现整个模型的端到端训练，使中心点与码能够联合优化以提升下游任务性能。
• 将标准嵌入层替换为 DPQ 模块作为即插即用的替代方案，在保持相同输入输出行为的同时显著减少参数量。
• 通过近似硬离散化过程为连续松弛形式，实现可微性，从而在反向传播中实现梯度流动，且无需随机采样。

实验结果

研究问题

• RQ1 能否使乘积量化可微，以实现离散嵌入码的端到端训练？
• RQ2 码大小 K 和码长度 D 的选择如何影响压缩比与模型性能？
• RQ3 DPQ 是否能在不造成性能下降的情况下实现高倍压缩比，尤其相较于不可微或两阶段方法？
• RQ4 与标准嵌入层相比，DPQ 在训练和推理过程中的计算开销如何？
• RQ5 DPQ 是否能消除先前工作中所需的蒸馏或预训练阶段？

主要发现

• DPQ 在 10 个多样化的 NLP 数据集上实现了 14–238× 的压缩比，与全精度嵌入层相比性能损失可忽略或无损失。
• 采用小 K（如 K=2）与大 D（如 D=128）的组合可获得优于大 K 与小 D 的性能与压缩效果，原因在于量化精度的提升。
• 使用 DPQ 训练平均仅增加 10% 的训练时间，且无额外内存占用，推理阶段无性能开销。
• 当 K 和 D 较大时，DPQ-VQ 相较于 DPQ-SX 展现出更高的计算效率，得益于优化的向量量化操作。
• DPQ 消除了对蒸馏或两阶段训练的需求，支持单阶段端到端学习，且收敛性与稳定性更优。
• 在 BERT 预训练的实证结果表明，DPQ 在显著压缩嵌入层的同时，仍能保持具有竞争力的性能。