[논문 리뷰] Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Zichang Liu, Jue Wang외 9명|arXiv (Cornell University)|2023. 10. 26.

Topic Modeling인용 수 19

한 줄 요약

LLMs에 맥락적 희소성을 도입하여 입력에 의존하는 작은 부분집합의 어텐션 헤드와 MLP 뉴런을 예측함으로써 추론 속도를 높이고, 재훈련 없이도 인-컨텍스트 학습을 희생하지 않으면서 실제 시간 기준 속도 향상을 가능하게 한다.

ABSTRACT

Large language models (LLMs) with hundreds of billions of parameters have sparked a new wave of exciting AI applications. However, they are computationally expensive at inference time. Sparsity is a natural approach to reduce this cost, but existing methods either require costly retraining, have to forgo LLM's in-context learning ability, or do not yield wall-clock time speedup on modern hardware. We hypothesize that contextual sparsity, which are small, input-dependent sets of attention heads and MLP parameters that yield approximately the same output as the dense model for a given input, can address these issues. We show that contextual sparsity exists, that it can be accurately predicted, and that we can exploit it to speed up LLM inference in wall-clock time without compromising LLM's quality or in-context learning ability. Based on these insights, we propose DejaVu, a system that uses a low-cost algorithm to predict contextual sparsity on the fly given inputs to each layer, along with an asynchronous and hardware-aware implementation that speeds up LLM inference. We validate that DejaVu can reduce the inference latency of OPT-175B by over 2X compared to the state-of-the-art FasterTransformer, and over 6X compared to the widely used Hugging Face implementation, without compromising model quality. The code is available at https://github.com/FMInference/DejaVu.

연구 동기 및 목표

• 사전 학습된 LLM에서 맥락적 희소성이 존재하며 추론 시 속도 향상을 위해 이를 활용할 수 있음을 동기 부여하고 입증한다.
• 레이어별로 맥락 의존적 희소 부분집합을 즉시 식별하는 저비용 예측기를 개발한다.
• 엔드-투-엔드 대기 시간을 줄이기 위한 비동기적이고 하드웨어 의식적인 구현을 제안한다.
• 여러 벤치마크에서 OPT-175B에 대한 속도향상과 정확도 보존을 검증한다.

제안 방법

• 맥락적 희소성 아래에서 희소화된 MLP 및 MHA(block) 블록을 정의한다.
• 맥락적 희소성이 입력당 약 85%의 구조적 희소성에 도달할 수 있음을 보이고, 입력당 최대 7배의 파라미터 감소를 가능하게 한다.
• 입력을 기반으로 중요한 뉴런/헤드를 식별하도록 훈련된 MLP용 및 어텐션용 이두 단계 희소성 예측기를 개발한다.
• MLP 희소성을 근접 이웃 탐색 문제(MaxIP)로 형식화하고 관련 뉴런을 선택하기 위한 빠른 신경 예측기를 구현한다.
• 마찬가지로 어텐션 헤드의 희소성도 근접 이웃 탐색으로 형식화하고 입력별로 사용할 헤드를 예측한다.
• 예측 오버헤드를 숨기기 위한 비동기적 룩어헤드 예측기를 도입하여 엔드-투-엔드 대기 시간 이득을 유지한다.
• GPU에서 속도향상을 실현하기 위한 하드웨어 인식 커널 융합 및 메모리 공동처리(memory-coalescing) 전략을 제공한다.

실험 결과

연구 질문

• RQ1 실용적 입력에 대해 사전 학습된 LLM에서 맥락적 희소성이 존재하는가?
• RQ2 재훈련 없이 즉석에서 맥락적 희소성을 정확히 예측할 수 있는가?
• RQ3 현대 하드웨어에서 모델 품질과 인-컨텍스트 학습을 보존하면서 엔드-투-엔드 대기 시간을 의미 있게 줄일 수 있는가?
• RQ4 맥락적 희소성으로 속도향상을 실현하기 위한 효과적인 하드웨어 인식 구현은 무엇인가?
• RQ5 제안된 접근 방식이 잔여 연결(residual connections)과 층 간 느리게 바뀌는 임베딩과 어떻게 상호 작용하는가?

주요 결과

• 사전 학습된 LLM에서 맥락적 희소성이 존재하며 주어진 입력에 대해 최대 약 85%의 파라미터가 맥락적으로 비활성화된다.
• 입력 및 층 맥 context에 기반한 예측기가 MLP 및 어텐션 블록의 희소 부분집합을 정확하게 식별할 수 있다.
• OPT-175B에 대해 FasterTransformer에 비해 엔드-투-엔드 대기 시간이 2배 이상 감소하고 Hugging Face에 비해 6배 이상 감소하며 품질 손실은 없다.
• 이 접근법은 인-컨텍스트 학습 능력을 보존하고 평가된 작업에서 제로샷/적은샷 성능을 유지한다.
• 비동기적 룩어헤드 예측기와 하드웨어 인식 구현이 실제로 벽시계 기준 속도향상을 달성하는 데 핵심이다.