[论文解读] MPCFormer: fast, performant and private Transformer inference with MPC

Dacheng Li, Rulin Shao+4 more|arXiv (Cornell University)|Nov 2, 2022

Cryptography and Data Security被引用 15

一句话总结

MPCFormer 通过结合 MPC 友好的函数近似与知识蒸馏，利用安全多方计算（MPC）实现了快速、私密且高性能的 Transformer 推理。其在 GLUE 上保持了 BERT_BASE 97% 的性能，同时相比 BERT_BASE 和 BERT_LARGE 实现了最高达 5.9× 的加速，且在多种模型与数据集上均保持了极低的精度损失。

ABSTRACT

Enabling private inference is crucial for many cloud inference services that are based on Transformer models. However, existing private inference solutions can increase the inference latency by more than 60x or significantly compromise the inference quality. In this paper, we design the framework MPCFORMER as a practical solution, using Secure Multi-Party Computation (MPC) and Knowledge Distillation (KD). Through extensive evaluations, we show that MPCFORMER significantly speeds up Transformer inference in MPC settings while achieving similar ML performance to the input model. On the IMDb dataset, it achieves similar performance to BERTBASE, while being 5.3x faster. On the GLUE benchmark, it achieves 97% performance of BERTBASE with a 2.2x speedup. MPCFORMER remains effective with different trained Transformer weights such as ROBERTABASE and larger models including BERTLarge. Code is available at https://github.com/MccRee177/MPCFormer.

研究动机与目标

• 解决使用 MPC 进行私有 Transformer 推理时高延迟的关键挑战，当前该方法使推理速度降低超过 60 倍。
• 克服 MPC 友好近似导致的性能下降问题，此类近似会干扰训练过程，并需要大量下游数据集。
• 实现在不牺牲模型准确率的前提下，将私有推理服务部署于大规模 Transformer 模型的实用化。
• 设计一种与多种预训练模型（如 BERT_BASE、RoBERTa_BASE、BERT_LARGE）兼容的框架，同时支持 MPC 友好的近似方法。
• 通过高效的知识蒸馏，在 MPC 环境中同时实现低推理延迟与高机器学习性能。

提出的方法

• 将预训练 Transformer 中的瓶颈函数替换为 MPC 友好的近似（例如，Quad+2ReLU，一种新型快速 Softmax 近似），以加速 MPC 推理。
• 应用知识蒸馏（KD）技术，利用中间隐藏表示将知识从原始教师模型迁移至近似后的学生模型。
• 在学生模型训练过程中，使用教师模型的 logits 和中间特征作为监督信号，即使在小规模下游数据集上也能实现有效蒸馏。
• 通过联合使用蒸馏损失与标准交叉熵损失训练学生模型，并对嵌入层与注意力层的超参数进行微调。
• 将蒸馏后的学生模型部署于 MPC 引擎（如 CrypTen）上，实现端到端的私有推理。
• 通过数据高效的训练策略优化蒸馏过程，降低对大规模下游数据集的依赖。

实验结果

研究问题

• RQ1 MPC 友好的近似是否能在不引起私有推理性能显著下降的前提下，有效应用于 Transformer 模型？
• RQ2 在 MPC 约束下训练近似后的 Transformer 模型时，知识蒸馏在多大程度上能够缓解性能损失？
• RQ3 MPCFormer 在不同模型规模（包括 BERT_LARGE）上的可扩展性如何？与无蒸馏基线相比表现如何？
• RQ4 在 MPC 环境中，实现高性能蒸馏所需的最小下游数据量是多少？
• RQ5 使用教师模型的权重初始化学生模型是否能提升训练稳定性与最终性能，尤其是在小数据集上？

主要发现

• 在 IMDb 数据集上，MPCFormer 使用 BERT_BASE 实现了 95.0% 的准确率，并实现了 5.3× 的加速，性能与原始模型相当。
• 在 IMDb 上使用 BERT_LARGE 时，MPCFormer 实现了 94.5% 的准确率，加速比达 5.9×，显著优于无蒸馏的基线模型。
• 在 GLUE 基准测试中，MPCFormer 实现了 BERT_BASE 平均得分的 97%（84.6 vs. 85.3），并实现了 2.2× 的加速，展现出强大的泛化能力。
• MPCFormer 在多种预训练模型（包括 RoBERTa_BASE）上均保持了高性能，表明其与多样化架构具有良好的兼容性。
• 知识蒸馏使小样本数据集上的有效训练成为可能：仅需中等规模数据集（如 QNLI）的 2% 和小规模数据集（如 MRPC）的 5% 即可实现高性能。
• 使用教师模型权重初始化学生模型，相比随机初始化能获得更优性能，尤其在小规模下游数据集上优势更明显。