기술 요약: Pulsar Attention

문제 정의

긴 시퀀스에 대한 대규모 언어 모델(LLM)의 추론은 셀프 어텐션(self-attention)의 이차 복잡도와 키-값(KV) 캐시의 선형적 증가로 인해 계산적으로 매우 부담스럽습니다. 이는 장치 메모리를 빠르게 고갈시킵니다. Star Attention(Acharya et al., 2025)과 같은 분산 방식은 컨텍스트를 여러 호스트에 분할하고 블록을 병렬로 처리함으로써 이를 완화하지만, 이 방식은 **정적이고 내용에 무관한 앵커 블록(static, content-blind anchor block)**에 의존합니다. 이 앵커는 모든 호스트의 입력 앞에 붙여지는 첫 번째 블록의 전체 복사본으로, 소프트맥스(softmax) 분포를 안정화하는 역할을 합니다.

저자들은 이 접근 방식에서 두 가지 결정적인 한계를 식별했습니다:

1. 중복 계산: 정적 앵커는 앵커의 관련성과 관계없이 컨텍스트 인코딩(Phase 1) 중에 각 호스트의 시퀀스 길이를 두 배로 늘립니다.
2. 정보 손실: 중간 블록(블록 2부터 $i-1$까지)은 인코딩 과정에서 블록 $i$에 정보를 제공하지 않습니다. 시퀀스 길이가 길어질수록 정적 앵커는 필요한 교차 블록(cross-block) 컨텍스트의 비중을 점차 줄여나가며, 이는 어텐션 품질의 저하로 이어집니다.

방법론: Pulsar Attention

본 논문은 정적 앵커를 대체하기 위해, KV 캐시 점유율은 유지하면서 연산을 줄일 수 있는 두 가지 경량화된 내용 인식형(content-aware) 구성 요소를 사용하는 Pulsar Attention을 제안합니다.

Phase 1: 통계적 컨텍스트 인코딩

첫 번째 블록 전체를 복제하는 대신, 각 호스트는 다음으로 구성된 입력 시퀀스를 구축합니다:

1. 어텐션 싱크 접두사 (Attention-Sink Prefix, $S_{sink}$): 시퀀스의 맨 시작 부분에서 추출한 작고 고정된 접두사(기본값 64 토큰)입니다. Xiao et al. (2023)을 따라, 이는 전체 앵커 블록을 사용하는 것보다 훨씬 적은 비용으로 소프트맥스 분포를 안정화합니다.
2. Max-IDF 블록 요약 ($\Sigma$): 모든 인과적 이전 블록(causally preceding block)에 대해 시스템은 압축된 요약본을 생성합니다.
   • 휴리스틱: 시스템은 Max-IDF(역문서 빈도) 휴리스틱을 사용하여 중복되지 않는 연속적인 청크(기본값 32 토큰)를 스코어링합니다. 청크의 점수는 해당 청크 내의 토큰 중 최대 IDF 값입니다.
   • 근거: Max-IDF는 전역적으로 희귀한 토큰(예: 엔티티 이름, 식별자, 수치 값)을 포함하는 청크를 우선시합니다. 이는 단일 희귀 토큰이 일반적인 단어들에 의해 평균화되어 버리는 평균화 휴리스틱(예: TF-IDF)의 신호 희석 문제를 방지합니다.
   • 선택: 이 점수를 기준으로 블록당 상위-$k$개의 청크를 선택하고, 이를 위치 순서대로 연결하여 요약본 $\Sigma_i$를 형성합니다.
   • 효율성: IDF 테이블은 CPU 상에서 토큰 ID에 대해 단 한 번의 $O(L)$ 패스를 통해 구축되며, 호스트 간 통신이나 GPU 시간을 필요로 하지 않습니다.

호스트 $i$의 입력은 $[S_{sink} \parallel \Sigma_0 \parallel \dots \parallel \Sigma_{i-1} \parallel B_i]$가 됩니다. 순전파(forward pass) 이후, 싱크와 요약본에 대한 KV 엔트리는 폐기되며, 로컬 블록 $B_i$에 대한 KV 상태만을 유지합니다. 이를 통해 유지되는 KV 캐시 점유율은 Star Attention과 동일하게 유지됩니다.

Phase 2: 글로벌 쿼리 인코딩

Phase 2는 Star Attention과 구조적으로 동일합니다. 쿼리는 모든 호스트로 브로드캐스트되며, 각 호스트는 유지된 KV 샤드에 대해 로컬 어텐션을 계산합니다. 이 점수들은 온라인 소프트웨어 머지(online softmax merge; Milakov and Gimelshein, 2018)를 통해 집계되어, 전체 KV 캐시를 전송하지 않고도 정확한 글로벌 어텐션을 복구합니다.

주요 기여

1. 정적 방식에서 내용 인식 방식으로의 교체: 저자들은 정적 앵커를 작은 어텐션 싱크 접두사와 Max-IDF 통계적 요약으로 대체하여, Star Attention 대비 Phase 1의 per-GPU FLOPs를 최대 3.3배 줄이면서도 동일한 KV 캐시 점유율을 유지합니다.
2. Max-IDF 스코어링 메커니즘: 본 논문은 교차 블록 컨텍스트를 선택하기 위한 Max-IDF 휴리스틱을 도입하고 검증합니다. 절제 연구(Ablation studies)에 따르면, 희귀 토큰의 존재를 기반으로 청크를 선택하는 것이 TF-IDF, BM25, 엔트로피, 균등 샘플링보다 효과적으로 교차 블록 컨텍스트를 포착하여 더 나은 성능을 보였습니다.
3. 실증적 성능: RULER 벤치마크에서 Llama-3.1-8B-Instruct를 사용했을 때, Pulsar Attention은 128K 토큰의 컨텍스트 길이까지 Star Attention 및 Dense Attention을 모두 능가합니다. 이는 Dense baseline 대비 최대 4.7%, Star Attention 대비 **6.1%**의 절대적인 정확도 향상을 달면합니다.
4. 교차 아키텍처 적용 가능성: 이 방법은 아키텍처 특화 튜닝 없이도 Qwen3-4B-Instruct에서 경쟁력 있는 정확도를 보여주며, 이의 이식성을 확인시켜 줍니다.
5. 효율성 검증: 엔드 투 엔드 프로파일링 결과, 128K 컨텍스트 길이에서 Dense Attention 대비 최대 6.24배의 속도 향상을 확인하였으며, 전체 플릿(fleet) 메모리 감소량은 최대 5.1배에 달했습니다.

결과 및 발견

• RULER 벤치마크: Pulsar Attention은 32K–128K 범위에서 Star Attention과 Dense Attention을 지속적으로 앞섭니다. 특히 집계(Aggregation) 작업(예: Common Words Extraction)에서 주목할 만한 이득이 관찰되었는데, Pulsar는 81.2%의 정확도를 달성한 반면 Dense 및 Star Attention은 거의 0에 가까운 성능을 보였습니다. 이는 내용 인식형 요약이 불필요한 방해 토큰(distractor tokens)을 효과적으로 필터링함을 시사합니다.
• BABILong 벤치마크: 멀티 홉 추론 작업(qa3, qa5)에서 Max-IDF를 적용한 Pulsar는 Star Attention을 유의미하게 능가했습니다(각각 +11%, +5%). 이는 희귀 엔티티 토큰을 포함하는 청크를 선택함으로써 팩트를 블록 간에 체이닝할 수 있는 능력을 입증합니다.
• 한계점: 이 방법은 단일 키와 연관된 여러 값을 검색해야 하는 작업(예: NIAH MultiValue)에서 어려움을 겪습니다. 수치 값은 대개 전역적 희귀성이 낮기 때문에 낮은 IDF 점수를 받게 되며, 이로 인해 요약 선택 과정에서 제외될 수 있습니다.
• 효율성: 고정된 싱크 접두사(64 토큰)와 가변적인 요약본의 조합은 컨텍스트 길이가 길어질수록 전체 블록 크기에서 차지하는 비중이 줄어들며, 이에 따라 컨텍스트가 커질수록 Dense Attention 대비 속도 향상 폭이 넓어집니다.

의의 및 주장

본 논문은 내용 인식형 요약이 계산량을 절감할 뿐만 아니라, 불필요한 토큰을 필터링함으로써 어텐션 품질을 능동적으로 개선한다고 주장합니다. 정적이고 내용에 무관한 앵커를 통계적이고 훈련이 필요 없는 메커니즘으로 대체함으로써, Pulsar Attention은 다음과 같은 기존의 분산 추론 파이프라인(Star Attention 등)에 대한 **드롭인 교체(drop-in replacement)**를 제공합니다:

• APB와 같은 다른 연구들과 달리 감독 학습(supervised training)이나 프록시 모델이 필요 없습니다.
• 새로운 모델 제품군에 즉각적인 이식성을 제공합니다.
• 인코딩 단계의 계산 병목을 줄이면서 긴 컨텍스트 시나리오에서 우수한 정확도를 달성합니다.

저자들은 현재의 Max-IDF 휴리스틱이 희귀 토큰 검색에는 매우 효과적이지만, 희귀하지 않은 반복되는 값을 요구하는 작업을 위한 쿼리 인식형(query-aware) 선택 전략에 대해서는 향후 연구가 필요하다고 결론지었습니다.