탄력적 어텐션 테스트 시 적응형 희소성 비율로 효율적인 트랜스포머 구현

초록

본 논문은 기존 하이브리드 어텐션이 고정된 희소성 비율을 사용해 다양한 작업에 최적화되지 못하는 문제를 해결하고자, 사전 학습된 모델에 경량화된 Attention Router를 삽입해 입력에 따라 어텐션 헤드의 전체 희소성을 동적으로 조절하는 Elastic Attention을 제안한다. 12시간·8×A800 GPU 학습만으로 LLM의 성능 저하 없이 추론 효율을 크게 향상시켰으며, 장문 처리 벤치마크에서 기존 방법들을 능가한다.

상세 분석

Elastic Attention은 두 가지 핵심 아이디어를 결합한다. 첫 번째는 “희소성‑강건(sparsity‑robust) 작업”과 “희소성‑민감(sparsity‑sensitive) 작업”이라는 작업군을 사전에 정의하고, 모델이 입력을 기반으로 어느 군에 속하는지를 판단해 전체 희소성 비율을 자동으로 조정한다는 점이다. 실험 결과, 요약과 같은 강건 작업은 희소 비율이 크게 변해도 성능이 유지되는 반면, QA와 같은 민감 작업은 일정 임계값을 넘으면 급격히 성능이 떨어진다. 따라서 고정된 희소성 비율을 사용하는 기존 하이브리드 어텐션은 모든 작업에 최적화되지 못한다는 한계가 명확히 드러난다.

두 번째는 경량화된 Attention Router 모듈이다. 각 레이어의 Key 텐서를 풀링해 얻은 헤드‑별 요약 벡터를 입력으로, Task MLP와 Router MLP를 차례로 통과시켜 헤드‑별 라우팅 로그it을 생성한다. 이 로그it은 Gumbel‑Softmax와 STE(직접‑통과 추정)를 이용해 연속적으로 이산 라우팅으로 변환된다. 라우팅 결과는 헤드가 Full‑Attention(FA) 혹은 Sparse‑Attention(SA) 중 어느 모드로 연산할지를 결정한다. 라우팅 결정은 학습 단계에서도 하드 라우팅을 사용해 추론 시와 동일한 동작을 보장하면서, Gumbel‑Softmax의 온도 스케줄링으로 점진적으로 확정성을 높인다.

희소성 비율 자체는 두 가지 메트릭으로 정의된다. Model Sparsity Ratio(Ω_MS​R)은 전체 헤드 중 SA 헤드의 비율을, Effective Sparsity Ratio(Ω_ES​R)은 실제 토큰 프루닝 비율을 반영한다. 학습 목표는 언어 모델링 손실에 sparsity regularization을 추가하고, 라우팅이 목표 희소성 구간(하한·상한) 내에 머무르도록 라그랑지안 승수를 동적으로 조정하는 min‑max 최적화이다. 이 설계는 모델의 기본 파라미터를 동결하고 Router 파라미터만 업데이트함으로써, 기존 LLM의 사전 지식을 그대로 보존한다.

효율성 측면에서 저자들은 모든 라우팅된 헤드를 하나의 fused kernel으로 병합해 단일 forward pass에서 동시에 처리한다. 이는 기존에 FA와 SA를 별도 연산으로 구현할 때 발생하는 메모리·연산 파편화를 크게 감소시킨다. 실험에서는 Qwen‑3 시리즈와 Llama‑3.1‑8B‑Instruct를 대상으로, 스트리밍 및 블록‑스파스 SA 패턴을 모두 적용했으며, 3가지 장문 벤치마크(LongBench‑V2, 실세계 검색, 장문 추론)에서 평균 0.60.8 수준의 Ω_MS​R 감소에도 불구하고 성능이 12% 상승했다. 특히, “FA‑SSA”(Full‑Attention + Sparse‑Self‑Attention)와 “FA‑XA”(Full‑Attention + eXternal‑Attention) 두 설정 모두 기존 DuoAttention, PruLong 등과 비교해 동일하거나 더 높은 정확도를 기록하면서, 추론 속도는 1.2~2배 가량 개선되었다.

학습 비용도 눈에 띈다. 전체 파라미터는 고정된 채 Router 파라미터만 0.27M per layer 정도 추가되며, 8×A800 GPU에서 12시간만에 수렴한다. 이는 대규모 파인튜닝이 필요 없는 효율적인 적응형 희소성 제어 메커니즘으로, 실제 서비스 환경에서 다양한 길이와 도메인의 입력에 대해 자동으로 최적의 연산량을 선택할 수 있게 한다.

요약하면, Elastic Attention은 (1) 작업별 희소성 특성을 사전 분석해 두 가지 범주로 단순화하고, (2) 입력‑조건부 헤드‑레벨 라우팅을 통해 실시간으로 희소성 비율을 조정하며, (3) 경량화된 라우터와 fused kernel을 통해 추론 효율성을 극대화한다는 세 가지 혁신을 제공한다. 이는 장문 처리 LLM의 확장성을 크게 향상시키면서도 기존 모델 성능을 유지하거나 개선하는 실용적인 솔루션이다.