Saved in:
Bibliographic Details
Main Authors: Wang, Ziteng, Chen, Jianfei, Zhu, Jun
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2402.17227
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917599099486208
author Wang, Ziteng
Chen, Jianfei
Zhu, Jun
author_facet Wang, Ziteng
Chen, Jianfei
Zhu, Jun
contents Sampling-based algorithms, which eliminate ''unimportant'' computations during forward and/or back propagation (BP), offer potential solutions to accelerate neural network training. However, since sampling introduces approximations to training, such algorithms may not consistently maintain accuracy across various tasks. In this work, we introduce a variance-controlled adaptive sampling (VCAS) method designed to accelerate BP. VCAS computes an unbiased stochastic gradient with fine-grained layerwise importance sampling in data dimension for activation gradient calculation and leverage score sampling in token dimension for weight gradient calculation. To preserve accuracy, we control the additional variance by learning the sample ratio jointly with model parameters during training. We assessed VCAS on multiple fine-tuning and pre-training tasks in both vision and natural language domains. On all the tasks, VCAS can preserve the original training loss trajectory and validation accuracy with an up to 73.87% FLOPs reduction of BP and 49.58% FLOPs reduction of the whole training process. The implementation is available at https://github.com/thu-ml/VCAS .
format Preprint
id arxiv_https___arxiv_org_abs_2402_17227
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Efficient Backpropagation with Variance-Controlled Adaptive Sampling
Wang, Ziteng
Chen, Jianfei
Zhu, Jun
Machine Learning
Sampling-based algorithms, which eliminate ''unimportant'' computations during forward and/or back propagation (BP), offer potential solutions to accelerate neural network training. However, since sampling introduces approximations to training, such algorithms may not consistently maintain accuracy across various tasks. In this work, we introduce a variance-controlled adaptive sampling (VCAS) method designed to accelerate BP. VCAS computes an unbiased stochastic gradient with fine-grained layerwise importance sampling in data dimension for activation gradient calculation and leverage score sampling in token dimension for weight gradient calculation. To preserve accuracy, we control the additional variance by learning the sample ratio jointly with model parameters during training. We assessed VCAS on multiple fine-tuning and pre-training tasks in both vision and natural language domains. On all the tasks, VCAS can preserve the original training loss trajectory and validation accuracy with an up to 73.87% FLOPs reduction of BP and 49.58% FLOPs reduction of the whole training process. The implementation is available at https://github.com/thu-ml/VCAS .
title Efficient Backpropagation with Variance-Controlled Adaptive Sampling
topic Machine Learning
url https://arxiv.org/abs/2402.17227