[논문 리뷰] Revisiting Locally Supervised Learning: an Alternative to End-to-end Training

Yulin Wang, Zanlin Ni외 3명|arXiv (Cornell University)|2021. 01. 26.

Advanced Neural Network Applications참고 문헌 65인용 수 24

한 줄 요약

이 논문은 정보 이론적 원리에 기반한 비그리디스러운 목표를 통해 초기 레이어에서 작업에 관련된 정보 붕괴를 방지하는 새로운 손실 함수인 InfoPro를 제안한다. InfoPro 손실의 계산 가능한 상한을 재구성 손실과 대조 손실의 조합으로 유도함으로써, 전체 네트워크를 통해 역전파가 필요 없는 국소적 목적함수를 도출하여, GPU 메모리 사용량이 40% 미만인 상태에서 엔드 투 엔드 학습과 유사한 성능을 달성하고, 이로 인해 비동기적이고 병렬적인 학습이 가능해진다.

ABSTRACT

Due to the need to store the intermediate activations for back-propagation, end-to-end (E2E) training of deep networks usually suffers from high GPUs memory footprint. This paper aims to address this problem by revisiting the locally supervised learning, where a network is split into gradient-isolated modules and trained with local supervision. We experimentally show that simply training local modules with E2E loss tends to collapse task-relevant information at early layers, and hence hurts the performance of the full model. To avoid this issue, we propose an information propagation (InfoPro) loss, which encourages local modules to preserve as much useful information as possible, while progressively discard task-irrelevant information. As InfoPro loss is difficult to compute in its original form, we derive a feasible upper bound as a surrogate optimization objective, yielding a simple but effective algorithm. In fact, we show that the proposed method boils down to minimizing the combination of a reconstruction loss and a normal cross-entropy/contrastive term. Extensive empirical results on five datasets (i.e., CIFAR, SVHN, STL-10, ImageNet and Cityscapes) validate that InfoPro is capable of achieving competitive performance with less than 40% memory footprint compared to E2E training, while allowing using training data with higher-resolution or larger batch sizes under the same GPU memory constraint. Our method also enables training local modules asynchronously for potential training acceleration. Code is available at: https://github.com/blackfeather-wang/InfoPro-Pytorch.

연구 동기 및 목표

• 중간 활성치를 저장해야 하는 탓에 GPU 메모리 사용량이 높은 엔드 투 엔드(E2E) 학습의 문제를 해결하기 위해 중간 활성치 저장 문제를 해결한다.
• 초기 레이어에서 유용한 입력 정보가 붕괴되는 문제를 야기하는 그리디하고 시야가 짧은 학습 방식으로 인한 성능 저하 문제를 해결한다.
• 국소 모듈 간에 작업에 관련된 정보를 유지하면서 불필요한(해로운) 성분은 제거하는 학습 목표를 개발한다.
• 모듈 간의 기울기 갱신을 분리함으로써 모델 병렬화와 비동기 학습을 가능하게 한다.
• 모델 정확도를 희생시키지 않은 채 실용적이고 메모리 효율적인 E2E 학습의 대안을 제공한다.

제안 방법

• 입력과 중간 특징 간의 상호정보량을 최대화하고, 불필요한 변수에 대한 의존도를 최소화하기 위해 정보 이론 원리에 기반한 정보 전파(InfoPro) 손실을 제안한다.
• InfoPro 손실의 계산 가능한 상한을 도출하여 계산 가능하게 만들며, 재구성 손실과 대조 손실의 조합으로 이루어진 대체 목표함수를 얻는다.
• 대체 InfoPro 손실을 사용하여 국소 모듈을 학습함으로써, 후속 레이어에 유용한 입력 정보를 유지하는 특징 표현을 장려한다.
• 전체 네트워크를 따라 역전파가 필요 없는 국소 목적함수로 최적화를 분해함으로써 기울기 격리 학습을 가능하게 한다.
• 후행 레이어의 오차 신호에 의존하지 않음으로써 국소 모듈의 비동기적 및 병렬 학습을 가능하게 한다.
• PyTorch에 해당 방법을 구현하고, CIFAR, SVHN, STL-10, ImageNet, Cityscapes를 포함한 다양한 벤치마크에서 검증한다.

실험 결과

연구 질문

• RQ1 딥 네트워크에서 그리디한 국소 감독 학습이 엔드 투 엔드 학습에 비해 성능 저하를 야기하는 이유는 무엇인가?
• RQ2 정보 이론적 손실 함수가 국소 감독 학습 중 초기 레이어에서 정보 붕괴를 완화할 수 있는가?
• RQ3 실용적인 딥 러닝 학습을 위해 비그리디스러운 정보 유지 목표를 효과적으로 근사할 수 있는가?
• RQ4 InfoPro 손실을 사용한 국소 감독 학습이 정확도를 유지하면서 GPU 메모리 사용량을 얼마나 줄일 수 있는가?
• RQ5 제안된 방법이 성능 저하 없이 네트워크 모듈의 비동기 또는 병렬 학습을 지원할 수 있는가?

주요 결과

• InfoPro 학습은 표준 국소 감독 학습에서 발생하는 주요 실패 원인인 초기 레이어에서의 작업에 관련된 정보 붕괴를 방지한다.
• CIFAR-10, SVHN, STL-10, ImageNet, Cityscapes에서 E2E 학습과 유사한 성능을 달성하면서 GPU 메모리 사용량이 40% 미만이다.
• 동일한 메모리 제약 조건 하에서 InfoPro는 E2E 학습 대비 50% 더 큰 배치 크기 또는 50% 더 높은 입력 해상도를 지원할 수 있다.
• 재구성 및 대조 항목으로 이루어진 대체 InfoPro 손실은 이론적 정보 이론적 목표를 효과적으로 근사한다.
• 국소 모듈의 비동기 학습을 지원하여 모델 병렬화를 통한 잠재적 성능 향상을 가능하게 한다.
• 분류 및 세그멘테이션 작업을 포함한 다양한 비전 벤치마크에서 제안된 접근법은 강건하고 효과적이다.