🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

이 논문은 히스토그램 기반 초기화와 스파스성 인식 GPU 구현을 통해 $\alpha$ -entmax 어텐션의 계산 오버헤드를 획기적으로 줄여, FlashAttention-2 와 경쟁 가능한 속도로 장문맥 학습 시 효율성을 극대화하는 'AdaSplash-2'를 제안합니다.

원저자: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

게시일 2026-04-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

📚 배경: AI 가 긴 글을 읽을 때 겪는 문제

지금까지 AI(트랜스포머 모델) 는 긴 문서를 읽을 때, 모든 단어를 서로 비교하는 방식을 썼습니다.

비유: 100 페이지짜리 책을 읽을 때, 사서가 1 페이지의 단어와 100 페이지의 단어를 모두 비교하고, 2 페이지의 단어와 100 페이지의 단어도 모두 비교하는 식입니다.
문제: 책이 100 페이지일 때는 괜찮지만, 10,000 페이지가 되면 사서가 할 일이 기하급수적으로 늘어나서 시간이 너무 오래 걸리고 메모리도 폭발합니다. (이를 '이차 복잡도'라고 합니다.)

💡 해결책 1: 'α-entmax'라는 새로운 사서 (기존 연구)

연구자들은 "모든 단어를 다 비교할 필요 없어! 중요한 단어만 골라보자!"라고 생각했습니다. 그래서 α-entmax라는 새로운 방식을 도입했습니다.

특징: 이 방식은 중요하지 않은 단어는 아예 0(영) 으로 만들어서 무시합니다. 즉, 사서가 "이 단어는 중요하지 않으니 책장 넘겨버려!"라고 말합니다.
장점: 불필요한 작업을 줄여서 메모리를 아낄 수 있습니다.
단점: 하지만 "어떤 단어가 중요하고 어떤 게 중요하지 않은지"를 결정하는 과정이 매우 복잡하고 계산이 느립니다. 마치 사서가 책을 한 번 더 훑어보면서 "이게 진짜 중요할까?"라고 고민하는 시간이 너무 길어서, 오히려 전체 속도가 느려지는 문제가 있었습니다.

🚀 ADASPLASH-2: 더 똑똑하고 빠른 사서 (이 논문의 핵심)

이 논문은 바로 그 느린 계산 과정을 획기적으로 개선했습니다. ADASPLASH-2는 다음과 같은 두 가지 마법 같은 기술을 사용합니다.

1. "한눈에 보는 히스토그램" (Histogram-based Initialization)

기존 방식은 정확한 답을 찾기 위해 수백 번의 계산을 반복했습니다. 하지만 ADASPLASH-2 는 **히스토그램 (도수 분포표)**을 활용합니다.

비유: 사서가 모든 단어를 하나하나 세지 않고, **"중요도 1~~10 점대 단어는 몇 개, 11~~20 점대는 몇 개"**라고 **대략적인 표 (히스토그램)**를 먼저 그립니다.
효과: 이 표를 보면 "아, 대략 15 점 정도가 기준선이겠구나!"라고 순간적으로 추측할 수 있습니다. 이 추측이 매우 정확해서, 이제부터는 1~2 번만 계산하면 정확한 답을 얻을 수 있습니다.
결과: 계산 횟수가 100 번에서 1~2 번으로 줄어든 셈입니다.

2. "빈 책장 건너뛰기" (Sparsity-aware GPU Implementation)

중요하지 않은 단어 (0 점) 가 많은 경우, 아예 그 부분을 계산하지 않고 건너뜁니다.

비유: 사서가 책장을 넘길 때, 빈 책장 (중요하지 않은 단어) 이 있는 구간은 아예 뛰어넘어서 다음 중요한 책장만 찾습니다.
기술: 이를 위해 **비트 (Bit)**라는 아주 작은 단위로 정보를 압축하여, GPU(컴퓨터의 두뇌) 가 빈 공간을 아주 빠르게 인식하고 스킵할 수 있게 했습니다.

🏆 결과: 왜 이것이 중요한가요?

속도 향상: 긴 문서 (긴 문맥) 를 다룰 때, 기존 방식보다 최대 2 배 이상 빨라졌습니다. 특히 문장이 길어질수록 (책이 두꺼워질수록) 이 기술의 이점이 더 커집니다.
성능 유지: 속도가 빨라졌다고 해서 AI 의 지능이 떨어지는 것은 아닙니다. 오히려 긴 문맥을 이해하는 능력이 기존 방식보다 더 좋아졌습니다.
실용성: 이 기술은 AI 가 긴 소설을 읽거나, 긴 회의록을 분석하거나, 긴 코드 파일을 이해할 때 실시간으로 더 빠르게 작동하게 해줍니다.

📝 한 줄 요약

"ADASPLASH-2 는 AI 가 긴 글을 읽을 때, '중요하지 않은 단어'를 대략적인 표로 빠르게 추측하고, 아예 건너뛰는 방식으로 계산 속도를 2 배 이상 높인 똑똑한 사서입니다."

이 기술 덕분에 앞으로 AI 는 훨씬 더 긴 내용을 처리하면서도 빠르고 정확하게 답변할 수 있게 될 것입니다.

ADASPLASH-2: 더 빠른 미분 가능한 희소 어텐션 (Faster Differentiable Sparse Attention) 기술 요약

이 논문은 트랜스포머 (Transformer) 모델의 장기 컨텍스트 학습 시 발생하는 이차적 (quadratic) 계산 비용 병목 현상을 해결하기 위해 제안된 ADASPLASH-2라는 새로운 하드웨어 인식 (hardware-aware) 희소 어텐션 메커니즘을 소개합니다. 기존 $\alpha$ -entmax 어텐션의 계산 오버헤드를 획기적으로 줄여, FlashAttention-2 와 경쟁하거나 더 빠른 속도를 내면서도 희소성 (sparsity) 의 이점을 유지하는 것이 핵심 목표입니다.

1. 문제 제기 (Problem)

트랜스포머의 병목 현상: 트랜스포머의 자기 어텐션 (self-attention) 메커니즘은 점곱 어텐션 점수 행렬 $S = QK^\top$ 을 계산하고 물리적으로 저장해야 하므로, 시퀀스 길이에 대해 이차적 ( $O(n^2)$ ) 시간 및 메모리 복잡도를 가집니다. 이는 장기 컨텍스트 학습의 주요 병목입니다.
Softmax 의 한계: Softmax 기반 어텐션은 모든 토큰에 양의 확률 질량을 분배하므로 본질적으로 밀집 (dense) 되어 있습니다. 장기 컨텍스트에서는 관련 없는 토큰에 주의가 분산되고 토큰 표현이 구별되지 않는 문제가 발생합니다.
$\alpha$ -entmax 의 계산적 비효율: $\alpha$ -entmax 는 입력에 의존적인 희소성 (exact zeros) 을 생성할 수 있는 미분 가능한 대안이지만, 정규화 상수 $\tau$ 를 계산하기 위해 반복적인 루트 찾기 (root-finding) 알고리즘이 필요합니다. 기존 방법들은 이 $\tau$ 계산을 위해 어텐션 점수를 여러 번 스캔해야 하므로, FlashAttention-2 와 같은 최적화된 Softmax 구현에 비해 학습 속도가 느렸습니다.

2. 방법론 (Methodology)

ADASPLASH-2 는 $\alpha$ -entmax 의 정규화 상수 $\tau$ 를 계산하는 비용을 줄이기 위해 온칩 SRAM 내의 히스토그램 기반 초기화를 도입했습니다.

핵심 아이디어

온칩 히스토그램 구축 (On-chip Histogram Construction):
- 키 (Key) 블록을 스트리밍하는 동안, 어텐션 점수의 coarse 히스토그램을 온칩 SRAM 내에서 실시간으로 구축합니다.
- 점수 구간을 이산화 (discretize) 하여 빈 (bin) 수 $B$ 개의 히스토그램을 유지하며, 이는 $O(B)$ 의 저장 공간만 필요로 합니다 ( $n$ 에 무관).
- 이 히스토그램은 실제 정규화 임계값 $\tau^*$ 에 대한 **하한 (lower bound)**을 제공하여, 매우 정확한 초기값을 제공합니다.
단일/이중 반복 정제 (One/Two-pass Refinement):
- 히스토그램 기반 초기 추정치 ( $\tau_h$ ) 를 사용하여, 이후의 루트 솔버 (root solver) 가 정확한 해 $\tau^*$ 에 수렴하는 데 필요한 반복 횟수를 보통 1~2 회로 줄입니다.
- $\alpha$ 값에 따라 Halley, Newton, Secant 방법을 혼합한 **보호된 하이브리드 솔버 (safeguarded hybrid solver)**를 사용하여 수치적 안정성을 보장합니다.
효율적인 희소성 활용 (Efficient Sparsity Exploitation):
- 비트 패킹 (Bit-packing): 0 이 아닌 어텐션 블록을 나타내는 이진 마스크를 32 비트 정수 (int32) 단위로 압축하여 저장합니다.
- GPU 네이티브 명령어: Forward/Backward pass 시, 0 인 블록을 건너뛰기 위해 GPU 의 find-next-set (fns) 및 population count (popc) 명령어를 사용하여 오버헤드 없이 효율적으로 트래버스합니다.

3. 주요 기여 (Key Contributions)

SRAM 기반 히스토그램을 통한 정규화: 밀집된 중간 결과를 생성하지 않고 온칩 메모리에서 히스토그램을 계산하여 $\tau$ 의 하한과 상한을 증명적으로 유도했습니다.
안전한 하이브리드 솔버와 단일 반복 정제: 히스토그램 초기화 덕분에 대부분의 경우 추가적인 스캔 없이 1 회 반복으로 정확한 $\tau$ 를 복원할 수 있습니다.
GPU 최적화 Triton 커널: 미세한 타일링 (fine-grained tiling) 과 경량 비트 패킹 인코딩을 통해 동적 희소성을 거의 오버헤드 없이 활용하는 커널을 설계했습니다.
강력한 실험 결과: 합성 데이터 및 언어 모델링 벤치마크에서 FlashAttention-2 와 비교하여 중등도~고도 희소성 구간에서 더 빠른 학습 속도를 보였으며, 장기 컨텍스트 성능에서도 Softmax 기반 모델을 능가했습니다.

4. 실험 결과 (Results)

성능 (Speed):
- 희소성 - 효율성 트레이드오프: 블록 희소성이 60% 이상인 중등도~고도 희소성 구간에서, FlashAttention-2 (CUDA 및 Triton 구현체) 보다 더 빠른 속도를 기록했습니다. 특히 희소성이 높을수록 속도가 2 배 이상 빨라졌습니다.
- Backward Pass: 희소성이 높은 환경에서 역전파 (backward pass) 시간이 크게 단축되어 전체 학습 시간을 줄였습니다.
모델 성능 (Accuracy):
- 단기 컨텍스트: 짧은 컨텍스트 길이 (4K) 에서도 Softmax 기반 모델과 동등하거나 더 나은 성능을 보였습니다.
- 장기 컨텍스트: 32K 이상의 긴 컨텍스트에서 RULER 및 HELMET 벤치마크를 통해 Softmax 모델보다 월등히 우수한 성능을 달성했습니다. 특히 변수 추적 (Variable Tracking) 및 단어 추출 (CWE, FWE) 과 같은 작업에서 큰 개선을 보였습니다.
- NAPE와의 시너지: $\alpha$ -entmax 와 NAPE (No Positional Encoding + ALiBi) 조합이 장기 컨텍스트 모델링에 가장 효과적이었습니다.

5. 의의 및 결론 (Significance)

ADASPLASH-2 는 $\alpha$ -entmax 와 같은 미분 가능한 희소 어텐션 메커니즘이 실제 대규모 언어 모델 (LLM) 학습에서 실용화될 수 있는 길을 열었습니다.

학습 효율성: 기존 $\alpha$ -entmax 구현체의 계산적 비효율성을 해결하여, FlashAttention-2 와 경쟁 가능한 속도를 제공함으로써 장기 컨텍스트 학습을 가능하게 합니다.
모델 성능 향상: 희소성을 통해 장기 컨텍스트에서의 정보 과부하 (attention dispersion) 를 방지하고, 모델이 관련 있는 정보에 집중하도록 하여 추론 능력을 향상시킵니다.
하드웨어 인식 설계: GPU 의 메모리 계층 구조 (HBM vs SRAM) 를 고려한 설계는 향후 다른 희소 어텐션 알고리즘 개발에도 중요한 시사점을 제공합니다.

결론적으로, ADASPLASH-2 는 장기 컨텍스트 처리가 필요한 차세대 트랜스포머 모델의 핵심 구성 요소로서, 속도와 정확도 모두에서 기존 밀집 어텐션을 대체할 수 있는 강력한 대안입니다.