블록 트랜스포머: 빠른 추론을 위한 전역-지역 언어 모델링

Block Transformer: Global-to-Local Language Modeling for Fast Inference

June 4, 2024

저자: Namgyu Ho, Sangmin Bae, Taehyeon Kim, Hyunjik Jo, Yireun Kim, Tal Schuster, Adam Fisch, James Thorne, Se-Young Yun

cs.AI

초록

본 논문은 자기 주의력(self-attention)의 추론 병목 현상을 완화하기 위해 계층적 전역-지역 모델링(global-to-local modeling)을 자기회귀 트랜스포머(autoregressive transformer)에 적용한 Block Transformer 아키텍처를 제안한다. 자기 주의력을 적용하기 위해서는 모든 이전 시퀀스의 키-값(Key-Value, KV) 캐시를 매 디코딩 단계마다 메모리에서 검색해야 한다. 이로 인해 배치 추론에서 KV 캐시 입출력(IO)이 주요 병목 현상으로 작용한다. 우리는 이러한 비용이 전역 문맥(global context)에 자기 주의력을 적용하는 데서 비롯된다는 점을 발견했으며, 따라서 전역 모델링의 고비용 병목 현상을 하위 계층으로 격리하고 상위 계층에서는 빠른 지역 모델링을 적용한다. 하위 계층에서 남아 있는 비용을 완화하기 위해 입력 토큰을 고정 크기 블록으로 집계한 후, 이 거친 수준에서 자기 주의력을 적용한다. 문맥 정보는 단일 임베딩으로 집계되어 상위 계층이 전역 주의력 없이 다음 토큰 블록을 디코딩할 수 있도록 한다. 전역 주의력 병목 현상에서 자유로운 상위 계층은 컴퓨팅 하드웨어를 최대한 활용하여 추론 처리량을 극대화할 수 있다. 전역 및 지역 모듈을 활용함으로써, Block Transformer 아키텍처는 동등한 복잡도(perplexity)를 가진 일반 트랜스포머 대비 10-20배의 추론 처리량 향상을 보여준다. 우리의 연구는 전역-지역 모델링의 새로운 적용을 통해 언어 모델 추론을 최적화하는 새로운 접근 방식을 소개한다. 코드는 https://github.com/itsnamgyu/block-transformer에서 확인할 수 있다.

English

This paper presents the Block Transformer architecture which adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks of self-attention. To apply self-attention, the key-value (KV) cache of all previous sequences must be retrieved from memory at every decoding step. Thereby, this KV cache IO becomes a significant bottleneck in batch inference. We notice that these costs stem from applying self-attention on the global context, therefore we isolate the expensive bottlenecks of global modeling to lower layers and apply fast local modeling in upper layers. To mitigate the remaining costs in the lower layers, we aggregate input tokens into fixed size blocks and then apply self-attention at this coarse level. Context information is aggregated into a single embedding to enable upper layers to decode the next block of tokens, without global attention. Free of global attention bottlenecks, the upper layers can fully utilize the compute hardware to maximize inference throughput. By leveraging global and local modules, the Block Transformer architecture demonstrates 10-20x gains in inference throughput compared to vanilla transformers with equivalent perplexity. Our work introduces a new approach to optimize language model inference through novel application of global-to-local modeling. Code is available at https://github.com/itsnamgyu/block-transformer.