기술 요약: SOCKET - 희소 어텐션을 위한 SOft Collision Kernel EsTimator

문제 제기

대형 언어 모델 (LLM) 은 긴 컨텍스트 추론 중 근본적인 확장성 한계에 직면합니다. 프리필링 (prefilling) 단계는 전체 입력 시퀀스를 처리하는 반면, 반복적 디코딩 단계에서는 새로 생성된 토큰마다 이전에 캐시된 모든 토큰에 주의를 기울여야 합니다. 컨텍스트 길이가 증가함에 따라 (2K 에서 수백만 토큰까지) 이 밀집 어텐션 메커니즘은 점점 더 메모리 병목 현상을 겪게 되어, 종종 GPU 용량을 초과하고 비용이 많이 드는 CPU-GPU 전송을 필요로 합니다.

기존의 희소 어텐션 방법들은 계산량을 토큰의 부분 집합으로 제한함으로써 이를 완화하려 시도합니다. 그러나 많은 방법들이 높은 첫 번째 토큰 생성 시간 (TTFT) 비용을 초래하고 분포 변화에 대한 견고성이 부족한 데이터 의존적 휴리스틱 (예: k-means 클러스터링) 에 의존합니다. 다른 방법들은 이진 충돌 신호를 할당하는 전통적인 국소 민감 해싱 (LSH) 을 활용합니다. 본 논문은 이러한 "하드 (hard)" 형식이 순위 안정성에 적합하지 않다고 주장합니다. 이진 충돌은 부분적 유사성을 표현할 수 없어 거칠고 노이즈가 많은 순위를 초래하며, 특히 엄격한 메모리 예산 하에서 검색 품질을 저하시킵니다.

방법론: SOCKET

저자들은 LSH 를 이진 후보 생성기에서 원칙적이고 유사성을 인식하는 스코어링 커널로 재해석하는 데이터 무관한 희소 어텐션 메커니즘인 SOCKET(SOft Collision Kernel EsTimator) 을 소개합니다.

핵심 메커니즘: 소프트 LSH

단일 토큰을 단일 이산 버킷에 할당하는 전통적인 LSH 와 달리, SOCKET 은 소프트 LSH를 사용합니다:

1. 사전 계산 (프리필링): 키를 무작위 투사를 사용하여 $L$개의 독립적인 해시 테이블에 해시합니다. 각 키는 테이블당 단일 버킷 ID 를 할당받고, 값 벡터의 노름이 캐시됩니다.
2. 소프트 할당 (디코딩): 새로운 쿼리가 도착하면 단일 버킷에 할당되지 않습니다. 대신, 각 해시 테이블의 모든 버킷에 대한 확률 분포를 유도합니다. 이는 쿼리의 무작위 초평면 투사를 기반으로 한 소프트 해시 점수를 계산하고, 온도 제어 소프트맥스 함수를 통과시킴으로써 달성됩니다.
3. 스코어링: 각 키에 대해, 모든 $L$개의 테이블에 걸쳐 해당 키의 특정 버킷에 할당된 확률 질량을 집계하여 연속적인 점수를 계산합니다. 이 점수는 키의 값 벡터 노름으로 가중치가 부여됩니다.
4. 선택: 상위 $k$개의 키가 이러한 소프트 점수를 기반으로 결정적으로 선택되며, 이 부분 집합에 대해 정확한 어텐션이 계산됩니다.

이론적 기반

본 논문은 소프트 LSH 점수를 사용하여 계산된 어텐션이 소프트맥스 어텐션과 밀접하게 관련된 커널 (구체적으로 각도 커널) 을 근사한다는 이론적 보장을 제공합니다.

• 안정성: 이론적 분석 (정리 3) 은 소프트 버킷화가 온도 $\tau \to 0$에 따라 사라지는 제어 가능한 편향 ($\epsilon_\tau$) 을 도입하면서도 부드러움을 유지함을 보여줍니다.
• 상관관계: 보조 정리 4 와 실험 결과는 소프트 점수가 하드 점수보다 투사된 쿼리 - 키 유사성의 방향 구조를 더 잘 보존함을 보여줍니다. 하드 LSH 는 $\ell_1$ 크기에 의존하고 사분면 내의 방향 정보를 폐기하여 유한한 $L$에서 불안정한 순위를 초래합니다. 소프트 LSH 는 작은 신호 영역에서 투사된 좌표와 선형 관계를 유지하여 실제 유사성과 더 높은 상관관계와 더 낮은 추정기 분산을 제공합니다.

주요 기여

1. 안정적인 랭커로서의 소프트 LSH: 저자들은 전통적인 하드 LSH 보다 훨씬 효과적인 랭킹 함수로 작용하는 데이터 무관한 소프트 LSH 메커니즘을 제안합니다. 이는 이진 충돌 신호를 확률적이고 유사성을 인식하는 집계로 대체하여, 임의의 투표 없이 안정적인 상위 $k$개 선택을 가능하게 합니다.
2. 커스텀 커널을 통한 효율성: 커스텀 CUDA 스코어링 커널과 Flash Decode Triton 백엔드가 구현되었습니다. 이를 통해 SOCKET 은 긴 컨텍스트에서 디코딩 시 FlashAttention 보다 최대 1.5 배 높은 처리량을 달성하면서도 기존 희소 방법들의 정확도를 유지하거나 능가할 수 있습니다.
3. 데이터 무관 설계: 클러스터링 기반 접근법과 달리 SOCKET 은 재학습, 보정 또는 오프라인 전처리가 필요하지 않습니다. 이는 더 빠른 TTFT 와 분포 변화에 대한 견고성을 가능하게 합니다.
4. 이론적 보장: 본 논문은 소프트 카운트 어텐션 추정기에 대한 엔드 투 엔드 오차 한계를 수립하여 오차를 샘플링 분산, 유한 테이블 근사 오차, 소프트 버킷화 편향으로 분해합니다.

실험 결과

저자들은 Llama-3.1-8B-Instruct, Llama-3.2-1B, 그리고 Qwen3모델 (4B, 8B, 30B) 을 두 가지 벤치마크인 LongBench(실제 세계 작업) 와 RULER(합성 진단 작업) 에서 평가했습니다.

• 정확도: LongBench 에서 SOCKET 은 Quest, PQCache, Double Sparsity, MagicPig 과 같은 베이스라인을 일관되게 능가했습니다. 예를 들어, 33 배 희소성에서 Llama-3.1-8B 의 경우 SOCKET 은 평균 점수 47.83 을 달성하여 가장 강력한 베이스라인 (Quest) 보다 약 0.84 점 더 높은 성능을 보였습니다. RULER-32K 에서는 특히 높은 희소성 수준 (예: 50 배) 에서 기존 방법과 맞먹거나 능가하는 성능을 보였습니다.
• 효율성: SOCKET 은 NVIDIA A100 및 H200 GPU 에서 상당한 처리량 향상을 보여주었습니다. H200 에서 컨텍스트 길이가 증가함에 따라 FlashAttention-2 및 FlashAttention-3 을 능가하여, 100K 토큰에서 FlashAttention-2 보다 1.58 배, 140K 토큰에서 FlashAttention-3 보다 1.58 배의 속도 향상을 달성했습니다 (33 배 희소성 기준).
• 메모리 오버헤드: 이 방법은 해시 테이블 인덱스와 값 노름을 저장하는 표준 KV 캐시 외에 약 15% 의 modest 한 메모리 오버헤드를 도입하며, 이는 동등한 정확도를 달성하기 위해 하드 LSH 가 필요로 하는 많은 수의 해시 테이블에 비해 훨씬 낮습니다.

중요성과 주장

본 논문은 SOCKET 이 기존 희소 어텐션 방법에 대한 원칙적이고 데이터 무관한 대안을 제공함으로써 긴 컨텍스트 LLM 추론의 결정적 병목 현상을 해결한다고 주장합니다. 하드, 이진 해싱에서 소프트, 확률적 스코어링으로 전환함으로써 SOCKET 은 우수한 정확도 - 효율성 트레이드오프를 달성합니다.

저자들은 그들의 접근 방식이 다음과 같음을 강조합니다:

• 데이터 의존적 전처리의 필요성 제거: 이를 통해 TTFT 를 줄이고 배포 유연성을 향상시킵니다.
• 이론적 안정성 제공: 제한된 수의 해시 테이블에서도 토큰 순위가 일관되게 유지되도록 보장하여 전통적 LSH 의 약점을 해결합니다.
• 실용적인 시스템 수준의 이득 제공: 스코어링 커널의 알고리즘적 개선이 모델 성능을 희생하지 않고 현대 GPU 아키텍처에서 직접적인 처리량 향상으로 이어질 수 있음을 입증합니다.

이 연구는 소프트 LSH 를 단순한 검색 도구가 아닌, 디코딩 단계에서 밀집 어텐션을 대체할 수 있는 견고한 스코어링 커널로 위치시킴으로써, LLM 이 수백만 토큰의 컨텍스트로 효율적으로 확장되도록 합니다.