DistrAttention: 최신 GPU에서 효율적이고 유연한 자기 주의 메커니즘

Transformer 아키텍처는 자연어 처리, 컴퓨터 비전, 시계열 예측 등 분야에서 최첨단 성능을 제공하며 딥러닝에 혁신을 가져왔습니다. 그러나 그 핵심 구성 요소인 자기 주의는 입력 시퀀스 길이에 대해 이차 시간 복잡도를 가지며, 이는 Transformer의 확장성을 제한합니다. 기존의 자기 주의 최적화 접근법은 전체 문맥 정보를 버리거나 유연성이 부족합니다. 본 연구에서는 전체 문맥을 유지하면서 효율적이고 유연한 자기 주의 메커니즘인 DistrAttention을 설계했습니다. DistrAttention은 임베딩 차원 수(d라고 부름)를 기준으로 데이터를 그룹화함으로써 이를 실현합니다. 우리는 유사 데이터를 그룹화하기 위해 지역 민감 해싱(locality-sensitive hashing)을 활용하는 경량 샘플링 및 융합 방법으로 DistrAttention을 구현했습니다. 또한 지역 민감 해싱으로 인한 오류를 제한하기 위한 블록 단위 그룹화 프레임워크도 설계했습니다. 블록 크기 선택을 최적화함으로써 DistrAttention은 FlashAttention-2와 쉽게 통합되어 최신 GPU에서 높은 성능을 얻을 수 있습니다. 광범위한 실험을 통해 DistrAttention을 평가했으며, 결과는 자기 주의 계산에서 FlashAttention-2보다 37% 빠름을 보여줍니다. ViT 추론에서는 근사 자기 주의 메커니즘 중에서 DistrAttention이 가장 빠르고 가장 정확합니다. Llama3-1B 환경에서는 1% 정확도 손실만으로도 가장 낮은 추론 시간을 달성합니다.