[논문 리뷰] Online normalizer calculation for softmax

Maxim Milakov, Natalia Gimelshein|arXiv (Cornell University)|2018. 05. 08.

Parallel Computing and Optimization Techniques참고 문헌 8인용 수 20

한 줄 요약

이 논문은 소프트맥스 함수에 대한 온라인 정규화 계산을 제안하여 요소당 메모리 접근 횟수를 네 번에서 세 번으로 줄여 더 빠른 추론을 가능하게 한다. 테슬라 V100에서의 벤치마크 결과, 소프트맥스만 적용할 경우 최대 1.3배의 성능 향상과 소프트맥스+TopK 병합 시 최대 5배의 성능 향상을 기록했으며, 이는 메모리 대역폭 부담 감소로 인한 성능 향상이다.

ABSTRACT

The Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothesize that this reduction in memory accesses should improve Softmax performance on actual hardware. The benchmarks confirm this hypothesis: Softmax accelerates by up to 1.3x and Softmax+TopK combined and fused by up to 5x.

연구 동기 및 목표

• 딥 러닝 추론에서 성능 저하 요소인 소프트맥스 계산의 메모리 접근 오버헤드를 줄이기 위해.
• 다양한 대안들이 존재하지만도 불구하고 고전적인 소프트맥스 함수에 대한 특화 최적화가 부족한 문제를 해결하기 위해.
• 정규화 및 선택 로직을 동시에 처리함으로써 소프트맥스와 TopK 연산의 효율적 융합을 가능하게 하기 위해.
• 단일 패assing 알고리즘을 통해 수치적 안정성을 유지하면서도 메모리 대역폭 사용을 줄이기 위해.
• 특히 GPU에서 현대 하드웨어에서 측정 가능한 성능 향상을 입증하기 위해.

제안 방법

• 요소당 네 번에서 세 번으로 메모리 접근 횟수를 줄이기 위해, 최댓값과 소프트맥스 정규화 항을 한 번의 반복 내에서 동시에 계산하는 단일 패assing 알고리즘을 도입한다.
• 오버플로우/언더플로우를 방지하기 위해 지수 연산 이전에 모든 로짓에서 최댓값을 빼는 수치적으로 안정적인 공식을 사용한다.
• 정규화 항에 대한 점진적 업데이트 규칙을 적용한다: $ d_j = d_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j} $, 여기서 $ m_j $ 는 누적 최댓값이다.
• 각 요소를 처리할 때마다 누적 최댓값 $ m_j $ 와 정규화된 합 $ d_j $ 를 점진적으로 갱신한다.
• 동일한 패assing 동안 상위-k 값도 추적함으로써 소프트맥스와 TopK 연산을 융합하여 중복된 메모리 접근을 제거한다.
• 메모리 대역폭 사용을 최소화하고 커널 융합을 가능하게 하여 GPU 성능 최적화를 수행한다.

실험 결과

연구 질문

• RQ1 소프트맥스 계산에서 메모리 접근 횟수를 줄이면 현대 하드웨어에서 측정 가능한 성능 향상이 이루어지는가?
• RQ2 수치적 안정성을 희생시키지 않고도 최댓값과 소프트맥스 정규화 항을 한 번의 패assing 내에서 계산할 수 있는가?
• RQ3 제안된 온라인 정규화 방법은 기존의 이중 또는 삼중 패assing 구현과 비교해 성능 및 정확도 측면에서 어떻게 다른가?
• RQ4 소프트맥스와 TopK 연산의 융합이 엔드 투 엔드 추론 속도를 얼마나 향상시킬 수 있는가?
• RQ5 GPU 아키텍처에서 벡터 크기와 배치 크기에 따라 성능 향상이 비례하는가?

주요 결과

• 제안된 온라인 정규화 방법은 요소당 메모리 접근 횟수를 네 번에서 세 번으로 줄여, 대규모 벡터 크기에서 테슬라 V100에서 소프트맥스만 적용할 경우 최대 1.3배의 성능 향상을 달성한다.
• TopK와 병합될 경우, 융합으로 인한 2.5배와 온라인 정규화로 인한 2배의 성능 향상으로 인해 소프트맥스+TopK 조합에서 최대 5배의 성능 향상을 기록한다.
• 성능 향상은 특히 메모리 대역폭이 제한 요소가 되는 대규모 배치 환경에서 두드러진다.
• 작은 배치 환경에서도 온라인 소프트맥스가 지연 시간과 메모리 접근 오버헤드를 줄여 1.5배에서 2.5배의 성능 향상을 기록한다.
• 수치적 안정성을 유지하며 기존 딥 러닝 프레임워크와 호환되어 즉시 적용 가능한 최적화를 제공한다.
• 성능 향상 효과는 히에라르키컬 소프트맥스, SVD-소프트맥스, 중요도 샘플링과 같은 다른 소프트맥스 최적화 기법과 독립적이므로, 병합 시 추가적인 가속을 가능하게 한다.