[论文解读] NxMTransformer: Semi-Structured Sparsification for Natural Language Understanding via ADMM

Connor Holmes, Minjia Zhang+2 more|arXiv (Cornell University)|Oct 28, 2021

Non-Destructive Testing Techniques参考文献 29被引用 6

一句话总结

NxMTransformer 提出了一种基于 ADMM 的新框架，用于在大规模预训练 Transformer 模型中引入 NxM 半结构化稀疏性，以实现自然语言理解任务，其 GLUE 得分比当前最佳实践高出 1.7 分，同时保持硬件效率。该方法通过约束优化联合优化模型微调与稀疏性约束，实现高精度、可部署的模型，适用于现代稀疏推理硬件。

ABSTRACT

Natural Language Processing (NLP) has recently achieved success by using huge pre-trained Transformer networks. However, these models often contain hundreds of millions or even billions of parameters, bringing challenges to online deployment due to latency constraints. Recently, hardware manufacturers have introduced dedicated hardware for NxM sparsity to provide the flexibility of unstructured pruning with the runtime efficiency of structured approaches. NxM sparsity permits arbitrarily selecting M parameters to retain from a contiguous group of N in the dense representation. However, due to the extremely high complexity of pre-trained models, the standard sparse fine-tuning techniques often fail to generalize well on downstream tasks, which have limited data resources. To address such an issue in a principled manner, we introduce a new learning framework, called NxMTransformer, to induce NxM semi-structured sparsity on pretrained language models for natural language understanding to obtain better performance. In particular, we propose to formulate the NxM sparsity as a constrained optimization problem and use Alternating Direction Method of Multipliers (ADMM) to optimize the downstream tasks while taking the underlying hardware constraints into consideration. ADMM decomposes the NxM sparsification problem into two sub-problems that can be solved sequentially, generating sparsified Transformer networks that achieve high accuracy while being able to effectively execute on newly released hardware. We apply our approach to a wide range of NLP tasks, and our proposed method is able to achieve 1.7 points higher accuracy in GLUE score than current practices. Moreover, we perform detailed analysis on our approach and shed light on how ADMM affects fine-tuning accuracy for downstream tasks. Finally, we illustrate how NxMTransformer achieves performance improvement with knowledge distillation.

研究动机与目标

• 为解决因模型规模过大导致的高延迟问题，以部署大规模预训练 NLP 模型。
• 在尊重硬件约束的同时，实现在 Transformer 中有效且半结构化的 NxM 稀疏性，同时保持模型精度。
• 开发一种系统化、硬件感知的稀疏化方法，避免昂贵的微调重训练，并保持在下游任务中的迁移能力。
• 研究将 NxM 稀疏性与知识蒸馏结合的兼容性及性能增益，以进一步压缩模型。
• 提供一种框架，在保持密集预训练表征的同时，实现在新兴 NxM 稀疏硬件上的高效推理。

提出的方法

• 将 NxM 稀疏性诱导建模为约束优化问题，将硬件感知的稀疏模式整合进微调目标。
• 采用交替方向乘子法（ADMM）将问题分解为两个顺序子问题：权重更新与稀疏性约束执行。
• 通过对偶上升法与增广拉格朗日函数，在微调下游 NLP 任务的同时迭代执行 NxM 约束。
• 将该方法应用于 BERT 和 DistilBERT 模型，实现稀疏化而无需重新训练原始预训练权重。
• 通过在蒸馏后的学生模型中应用 NxM 稀疏性，与知识蒸馏集成，实现参数减少的同时保持高精度。
• 通过保留原始预训练权重，仅在微调阶段进行稀疏化，保持与现有密集模型的兼容性。

实验结果

研究问题

• RQ1 基于 ADMM 的优化能否在不牺牲下游性能的前提下，有效诱导大规模预训练 Transformer 中的 NxM 半结构化稀疏性？
• RQ2 与标准微调及现有稀疏化方法相比，ADMM 实现的 NxM 稀疏性在精度与硬件效率方面表现如何？
• RQ3 NxM 稀疏性与知识蒸馏结合的程度如何，能否在保持高精度的同时进一步压缩模型？
• RQ4 基于 ADMM 的稀疏化方法是否能保留模型在低资源下游 NLP 任务中的泛化能力？
• RQ5 稀疏性对蒸馏模型有何影响？该方法是否揭示了即使在已压缩的模型（如 DistilBERT）中也存在结构冗余？

主要发现

• NxMTransformer 在标准 NLP 基准测试中，平均 GLUE 得分比当前最佳实践高出 1.7 分。
• 该方法在应用 NxM 稀疏性的同时，保持了原始 BERT base 模型 98.4% 的精度，展现出优异的性能保持能力。
• 应用于 DistilBERT 时，NxMTransformer 在参数减少 32%（从 66.6M 减少到 45.3M）的同时，保持了 97.6% 的蒸馏模型精度。
• 与标准 DistilBERT 相比，该方法在平均 GLUE 上高出 2 分，表明通过层内稀疏性可更有效地捕捉冗余信息。
• 该方法无需重新训练预训练权重，即可实现高精度、硬件优化的模型，同时保持迁移能力。
• 结果表明，即使在知识蒸馏后高度压缩的模型中，仍可能存在‘彩票’结构。