[论文解读] Distilled embedding: non-linear embedding factorization using knowledge distillation

Vasileios Lioutas, Ahmad Rashid+3 more|arXiv (Cornell University)|Oct 2, 2019

Topic Modeling参考文献 25被引用 8

一句话总结

本文提出 Distilled Embedding，一种基于知识蒸馏的非线性词嵌入压缩方法，结合低秩矩阵分解。通过先对分解后的嵌入进行预训练以重建完整嵌入，再在下游任务上通过知识蒸馏进行微调，该方法在机器翻译任务中，尤其在高倍率压缩下，性能优于低秩分解和近期的压缩方法。

ABSTRACT

Word-embeddings are a vital component of Natural Language Processing (NLP) systems and have been extensively researched. Better representations of words have come at the cost of huge memory footprints, which has made deploying NLP models on edge-devices challenging due to memory limitations. Compressing embedding matrices without sacrificing model performance is essential for successful commercial edge deployment. In this paper, we propose Distilled Embedding, an (input/output) embedding compression method based on low-rank matrix decomposition with an added non-linearity. First, we initialize the weights of our decomposition by learning to reconstruct the full word-embedding and then fine-tune on the downstream task employing knowledge distillation on the factorized embedding. We conduct extensive experimentation with various compression rates on machine translation, using different data-sets with a shared word-embedding matrix for both embedding and vocabulary projection matrices. We show that the proposed technique outperforms conventional low-rank matrix factorization, and other recently proposed word-embedding matrix compression methods.

研究动机与目标

• 解决词嵌入矩阵高内存占用的问题，该问题阻碍了自然语言处理模型在内存受限的边缘设备上的部署。
• 通过引入非线性和知识蒸馏，改进传统低秩矩阵分解方法，以在压缩过程中保持模型性能。
• 在序列到序列模型中实现共享词嵌入和词表投影矩阵的有效压缩，且不造成性能下降。
• 证明所提出方法在多个数据集和压缩率下均优于现有嵌入压缩技术。

提出的方法

• 对完整词嵌入矩阵应用低秩矩阵分解，将其分解为两个较小矩阵，并引入非线性变换以增强表征能力。
• 通过最小化原始嵌入矩阵与分解后矩阵乘积之间的重建损失，对分解后的嵌入进行预训练。
• 在下游任务上使用知识蒸馏对压缩后的嵌入进行微调，其中学生模型（分解后的嵌入）从教师模型（完整嵌入）的logits中学习。
• 在序列到序列模型中采用共享嵌入和投影矩阵的设置，以在压缩过程中保持一致性并减少冗余。
• 端到端优化整个流程，使用下游任务的交叉熵损失，使压缩后的嵌入与特定任务的表示对齐。
• 采用残差连接或跳跃连接机制，以在分解结构的非线性变换过程中保持信息流动。

实验结果

研究问题

• RQ1 非线性低秩分解结合知识蒸馏是否能在压缩词嵌入方面优于标准低秩分解？
• RQ2 在高倍率压缩下，所提出方法在下游自然语言处理任务中的性能保持程度如何？
• RQ3 预训练与知识蒸馏的结合在机器翻译任务中对压缩嵌入泛化能力的影响如何？
• RQ4 当同一嵌入矩阵在序列到序列模型中同时用于输入和输出投影时，该方法是否仍能保持性能？

主要发现

• 所提出的 Distilled Embedding 方法在所有测试压缩率下均优于传统低秩矩阵分解。
• 在机器翻译任务中，尤其在高倍率压缩下，其性能优于近期提出的词嵌入压缩技术。
• 微调阶段的知识蒸馏显著提升了压缩嵌入在下游任务上的泛化能力。
• 即使在输入和输出投影共享同一嵌入矩阵的常见序列到序列模型设置下，该方法仍能保持优异性能。
• 在完整嵌入重建任务上的预训练阶段，使分解模型能够更准确地学习原始嵌入空间的近似表示。
• 分解中的非线性组件相比仅使用线性分解，能带来更优的表征学习效果。