[논문 리뷰] SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Hanrui Wang, Zhekai Zhang외 1명|arXiv (Cornell University)|2020. 12. 17.

Advanced Neural Network Applications참고 문헌 124인용 수 12

한 줄 요약

SpAtten는 메모리 액세스와 계산을 극적으로 줄이기 위해 캐스케이드 토큰 및 헤드 프루닝과 점진적 양자화를 결합한 하드웨어-소프트웨어 공동 설계의 희소 어텐션 아키텍처를 제안한다. 실시간으로 중요도가 낮은 토큰과 헤드를 동적으로 프루닝하고 신뢰도에 따라 비트 폭을 조정함으로써, 정확도 손실 없이 TITAN Xp GPU 대비 최대 4059배의 에너지 절감과 162배의 속도 향상을 달성한다.

ABSTRACT

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, attention becomes the compution bottleneck because of its quadratic computational complexity to input length, complicated data movement and low arithmetic intensity. Moreover, existing NN accelerators mainly focus on optimizing convolutional or recurrent models, and cannot efficiently support attention. In this paper, we present SpAtten, an efficient algorithm-architecture co-design that leverages token sparsity, head sparsity, and quantization opportunities to reduce the attention computation and memory access. Inspired by the high redundancy of human languages, we propose the novel cascade token pruning to prune away unimportant tokens in the sentence. We also propose cascade head pruning to remove unessential heads. Cascade pruning is fundamentally different from weight pruning since there is no trainable weight in the attention mechanism, and the pruned tokens and heads are selected on the fly. To efficiently support them on hardware, we design a novel top-k engine to rank token and head importance scores with high throughput. Furthermore, we propose progressive quantization that first fetches MSBs only and performs the computation; if the confidence is low, it fetches LSBs and recomputes the attention outputs, trading computation for memory reduction. Extensive experiments on 30 benchmarks show that, on average, SpAtten reduces DRAM access by 10.0x with no accuracy loss, and achieves 1.6x, 3.0x, 162x, 347x speedup, and 1,4x, 3.2x, 1193x, 4059x energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

연구 동기 및 목표

• 일반 목적 플랫폼인 CPU와 GPU에서 어텐션 메커니즘의 비효율성, 즉 높은 메모리 대역폭과 낮은 산술 집중도 문제를 해결한다.
• CNN과 RNN에 최적화된 기존 가속기들이 어텐션 메커니즘에는 적합하지 않은 한계를 극복한다.
• 인간 언어와 어텐션 헤드의 구조적 부재를 활용해 DRAM 액세스와 계산을 줄인다.
• 추론 시 입력 기반으로 동적으로 토큰과 헤드를 프루닝할 수 있는 도메인 특화 가속기를 설계하여 실시간 효율성 향상을 가능하게 한다.
• 메모리와 에너지 오버헤드를 최소화하여 BERT와 GPT-2와 같은 대규모 NLP 모델을 자원 제약이 있는 장치에 효율적으로 구현할 수 있도록 한다.

제안 방법

• 문장 길이에 따라 적응 가능한 프루닝 비율을 갖는 누적 어텐션 확률 기반의 캐스케이드 토큰 프루닝을 제안한다.
• 층을 가로질러 누적된 헤드 크기 점수를 사용해 재중복 어텐션 헤드를 순위 매기고 제거하는 캐스케이드 헤드 프루닝을 도입한다.
• 실시간으로 토큰과 헤드 중요도 점수를 순위 매기기 위해 O(n) 복잡도를 갖는 고처리량 병렬 top-k 엔진을 설계하여 실시간 프루닝을 가능하게 한다.
• 가장 중요한 비트(MSB)로만 계산을 수행하고, 신뢰도가 낮을 경우 가장 낮은 비트(LSB)를 확보해 재계산함으로써 메모리 감소를 위한 점진적 양자화를 구현한다.
• 이론적 절감 효과를 실현하기 위해 파ipel라인 데이터패스, 최적화된 메모리 계층 구조, 동적 희소성 및 가변 비트 폭을 지원하는 전용 하드웨어 아키텍처를 구축한다.
• 프루닝과 양자화가 입력 기반이며 추론 시점에 적용되도록 알고리즘과 하드웨어를 공동 설계함으로써 유연성과 최소한의 정확도 손실을 확보한다.

실험 결과

연구 질문

• RQ1 동적으로 입력 기반으로 토큰과 헤드를 프루닝하는 것이 정확도 손실 없이 어텐션 계산과 메모리 액세스를 줄일 수 있는가?
• RQ2 저지연성의 top-k 엔진을 활용해 실시간으로 중요한 토큰과 헤드를 순위 매기고 선택하는 하드웨어 가속기가 어떻게 효율적으로 지원될 수 있는가?
• RQ3 점진적 양자화가 추론 신뢰도를 유지하면서 메모리 대역폭과 에너지 소비를 얼마나 줄일 수 있는가?
• RQ4 상태 기반 가속기인 A³과 MNNFast에 비해 캐스케이드 프루닝과 점진적 양자화가 성능과 에너지 효율성에서 얼마나 우수한가?
• RQ5 제안된 공동 설계 방법론이 분류형(BERT 등)과 생성형(GPT-2 등) 트랜스포머 모델 전반에 일반화될 수 있는가?

주요 결과

• SpAtten는 30개 벤치마크 평균으로 정확도 손실 없이 DRAM 액세스를 10.0배 감소시켰다.
• A³과 MNNFast 가속기 대비 각각 1.6배와 3.0배의 속도 향상, 1.4배와 3.2배의 에너지 절감을 달성했다.
• TITAN Xp GPU에서 SpAtten는 162배의 속도 향상과 1193배의 에너지 절감을 달성했다.
• Xeon CPU에서 SpAtten는 347배의 속도 향상과 4059배의 에너지 절감을 달성했다.
• top-k 엔진은 O(n) 시간 복잡도를 갖는 고처리량 토큰 및 헤드 중요도 순위 매기기를 효율적으로 가능하게 하여 실시간 프루닝에 핵심적인 역할을 한다.
• 점진적 양자화를 통해 동적 비트 폭 적응이 가능해져 메모리 대역폭을 감소시키면서도 반복적 정밀화를 통해 추론 신뢰도를 유지한다.