[논문 리뷰] Efficient Neural Network Training via Forward and Backward Propagation Sparsification

Xiao Zhou, Weizhong Zhang외 3명|arXiv (Cornell University)|2021. 11. 10.

Domain Adaptation and Few-Shot Learning참고 문헌 51인용 수 15

한 줄 요약

이 논문은 깊이 신경망을 위한 완전히 희소화된 훈련 방법을 제안하며, 가중치와 아키텍처 파라미터 업데이트를 분리함으로써 효율적인 순방향 및 역방향 전파를 달성한다. 두 번의 순방향 전파만을 요구하는 분산 감소 정책 기반 기울기 추정기(Variancе-reduced Policy Gradient Estimator, VR-PGE)를 도입하여 완전히 희소화된 역방향 전파를 가능하게 하고, 이로 인해 이전 방법 대비 최대 한 단계 빠른 훈련을 실현한다.

ABSTRACT

Sparse training is a natural idea to accelerate the training speed of deep neural networks and save the memory usage, especially since large modern neural networks are significantly over-parameterized. However, most of the existing methods cannot achieve this goal in practice because the chain rule based gradient (w.r.t. structure parameters) estimators adopted by previous methods require dense computation at least in the backward propagation step. This paper solves this problem by proposing an efficient sparse training method with completely sparse forward and backward passes. We first formulate the training process as a continuous minimization problem under global sparsity constraint. We then separate the optimization process into two steps, corresponding to weight update and structure parameter update. For the former step, we use the conventional chain rule, which can be sparse via exploiting the sparse structure. For the latter step, instead of using the chain rule based gradient estimators as in existing methods, we propose a variance reduced policy gradient estimator, which only requires two forward passes without backward propagation, thus achieving completely sparse training. We prove that the variance of our gradient estimator is bounded. Extensive experimental results on real-world datasets demonstrate that compared to previous methods, our algorithm is much more effective in accelerating the training process, up to an order of magnitude faster.

연구 동기 및 목표

• 희소화된 순방향 전파는 이루어지지만 여전히 밀집된 역방향 전파가 필요한 기존의 희소 훈련 방법의 비효율성 문제를 해결하기 위해.
• 아키텍처 파라미터에 대한 기울기 추정 방식을 재고함으로써 완전히 희소화된 순방향 및 역방향 전파를 가능하게 하기 위해.
• 메모리 오버헤드를 증가시키지 않으면서 더 큰, 과도하게 파rameter화된 네트워크의 훈련을 지원할 수 있는 확장 가능하고 메모리 효율적인 방법을 개발하기 위해.
• 기본적인 딥러닝 프레임워크인 PyTorch 및 TensorFlow에서 실질적인 계산 속도 향상을 달성하기 위해.

제안 방법

• 전체 희소성 조건 하에 희소 신경망 훈련을 제약 최적화 문제로 재정의하며, 가중치와 아키텍처 파라미터(베르누이 마스크 확률)를 동시에 최적화한다.
• 훈련을 두 단계로 분리한다: 표준 체인 룰 기반 역방향 전파를 통한 가중치 업데이트(마스크 구조에 의해 희소화됨), 그리고 새로운 분산 감소 정책 기반 기울기 추정기(VR-PGE)를 통한 아키텍처 파라미터 업데이트.
• VR-PGE를 사용해 두 번의 순방향 전파만으로 아키텍처 파라미터의 기울기를 추정함으로써, 역방향 전파가 필요 없게 하고 완전히 희소화된 계산을 가능하게 한다.
• 베르누이 파라미터의 합을 제약함으로써 전체 희소성 수준을 제어하여, 훈련 전반에 걸쳐 일관된 네트워크 크기를 유지한다.
• 순방향 및 역방향 전파 모두에서 희소 서브넷을 사용하여, FLOPs를 밀집 기준 대비 약 ρ²로 감소시킨다. 여기서 ρ는 채널 유지 비율이다.
• 새로운 서브넷 샘플링 빈도를 낮추는 전략(예: 50회 반복마다 한 번)을 도입하여, 전체 모델을 CPU에 저장하고 서브넷을 GPU에 로드함으로써 메모리 부담을 줄인다.

실험 결과

연구 질문

• RQ1 밀집된 역방향 전파에 의존하지 않고도 딥러닝 네트워크에서 순방향 및 역방향 전부를 완전히 희소화한 훈련을 달성할 수 있는가?
• RQ2 분산 감소 정책 기반 기울기 추정기가 기울기 체인 룰 기반 추정 방식을 대체하면서도 훈련의 안정성과 수렴성을 유지할 수 있는가?
• RQ3 완전히 희소화된 훈련이 실질적으로 표준 딥러닝 프레임워크에서 높은 계산 속도 향상을 이끌 수 있는가?
• RQ4 제안된 방법이 과도하게 큰 네트워크로 확장될 수 있으며, 메모리나 계산 비용이 기하급수적으로 증가하지는 않는가?

주요 결과

• VGG-19에서 0.8%의 가중치만 유지되는 조건(ρ = 0.8)에서 훈련 계산 시간에 대해 최대 7.41배의 속도 향상을 달성한다.
• ResNet32를 사용한 CIFAR-10에서 3.06%의 가중치만 유지되며 4.50배의 속도 향상을 기록했고, 검증 정확도는 90.82%를 유지한다.
• 제안된 VR-PGE 추정기의 분산은 증명 가능하게 유계이므로 안정적이고 신뢰할 수 있는 기울기 업데이트를 보장한다.
• 모든 희소성 수준에서 높은 정확도를 유지하며, 매우 낮은 채널 유지 비율에서도 미미한 정확도 저하(예: VGG-19에서 93.84%에서 90.82%로)를 보인다.
• 실증 결과에 따르면, 이전의 희소 훈련 기반 방법보다 현저히 빠르며 최대 한 단계 빠른 속도 향상을 기록한다.
• 서브넷 샘플링 주기를 50회 반복마다 한 번으로 줄여도 최종 정확도가 유지되며, 이는 효율적인 메모리 관리와 더 큰 모델로의 확장 가능성을 보장한다.