Dynamic Ultrametric Attention: 하드웨어 최적화 희소성을 스스로 학습하는 Transformer

섹션별 상세

Dense self-attention은 시퀀스 길이에 따라 제곱으로 연산량이 증가하여 긴 문맥 처리에 병목을 유발한다. 이를 해결하기 위해 제안된 Dynamic Ultrametric Attention은 Transformer가 훈련 과정에서 헤드별 블록 희소성 라우팅 토폴로지를 스스로 학습하게 한다.

이 프레임워크는 훈련 중 Gumbel-Sigmoid depth gates를 사용하여 울트라메트릭 거리 행렬 기반의 희소성 패턴을 학습한다. 학습된 패턴은 추론 시 커스텀 Triton 블록 희소 커널로 오프로드되어, 어텐션이 필요 없는 블록에 대한 SRAM 로드를 건너뛰어 연산 효율을 높인다.

실험 결과, 2048 토큰에서 11.59배, 8192 토큰에서 28배의 추론 속도 향상을 달성했다. 또한 메모리 사용량은 98.4% 절감되었으며, 희소 PagedAttention 디코딩 커널은 KV 캐시 블록 로드를 조건부로 건너뛰어 8배의 유효 메모리 대역폭을 확보했다.

이 방식은 사후 가지치기나 증류 없이 LLM이 스스로 하드웨어 최적화된 희소성 패턴을 학습하는 첫 사례이다. 실제 자연어 작업에서 88% 이상의 희소성을 유지하면서도 교차 엔트로피 손실을 10.9에서 1.55로 낮추어 성능과 효율을 동시에 입증했다.

언급된 도구

Triton추천

커스텀 블록 희소 커널 구현

PyTorch중립

Dense attention 성능 비교 베이스라인

PagedAttention추천

희소 디코딩 커널 구현

언급된 리소스

논문Learning to Skip Blocks