기술 요약: 토큰 순열을 통한 더 희소한 블록-희소 어텐션

문제 제기

대규모 언어 모델 (LLM) 의 컨텍스트 길이를 확장하는 것은 주로 시퀀스 길이에 대한 $O(N^2)$ 복잡도를 갖는 자기 어텐션 (self-attention) 메커니즘으로 인해 계산 비용이 매우 많이 듭니다. 시퀀스를 타일링하여 메모리 오버헤드를 완화하는 하드웨어 인식 최적화 (예: FlashAttention) 가 존재하지만, 여전히 전체 어텐션 행렬을 계산합니다. 블록-희소 어텐션은 시퀀스를 블록으로 분할하고 사전 계산된 마스크에 따라 특정 블록을 건너뛰어 계산을 더 줄이려고 시도합니다.

그러나 기존 블록-희소 방법은 정보 분열 (information fragmentation) 으로 고통받고 있습니다. 긴 시퀀스에서 주어진 쿼리에 대한 중요한 키 토큰 (종종 "헤비 히터") 은 종종 여러 개의 비연속 블록에 흩어져 있습니다. 결과적으로 블록 기반 방법은 충분한 어텐션 질량을 포착하기 위해 많은 수의 블록을 검색해야 하며, 이는 계산적 중복과 비최적의 희소성으로 이어집니다. 현재 접근 방식은 분열된 행렬에서 수동으로 블록을 선택하는 반면, 행렬 구조 자체를 최적화하지는 않습니다.

방법론: 순열 블록-희소 어텐션 (PBS-Attn)

저자들은 어텐션 메커니즘의 순열 불변성을 활용하여 어텐션 행렬을 재구성함으로써 흩어진 어텐션 질량을 연속적인 고밀도 블록으로 통합하는 순열 블록-희소 어텐션 (PBS-Attn) 을 제안합니다. 이는 플러그 앤 플레이 방식의 방법입니다.

1. 이론적 기반

이 방법은 어텐션의 두 가지 핵심 속성에 의존합니다:

• 키 - 값 쌍 순열 불변성: 키 - 값 쌍의 짝이 유지된다면, 키 - 값 쌍의 순서가 순열되어도 출력은 변하지 않습니다.
• 쿼리 순열 공변성: 쿼리 시퀀스를 순열하면 출력도 해당되게 순열됩니다.
• 역변환: 쿼리와 키/값을 순열한 후 출력에 역순열을 적용하면 원래의 비순열 결과를 복원할 수 있습니다.

2. 인과성을 위한 세그먼트 순열

전역 순열은 쿼리가 이전 키에만 어텐션할 수 있어야 하는 자기회귀 LLM 에 필요한 인과적 마스킹을 파괴합니다. 이를 해결하기 위해 PBS-Attn 은 세그먼트 순열을 도입합니다:

• 입력 시퀀스는 크기가 $S$인 연속적이고 겹치지 않는 세그먼트로 나뉩니다.
• 순열은 세그먼트 내부 (intra-segment) 에 적용되지만 세그먼트 간 인과성 (inter-segment causality) 은 엄격히 유지됩니다.
• 이를 통해 모델은 중요한 키들을 클러스터링하기 위해 세그먼트 내의 토큰을 재배열할 수 있으면서도, 세그먼트 간의 의존성이 유효하도록 보장합니다 (즉, 세그먼트 $i$의 쿼리는 세그먼트 $j > i$의 키에 어텐션할 수 없습니다).

3. 전역 중요도 기반 키 순열

각 세그먼트에 대한 최적의 순열을 결정하기 위해, 이 방법은 경량의 프록시 전략을 사용합니다:

• 헤비 히터 클러스터링: 세밀한 지역 정렬 대신, 많은 쿼리로부터 어텐션을 끌어당기는 "전역적으로 중요한" 키 토큰 (헤비 히터) 을 식별합니다.
• 프록시 점수 매기기: 키의 전역 중요도는 쿼리의 작은 부분집합 (특히, 마지막 쿼리 블록) 을 사용하여 추정됩니다. 전역 중요도 점수 벡터는 어텐션 점수의 평균 풀링을 통해 계산됩니다.
• 정렬: 각 세그먼트 내에서 키는 전역 중요도 점수의 내림차순으로 정렬됩니다. 이는 가장 중요한 토큰들을 각 세그먼트의 선두 블록으로 클러스터링하여, 더 적은 수의 블록으로 포착할 수 있는 어텐션 질량의 조밀한 영역을 생성합니다.
• 쿼리 처리: 저자들은 쿼리를 순열하는 것은 미미한 이득을 제공하고 그룹화 쿼리 어텐션 (GQA) 과 관련하여 비효율성을 초래할 수 있음을 발견했으므로, 기본 전략은 쿼리의 자연스러운 순서를 유지합니다.

4. 구현

• 블록 선택: 순열된 키를 기반으로 블록 중요도를 추정하기 위해 평균 풀링 전략이 사용됩니다.
• 커널: 이 방법은 Triton 에서 커스텀 순열-FlashAttention 커널을 사용하여 구현되며, 이는 순열, 블록 선택, 역순열을 효율적으로 처리합니다.

주요 기여

1. 구조적 최적화: 분열된 행렬에서의 수동 블록 선택에서 토큰 순열을 통한 능동적 구조 통합으로 최적화 패러다임을 전환합니다.
2. 세그먼트 순열 전략: 자기회귀 LLM 의 인과적 제약을 엄격히 유지하면서 순열을 적용하는 새로운 메커니즘을 제안합니다.
3. 플러그 앤 플레이 설계: 이 방법은 재학습이 필요 없으며, 아키텍처 변경 없이 기존 사전 훈련된 모델 (예: Llama-3.1, Qwen-2.5) 에 적용할 수 있습니다.
4. 효율성 향상: 헤비 히터를 클러스터링하는 것이 높은 어텐션 커버리지를 유지하는 데 필요한 블록 수를 크게 줄인다는 것을 입증합니다.

실험 결과

이 방법은 LongBench, LongBenchv2, RULER 벤치마크에서 Llama-3.1-8B 및 Qwen-2.5-7B-1M에 대해 평가되었습니다.

• 정확도: PBS-Attn 은 기존 블록-희소 방법 (Minference, FlexPrefill, XAttention) 보다 일관되게 우수한 성능을 보이며 전체 어텐션 기준과 밀접하게 일치합니다. LongBench 에서 두 모델 모두에서 가장 좋은 전체 성능을 달성했습니다. LongBenchv2 에서 전체 어텐션에 비해 최소한의 저하를 보였으며 다른 희소 방법보다 훨씬 우수한 성능을 보였습니다.
• 희소성: 순열은 블록 수준의 희소성을 증가시킵니다. 128K 컨텍스트 길이에서 이 방법은 97.5% 어텐션 커버리지를 유지하면서 비순열 기준 대비 선택된 블록 수를 14.4% 줄였습니다.
• 속도: 커스텀 커널을 기반으로 PBS-Attn 은 256K 컨텍스트 길이에서 프렐링 (prefilling) 시간 (첫 번째 토큰 도달 시간) 에서 최대 2.75 배의 엔드 투 엔드 속도 향상을 달성합니다. 속도 향상은 컨텍스트 길이에 따라 확장되는 반면, 다른 방법들은 극단적인 길이에서 정체되거나 제한된 이득을 보입니다.
• 강건성: 이 방법은 다양한 모델 크기 (최대 14B 파라미터) 에서 효과적이며, 사용되는 특정 블록 선택 알고리즘에 대해 강건합니다.

중요성 및 주장

이 논문은 PBS-Attn 이 효율적인 긴 컨텍스트 LLM 추론을 위한 실용적이고 실행 가능한 경로를 제공한다고 주장합니다. 단순한 더 나은 선택 휴리스틱이 아닌 구조적 재구성을 통해 비최적 희소성의 근본 원인 (정보 분열) 을 해결함으로써, 이 방법은 모델 성능을 훼손하지 않으면서 상당한 계산적 절감을 달성합니다.

저자들은 이 접근 방식의 오버헤드가 미미함을 강조합니다. 순열 및 점수 매기기 단계는 어텐션 계산 속도의 막대한 향상과 비교할 때 시간 및 메모리 비용을 무시할 수준으로 도입합니다 (예: 512K 컨텍스트에서 약 3% 오버헤드). 이 연구는 어텐션의 순열 속성을 활용하는 것이 긴 컨텍스트 처리를 최적화하는 데 있어 중요하지만 아직 충분히 활용되지 않은 축임을 시사합니다.