用零阶信息改进直通估计器，提升量化模型训练效率与精度
Improving the Straight-Through Estimator with Zeroth-Order Information
- 结合一阶引导的零阶梯度下降，降低直通估计偏差
- 在相同迭代次数下，模型准确率提升1-22个百分点
- 大幅减少计算量，比传统零阶方法快796倍
我们研究量化参数神经网络的训练问题。通过直通估计器（STE）实现梯度计算虽能支持反向传播（一阶方法），但存在偏差；而零阶（ZO）梯度下降虽无偏但计算成本高。为此提出一阶引导的零阶梯度下降（FOGZO），在降低STE偏差的同时减少计算开销。实验表明，相比标准STE，在相同迭代次数下，DeiT Tiny/Small模型准确率提升1-8%，ResNet 18/50提升1-2%，LLaMA模型（最大0.3亿参数）困惑度降低1-22点。对于两层MLP在MNIST上的任务，相同损失下，FOGZO相较n-SPSA减少796倍计算量。代码已开源：https://github.com/1733116199/fogzo。
原文摘要 · Abstract (English)
We study the problem of training neural networks with quantized parameters. Learning low-precision quantized parameters by enabling computation of gradients via the Straight-Through Estimator (STE) can be challenging. While the STE enables back-propagation, which is a first-order method, recent works have explored the use of zeroth-order (ZO) gradient descent for fine-tuning. We note that the STE provides high-quality biased gradients, and ZO gradients are unbiased but can be expensive. We thus propose First-Order-Guided Zeroth-Order Gradient Descent (FOGZO) that reduces STE bias while reducing computations relative to ZO methods. Empirically, we show FOGZO improves the tradeoff between quality and training time in Quantization-Aware Pre-Training. Specifically, versus STE at the same number of iterations, we show a 1-8\% accuracy improvement for DeiT Tiny/Small, 1-2\% accuracy improvement on ResNet 18/50, and 1-22 perplexity point improvement for LLaMA models with up to 0.3 billion parameters. For the same loss, FOGZO yields a 796$\times$ reduction in computation versus n-SPSA for a 2-layer MLP on MNIST. Code is available at https://github.com/1733116199/fogzo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。