具有方差控制的自适应采样高效反向传播
机器学习
2024-02-28 v1
摘要
基于采样的算法通过在正向和/或反向传播 (BP) 过程中消除“不重要”的计算，为加速神经网络训练提供了潜在的解决方案。然而，由于采样引入了对训练的近似，此类算法可能无法在各种任务中始终保持一致的准确性。在这项工作中，我们引入了一种方差控制的自适应采样 (VCAS) 方法，旨在加速 BP。VCAS 在数据维度上计算激活梯度的细粒度逐层重要性采样，并在 token 维度上计算权重梯度的杠杆得分采样，从而计算出无偏的随机梯度。为了保持准确性，我们在训练过程中联合学习采样率与模型参数，以控制额外的方差。我们在视觉和自然语言领域的多个微调和预训练任务上评估了 VCAS。在所有任务中，VCAS 都能保持原始的训练损失轨迹和验证准确率，同时 BP 的 FLOPs 最多减少 73.87%，整个训练过程的 FLOPs 减少 49.58%。实现代码可在 https://github.com/thu-ml/VCAS 获取。
引用
@article{arxiv.2402.17227,
title = {Efficient Backpropagation with Variance-Controlled Adaptive Sampling},
author = {Ziteng Wang and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2402.17227},
year = {2024}
}
备注
ICLR 2024