컨텍스트 집중을 통한 장문 확산 LLM 가속화

초록

Focus-dLLM은 이전 디노이징 단계의 토큰 신뢰도(Confidence)를 이용해 다음 단계에서 마스크 해제될 토큰을 예측하고, 중요한 “어텐션 싱크”를 보존하면서 희소 어텐션을 적용한다. KV 캐시와 결합한 이 방법은 32K 길이의 컨텍스트에서 29배 이상의 손실 없는 속도 향상을 달성한다.

상세 분석

본 논문은 확산 기반 대형 언어 모델(dLLM)의 비자율적 디코딩 특성 때문에 전체 토큰에 대해 양방향 풀 어텐션을 수행해야 하는 근본적인 병목을 해결하고자 한다. 두 가지 핵심 관찰이 제안의 토대가 된다. 첫째, 인접 디노이징 단계 사이의 토큰 신뢰도 점수는 강한 양의 상관관계를 보이며, 현재 단계에서 마스크 해제될 토큰은 바로 이전 단계에서 높은 신뢰도를 가진 토큰과 높은 겹침을 보인다. 실험적으로 상위 4개의 고신뢰 토큰을 이전 단계에서 선택했을 때, 다음 단계에서 실제 마스크 해제되는 토큰을 96.1% 정확도로 예측할 수 있음을 확인하였다. 이를 바탕으로 “과거 신뢰도 기반 인디케이터(past confidence‑guided indicator)”를 설계해, 이전 단계의 신뢰도 순위에 따라 현재 단계의 후보 쿼리 토큰 집합 I_focus를 선정하고, 로컬 윈도우 확장을 통해 I_active를 만든다. 이렇게 하면 불필요한 원거리 토큰에 대한 어텐션 계산을 크게 줄일 수 있다.

둘째, dLLM의 어텐션 맵에서 “어텐션 싱크(attention sink)”라 불리는 특정 토큰이 여러 레이어에 걸쳐 지속적으로 높은 어텐션 가중치를 받는 현상을 발견했다. 시각화 결과, 서로 다른 레이어(예: Layer 9와 Layer 19)에서 싱크 토큰의 인덱스가 거의 일치한다는 점을 확인하였다. 이는 싱크 토큰을 중간 레이어에서 한 번 식별하면 이후 깊은 레이어에서 재사용할 수 있음을 의미한다. 논문은 이를 “sink‑aware pruning” 전략으로 구현한다. 초기 l개의 레이어는 완전 어텐션을 유지하고, 그 이후 레이어에서는 이전 레이어에서 추출한 싱크 인덱스를 고정하여 어텐션 키/밸류를 선택적으로 유지한다. 동시에, 프롬프트 히스토리 중에서도 동적 블록‑와이즈 Top‑K 선택을 통해 가장 관련성 높은 키/밸류만을 남긴다.

이 두 가지 메커니즘을 결합한 Focus‑dLLM은 KV 캐시를 근사적으로 재사용하면서, 쿼리 토큰을 과거 신뢰도 기반으로 제한하고, 어텐션 싱크와 중요한 히스토리 토큰만을 유지한다. 실험에서는 32K 컨텍스트 길이에서 기존 Fast‑dLLM 대비 2.05배, 그리고 전체 풀 어텐션 대비 29배 이상의 속도 향상을 보였으며, BLEU·ROUGE·GPT‑4 평가에서 성능 저하 없이 동일하거나 약간 개선된 결과를 기록했다. 또한, 다양한 모델(LLaDA‑8B, UltraLLaMA‑13B 등)과 데이터셋(GSM‑8K, WikiText‑103)에서 일관된 효율성을 입증하였다.

핵심 기여는 (1) 토큰 신뢰도와 어텐션 싱크의 시간·공간적 일관성을 정량적으로 분석하고, (2) 이를 활용한 훈련‑프리 희소 어텐션 프레임워크를 제안했으며, (3) 실제 대규모 장문 생성 작업에서 손실 없는 가속을 달성했다는 점이다.