기술 요약: 메모리 제약이 있는 LLM 추론을 위한 동적 계층적 희소 어텐션

문제 제기

자기 어텐션 메커니즘의 이차적 계산 복잡도 ($O(L^2)$) 는 엄격한 하드웨어 메모리 예산 하에서 장문맥으로 대규모 언어 모델 (LLM) 을 확장하는 데 있어 치명적인 병목 현상을 초래합니다. 이전 연구들은 LLM 의 어텐션 행렬이 본질적으로 희소하다는 것을 입증했으나, 기존 솔루션들은 다음과 같은 중대한 한계에 직면해 있습니다:

• 정적 희소 방법: Longformer 및 BigBird 와 같은 접근 방식은 고정된 패턴 (예: 슬라이딩 윈도우, 확장 어텐션) 에 의존합니다. 이는 어텐션 분포의 입력 의존적 변이에 적응할 수 없어 다양한 작업에서 성능 저하를 초래합니다.
• 동적 휴리스틱 방법: 최근의 동적 접근 방식 (예: MInference, H2O) 은 추론을 가속화하기 위해 사전 정의된 템플릿이나 휴리스틱 규칙 (예: "수직-대각선" 패턴) 을 활용합니다. 그러나 이러한 방법들은 매우 입력 의존적인 어텐션 희소성을 포착할 유연성이 부족하여 다양한 작업과 입력에 대한 일반화 능력을 제한합니다.

핵심적인 과제는 LLM 백본의 재학습 없이 인과적 의존성을 보존할 수 있으면서도 입력 적응형이고 계산 효율적인 희소 어텐션 메커니즘을 개발하는 것입니다.

방법론: 동적 계층적 희소 어텐션 (DHSA)

저자들은 표준 디코더 전용 트랜스포머를 위한 플러그인 모듈인 **동적 계층적 희소 어텐션 (DHSA)**을 제안합니다. DHSA 는 LLM 백본을 고정된 상태로 유지하면서 온라인으로 어텐션 희소성을 예측하도록 작동합니다. 이 프레임워크는 중요한 토큰 상호작용을 식별하는 복잡도를 줄이기 위해 2 단계 계층적 라우팅 프로세스를 사용합니다.

1. 계층적 라우팅 공식화

모든 $L \times L$ 토큰 쌍을 스코어링하는 대신, DHSA 는 문제를 2 단계 프로세스로 축소합니다:

• 1 단계: 청크 수준 예측: 입력 토큰 시퀀스는 가변 길이의 청크로 분할됩니다. 경량 경계 예측기는 로컬 키 벡터 윈도우를 기반으로 동적으로 청크 경계를 결정합니다. 그런 다음 청크 간 상호작용의 중요도를 추정하기 위해 청크 수준의 유사성 행렬이 구성됩니다.
• 2 단계: 토큰 수준 선택: 청크 수준 유사성 점수에 기반하여, 시스템은 각 쿼리 청크를 관련 키 토큰 인덱스의 소형 집합으로 라우팅합니다. 이 선택은 고점수 청크를 구성 토큰으로 확장하면서도 인과적 제약을 강제하여, 가장 영향력 있는 쿼리 - 키 상호작용을 보존하는 희소 어텐션 마스크를 생성합니다.

2. 주요 기술 구성 요소

• 동적 경계 감지: 신경망 기반 예측기 (인코더, 특징 융합, MLP 로 구성) 는 로컬 키 윈도우를 분석하여 의미론적 경계를 감지합니다. 이는 수동 주석이 필요하지 않도록 기본 모델의 밀집 어텐션 패턴에서 파생된 자동 소프트 레이블을 사용하여 학습됩니다. 예측기는 Non-Maximum Suppression (NMS) 을 사용하여 경계를 최종화하여, 경계가 잘 분리되고 노이즈에 강건하도록 보장합니다.
• 강건한 청크 표현: 가변 길이의 청크를 효과적으로 처리하기 위해 DHSA 는 길이 정규화 풀링 방식을 사용합니다. 임베딩의 접합 합을 계산하여 실제 청크 길이로 나눈 후 길이 정규화를 수행합니다. 이는 제로 패딩이 평균 표현을 희석하는 것을 방지하고 다양한 크기의 청크 간 안정적인 유사성 추정을 보장합니다.
• 하드웨어 중립적 백엔드: DHSA 는 다양한 하드웨어를 지원하기 위해 두 가지 상호 보완적인 백엔드로 구현됩니다:
  • 모델 패밀리 및 플랫폼 (CPU 포함) 전반의 광범위한 호환성을 위한 PyTorch SDPA 백엔드.
  • 전체 어텐션 행렬을 물리화하지 않고 스트리밍 소프트맥스를 활용하여 GPU 에서 높은 효율성을 제공하는 타일링 온라인-소프트맥스 백엔드.

주요 기여

1. 계층적 라우팅 공식화: 본 논문은 희소화를 청크 수준 라우팅 문제로 재정의합니다. 청크 - 청크 유사성 행렬을 예측하고 이를 사용하여 어텐션을 소형 토큰 집합으로 라우팅함으로써, DHSA 는 계산 오버헤드를 크게 줄이면서도 인과적 의미론을 보존합니다.
2. 콘텐츠 인식 분할: 저자들은 고정된 길이가 아닌 콘텐츠 변화에 기반하여 시퀀스를 분할하는 동적 경계 예측기를 도입합니다. 이는 청크 크기에 관계없이 안정적인 유사성 추정을 보장하는 길이 강건 표현 방식과 결합됩니다.
3. 실용적인 플러그인 모듈: DHSA 는 쿼리/키 프로젝션과 어텐션 백엔드 사이에 개입하는 플러그인 모듈로 설계되었습니다. LLM 백본의 미세 조정이 필요 없으며, 여러 오픈 가중치 모델 패밀리를 지원하고 GPU 및 CPU 환경 모두에서 실행됩니다.

실험 결과

LLaMA-3.1-8B(4 비트), Qwen2.5-3B, Gemma-2-2B 와 같은 모델을 사용하여 Needle-in-a-Haystack, LongBench, RULER 벤치마크에서 광범위한 평가가 수행되었습니다.

• 정확도: DHSA 는 매우 높은 희소 영역에서도 밀집 어텐션에 버금가는 정확도를 유지합니다. LongBench 에서 DHSA 는 유사한 프리필 비용 대비 블록 희소 어텐션보다 12~20% 의 상대적 정확도 향상을 달성합니다. Needle-in-a-Haystack 테스트에서 DHSA 는 고정 밀도 6.25% 로 100K 문맥 길이까지 높은 검색 정확도를 유지하는 반면, 정적 방법들은 중요한 정보가 제한된 범위 밖으로 떨어질 때 급격한 성능 저하를 겪습니다.
• 효율성:
  • 가속: 메모리 효율적인 타일링 백엔드를 사용하여 DHSA 는 128K 문맥 길이에서 밀집 어텐션 대비 최대 10 배의 프리필 가속을 제공합니다.
  • 메모리 제약: 단일 24GB GPU 에서 DHSA 는 Out-Of-Memory (OOM) 오류로 인해 밀집 어텐션이 실패하는 영역인 100K 문맥 길이까지 LLaMA-3.1-8B(4 비트) 를 확장할 수 있게 합니다.
  • 배치 처리: 메모리 제약 환경에서 DHSA 는 경량 for 루프를 통한 배치의 순차적 처리를 통해 배치형 FlashAttention-2 에서 흔히 발생하는 OOM 실패를 방지하고 더 빠른 추론을 제공합니다.

중요성 및 주장

본 논문은 DHSA 를 메모리 제약이 있는 장문맥 LLM 추론을 위한 효율적이고 적응적인 솔루션으로 위치시킵니다. DHSA 의 주요 중요성은 어텐션의 이론적 희소성과 실용적인 입력 적응형 구현 사이의 간극을 메우는 데 있습니다.

• 적응성: 템플릿 기반 방법과 달리, DHSA 는 토큰 임베딩에서 온라인으로 콘텐츠 적응형 희소성을 학습하여 수동 튜닝 없이 다양한 작업에서 가속화를 제공합니다.
• 확장성: 레이어당 어텐션 비용을 $O(L^2)$에서 약 $O(L \cdot N_b)$ (여기서 $N_b$는 토큰 예산) 로 줄임으로써, DHSA 는 소비자 등급 하드웨어에서 극도로 긴 문맥을 처리할 수 있게 합니다.
• 일반성: 이 프레임워크는 모델 중립적이고 하드웨어 호환성이 있어 GPU 및 CPU 백엔드를 모두 지원하므로, 다양한 오픈 가중치 모델 패밀리에서 장문맥 기능을 배포하기 위한 실용적인 도구가 됩니다.

저자들은 DHSA 가 계산 효율성과 검색 정확도 사이의 균형을 효과적으로 해결하여, 장문맥 모델링을 위한 정적 및 휴리스틱 희소 어텐션 방법에 대한 견고한 대안을 제공한다고 결론지었습니다.