[논문 리뷰] Parallel Training of Deep Networks with Local Updates

Michael Laskin, Luke Metz외 6명|arXiv (Cornell University)|2020. 12. 07.

Advanced Neural Network Applications참고 문헌 66인용 수 10

한 줄 요약

이 논문은 전역 역전파를 잘라서 계층 단위로 역전파하는 방법인 로컬 병렬성(local parallelism)을 제안한다. 이는 깊은 신경망의 완전히 비동기적이고 계산 효율적인 훈련을 가능하게 한다. 완전한 순방향 전파가 끝나지 않은 상태에서 각 계층을 독립적으로 업데이트함으로써, 통신 및 메모리 오버헤드를 줄이며, MLP에서는 최대 10배의 속도 향상과 Transformer에서는 2배의 속도 향상을 달성한다. ResNet에서는 전통적인 역전파보다 40% 높은 샘플 처리량을 기록한다.

ABSTRACT

Deep learning models trained on large data sets have been widely successful in both vision and language domains. As state-of-the-art deep learning architectures have continued to grow in parameter count so have the compute budgets and times required to train them, increasing the need for compute-efficient methods that parallelize training. Two common approaches to parallelize the training of deep networks have been data and model parallelism. While useful, data and model parallelism suffer from diminishing returns in terms of compute efficiency for large batch sizes. In this paper, we investigate how to continue scaling compute efficiently beyond the point of diminishing returns for large batches through local parallelism, a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation. Local parallelism enables fully asynchronous layer-wise parallelism with a low memory footprint, and requires little communication overhead compared with model parallelism. We show results in both vision and language domains across a diverse set of architectures, and find that local parallelism is particularly effective in the high-compute regime.

연구 동기 및 목표

• 큰 배치 크기에서 데이터 병렬성과 모델 병렬성의 수익 감소 문제를 해결하기 위해, 병렬화 전략의 대안을 탐색한다.
• 역전파에서 순방향 및 역방향 잠금(locking) 문제를 해결하여 병렬성 제한과 메모리 및 동기화 비용 증가를 완화한다.
• 엄밀히 기울기 기반은 아니지만, 로컬 계층 단위 업데이트가 깊은 신경망의 확장 가능하고 효율적인 훈련을 가능하게 하는지 조사한다.
• 실제 구현을 통해 시각 및 언어 아키텍처에서 로컬 병렬성의 실현 가능성과 이점을 입증한다.
• 로컬 병렬성과 역전파를 통신, 메모리, 훈련 효율성 측면에서 비교하여 고계산 환경에서의 이점 수량화

제안 방법

• 전역 역전파를 잘라서 계층 단위로 역전파하는 것으로 대체하여, 각 계층이 완전한 순방향 전파가 끝나지 않은 상태에서도 독립적으로 업데이트될 수 있도록 한다.
• 청크 기반 로컬 병렬성 구현: 각 계층이 로컬 배치를 처리하고, 하류 계층의 순방향 전파가 완료되기 전에 가중치를 업데이트한다.
• 각 계층에 보조 분류기(auxiliary classifiers)를 도입하여 로컬 감독 신호를 제공하고, 가중치 업데이트를 전역 기울기 계산에서 분리한다.
• 계층 간 비동기 훈련을 적용하여 역방향 잠금을 제거하고, 독립적이고 파ipelinable한 계산을 가능하게 한다.
• 이웃 계층 간에 필요한 활성화 및 기울기 데이터만 전송함으로써 프로세서 간 통신을 최소화하여, 파이프라인 역전파 대비 약 50% 감소된 총 데이터 전송량을 달성한다.
• 정확한 비교를 위해 역전파에서 활성화 재계산(activation recomputation)을 통합하지만, 로컬 병렬성에서는 낮은 메모리 프로파일을 유지한다.

실험 결과

연구 질문

• RQ1 로컬 계층 단위 업데이트가 깊은 신경망에서 기존 역전파 대비 상당한 훈련 속도 향상을 제공할 수 있는가?
• RQ2 로컬 병렬성은 통신 오버헤드와 메모리 소비 측면에서 데이터 병렬성과 모델 병렬성과 비교해 어떻게 성능을 내는가?
• RQ3 로컬 업데이트가 전역 역전파의 유용한 특징 및 기울기 표현을 어느 정도 유지하는가?
• RQ4 로컬 병렬성은 실세계의 시각 및 언어 모델에서 실제로 어떤 성능 향상을 가져오는가, 특히 고계산 환경에서?
• RQ5 로컬 병렬성은 정확도를 유지하면서 하드웨어 활용도와 처리량을 높일 수 있는가?

주요 결과

• 로컬 병렬성은 단순한 MLP에서 순차적 계산 단계를 최대 10배 줄였고, Transformer 아키텍처에서는 2배의 속도 향상을 달성하여 훈련 효율성 향상이 뚜렷하게 입증되었다.
• ImageNet에서 ResNet34를 훈련할 때, 로컬 병렬성은 전통적인 역전파 대비 샘플 처리량(초당 훈련 포인트 수)을 40% 향상시켰다. 이는 더 높은 하드웨어 활용도 덕분이었다.
• 로컬 병렬성에서는 파이프라인 역전파 대비 약 50% 감소된 프로세서 간 통신이 이루어졌으며, 4개의 IPU에서 총 43.2MB의 데이터 전송량을 기록한 반면, 역전파에서는 86.4MB였다.
• 로컬 병렬성에서는 메모리 소비가 항상 낮거나 유사했으며, 특히 고미크로배치 및 다중 프로세서 환경에서 활성화 저장 필요량 감소로 인해 유리했다.
• 기울기 기반은 아니지만, 로컬 업데이트가 역전파와 유사한 특징 및 기울기 표현을 유지하여 최적화에서의 효과성을 뒷받침했다.
• 활성화 재계산는 역전파의 메모리 사용량을 줄였지만, 더 큰 미크로배치와 더 많은 프로세서를 사용할 경우 스케일링 단계에서 로컬 병렬성의 메모리 우수성을 상쇄하지 못했다.