[논문 리뷰] Top-KAST: Top-K Always Sparse Training

Siddhant M. Jayakumar, Razvan Pascanu외 3명|arXiv (Cornell University)|2021. 06. 07.

Advanced Neural Network Applications인용 수 21

한 줄 요약

Top-KAST는 활성화에 대해 상위-K 크기의 가중치를 선택하고 기울기 갱신을 위해 더 큰 집합을 사용함으로써 전방 및 역방향 전파에서 일정한 희소성을 유지하는 완전 희소 학습 방법이다. 하위최적 마스크로의 수렴을 방지하기 위해 탐색 손실을 도입한다. ImageNet과 enwik8에서 최신 기술 성능을 달성하며, 밀도 모델과 80% 희소성까지 동등한 성능를 보이며, 밀도 매개변수나 기울기 계산 없이 효율적이고 확장 가능한 학습을 가능하게 한다.

ABSTRACT

Sparse neural networks are becoming increasingly important as the field seeks to improve the performance of existing models by scaling them up, while simultaneously trying to reduce power consumption and computational footprint. Unfortunately, most existing methods for inducing performant sparse models still entail the instantiation of dense parameters, or dense gradients in the backward-pass, during training. For very large models this requirement can be prohibitive. In this work we propose Top-KAST, a method that preserves constant sparsity throughout training (in both the forward and backward-passes). We demonstrate the efficacy of our approach by showing that it performs comparably to or better than previous works when training models on the established ImageNet benchmark, whilst fully maintaining sparsity. In addition to our ImageNet results, we also demonstrate our approach in the domain of language modeling where the current best performing architectures tend to have tens of billions of parameters and scaling up does not yet seem to have saturated performance. Sparse versions of these architectures can be run with significantly fewer resources, making them more widely accessible and applicable. Furthermore, in addition to being effective, our approach is straightforward and can easily be implemented in a wide range of existing machine learning frameworks with only a few additional lines of code. We therefore hope that our contribution will help enable the broader community to explore the potential held by massive models, without incurring massive computational cost.

연구 동기 및 목표

• 학습 중에 밀도 매개변수나 기울기 계산을 피하는 완전 희소 학습 방법을 개발함으로써 확장 가능하고 효율적인 모델 학습을 가능하게 하는 것.
• 큰 밀도 모델의 학습에 따른 계산 및 환경 비용을 줄이기 위해 학습 전반에 걸쳐 희소성을 유지하는 것.
• 일정한 희소성을 통해 하드웨어 및 에너지 요구량을 감소시킴으로써 대규모 모델에 대한 접근성을 민주화하는 것.
• 특히 자연어 처리 분야에서 자원이 제한된 환경에서 트랜스포머와 같은 대규모 아키텍처의 효과적인 희소화를 가능하게 하는 것.
• 기존 딥 러닝 라이브러리에 쉽게 통합될 수 있도록 단순하고 프레임워크 호환성 있는 방법을 제공하는 것.

제안 방법

• 전방 전파에서 상위-K 크기의 매개변수를 크기 기반으로 선택하여 추론 중에도 희소성을 유지한다.
• 역방향 전파에서 기울기는 더 큰 매개변수 집합 B에 적용되며, 이 집합 B는 상위-K 집합 A를 포함한다. 이를 통해 더 효과적인 가중치 갱신이 가능하다.
• 희소성 마스크의 동적 적응을 장려하고 하위최적 가중치 하위집합에 고착되는 것을 방지하기 위해 보조 탐색 손실을 도입한다.
• 이 방법은 밀도 매개변수나 기울기 계산을 완전히 회피하여 전방 및 역방향 전파에서 일정한 희소성을 보장한다.
• 사용자 정의 하이퍼파라미터를 통해 상위-K 및 탐색 집합 크기를 제어함으로써 자원 제약에 따라 민감하게 조정 가능한 유연성을 제공한다.
• 기존 딥 러닝 프레임워크에 최소한의 코드 변경으로 쉽게 통합될 수 있도록 설계되어 있다.

실험 결과

연구 질문

• RQ1 밀도 매개변수나 기울기를 전혀 생성하지 않으면서도 밀도 모델과 비교할 수 있는 성능을 유지할 수 있는 완전 희소 학습 방법이 가능한가?
• RQ2 ImageNet과 enwik8와 같은 표준 벤치마크에서 기존의 희소-희소 및 밀도-희소 학습 방법과 비교해 Top-KAST는 어떤 성능을 보이는가?
• RQ3 Top-KAST는 언어 모델링에서 대규모 트랜스포머 아키텍처를 효과적으로 희소화할 수 있으며 성능 저하를 최소화할 수 있는가?
• RQ4 탐색 손실의 포함이 일반화 성능 향상과 나쁜 희소 하위네트워크로의 수렴 방지를 위해 유의미한가?
• RQ5 실제로 Top-KAST는 모델 크기를 얼마나 확장할 수 있으며, 계산 및 메모리 오버헤드를 얼마나 줄일 수 있는가?

주요 결과

• Top-KAST는 다양한 FLOP 예산에서 ImageNet 벤치마크에서 최신 기술 성능을 달성하며, 기존의 희소-희소 방법을 능가하고 RigL과 동등한 성능을 기록한다.
• enwik8 언어 모델링 작업에서 Top-KAST로 학습된 Transformer-XL은 5500만 파라미터에서 1.00 비트/문자(bpc) 성능을 달성했으며, 이는 2억 7700만 파라미터의 더 작은 밀도 모델(0.99 bpc)을 능가하는 성능이다.
• Top-KAST는 ImageNet에서 80% 희소성까지 높은 성능을 유지하며 성능 저하가 최소한이었으며, 90% 희소성 이상에서는 성능이 점차 떨어지기 시작한다.
• Top-KAST는 80% 희소성까지 밀도 모델 성능을 따라잡는 희소 모델 학습을 가능하게 하여 뛰어난 효율성과 확장성을 입증한다.
• 이 방법은 비전 및 언어 모델링 모두에 효과적이며, 다양한 모델 유형과 도메인에 걸쳐 광범위한 적용 가능성을 보여준다.
• 기존 프레임워크에 몇 줄의 추가 코드만으로 쉽게 구현 가능하여 널리 보급될 수 있다.