arXiv:2502.20853cs.LGcs.AI2025-02ICML被引 22

解决视觉Transformer 4-bit训练精度下降问题,提升模型性能。

Oscillation-Reduced MXFP4 Training for Vision Transformers

  • 提出TetraJet训练方法,通过EMA量化器与自适应调度优化器缓解权重振荡。
  • 相比基线精度损失降低超50%,部分任务接近全精度训练表现。
  • 适合追求高效低精度推理的视觉模型研发人员使用。

在FP4精度下预训练Transformer正成为显著提速的可行方案,但伴随明显精度下降。微缩放(MX)数据格式提供细粒度分组量化能力,支持下一代Blackwell GPU架构。然而,使用MXFP4训练仍导致显著性能退化,且缺乏系统性分析。本文提出TetraJet训练方法,全面评估训练中所有量化器,发现前向传播中的权重振荡是主要成因。为此,引入两种新方法:EMA量化器(Q-EMA)和自适应调度优化器(Q-Ramping),有效抑制振荡。大量实验表明,TetraJet持续优于现有4比特训练方法,且Q-EMA与Q-Ramping可进一步提升性能。相比基线,精度损失减少超过50%,甚至在某些任务上达到与全精度训练相当的表现。代码已开源:https://github.com/thu-ml/TetraJet-MXFP4Training。

原文摘要 · Abstract (English)

Pre-training Transformers in FP4 precision is becoming a promising approach to gain substantial speedup, but it comes with a considerable loss of accuracy. Microscaling (MX) data format provides a fine-grained per-group quantization method to improve the representation ability of the FP4 format and is supported by the next-generation Blackwell GPU architecture. However, training with MXFP4 data format still results in significant degradation and there is a lack of systematic research on the reason. In this work, we propose a novel training method TetraJet for a more accurate FP4 training. We comprehensively evaluate all of the quantizers involved in the training, and identify the weight oscillation problem in the forward pass as the main source of the degradation in MXFP4 training. Therefore, we introduce two novel methods, EMA Quantizer (Q-EMA) and Adaptive Ramping Optimizer (Q-Ramping), to resolve the oscillation problem. Extensive experiments on Vision Transformers demonstrate that TetraJet consistently outperforms the existing 4-bit training methods, and Q-EMA & Q-Ramping can provide additional enhancement by effectively reducing oscillation. We decreased the accuracy degradation by more than $50\%$ compared to the baseline, and can even achieve competitive performance compared to full precision training. The codes are available at https://github.com/thu-ml/TetraJet-MXFP4Training

低精度训练视觉Transformer量化FP4

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。