中文

I-BERT:仅整数运算的 BERT 量化方法

计算与语言 2022-05-02 v3

摘要

基于 Transformer 的模型,如 BERT 和 RoBERTa,已在许多自然语言处理任务中取得最先进的结果。然而,它们的内存占用、推理延迟和功耗对于边缘端乃至数据中心的高效推理而言都高得令人望而却步。虽然量化可以作为一个可行的解决方案,但先前关于 Transformer 基模型量化的工作在后向推理期间使用浮点运算,无法有效利用仅整数逻辑单元,例如近期的 Turing Tensor Cores 或传统的仅整数 ARM 处理器。在本工作中,我们提出 I-BERT,一种用于 Transformer 基模型的新型量化方案,其以仅整数运算量化整个推理过程。基于针对非线性运算(如 GELU、Softmax 和层归一化)的轻量级仅整数近似方法,I-BERT 执行端到端的仅整数 BERT 推理,无任何浮点计算。我们使用 RoBERTa-Base/Large 在 GLUE 下游任务上评估我们的方法。我们表明,在这两种情况下,I-BERT 相较于全精度基线取得了相似(且略高)的准确率。此外,我们对 I-BERT 的初步实现显示,在 T4 GPU 系统上,INT8 推理相比 FP32 推理实现了 2.4–4.0 倍的加速。该框架已在 PyTorch 中开发并已开源。

关键词

引用

@article{arxiv.2101.01321,
  title  = {I-BERT: Integer-only BERT Quantization},
  author = {Sehoon Kim and Amir Gholami and Zhewei Yao and Michael W. Mahoney and Kurt Keutzer},
  journal= {arXiv preprint arXiv:2101.01321},
  year   = {2022}
}