DLLM 추론 효율 혁신: 토큰 집중으로 컴퓨팅 한계 돌파

초록

본 논문은 확산 기반 대형 언어 모델(DLLM)의 디코딩 과정이 대부분 비효율적인 연산으로 가득 차 있다는 점을 지적한다. 저자는 초기 레이어의 어텐션 점수가 토큰의 디코딩 가능성과 강하게 연관된다는 사실을 발견하고, 이를 기반으로 ‘FOCUS’라는 토큰 선택·제거 시스템을 제안한다. FOCUS는 디코딩 가능성이 높은 토큰만을 남겨 연산량을 65‑80% 줄이고, 기존 엔진 대비 최대 3.52배의 처리량 향상을 달성하면서도 생성 품질을 유지한다.

상세 분석

DLLM은 블록 단위로 토큰을 동시에 처리하지만, 각 디퓨전 단계에서 실제로 디코딩되는 토큰은 전체 블록의 약 10%에 불과하다. 이는 전체 블록에 대해 KV 캐시와 어텐션 연산을 수행하므로, 대부분의 FLOP이 ‘무용지물’ 연산에 소비된다는 의미다. 기존 연구들은 KV 캐시 재사용이나 근사 캐시 등을 통해 일부 병목을 완화했지만, 블록 전체를 매 단계마다 재계산하는 구조적 비효율을 근본적으로 해결하지 못했다.

저자들은 먼저 어텐션 스코어를 토큰 중요도 지표로 정의하고, 레이어 0과 레이어 1 사이의 중요도 차이(ΔI)를 계산한다. 실험 결과, ΔI가 높은 토큰은 디코딩 확률이 크게 상승하고, 반대로 낮은 ΔI를 가진 토큰은 거의 디코딩되지 않는다. 이는 초기 레이어에서 토큰이 의미적 ‘정착’ 과정을 겪으며, 의미가 명확해진 토큰이 이후 단계에서 지속적으로 처리될 가능성이 높아진다는 것을 시사한다.

이러한 통찰을 바탕으로 FOCUS는 두 가지 핵심 메커니즘을 구현한다. 첫째, ‘Dynamic Budgeting’은 현재 단계에서 예상되는 디코딩 토큰 수(N̄decoded)와 사전 통계에 기반해 유지할 토큰 수 K를 동적으로 결정한다. 둘째, ‘Delta Calculation’은 레이어 0과 1의 어텐션 점수 차이를 실시간으로 계산해 토큰을 ‘evict’(제거)하거나 유지한다. 이 과정은 추가 학습 없이도 수행 가능하며, 토큰이 제거되면 이후 레이어의 어텐션 및 FFN 연산에서도 해당 토큰을 완전히 배제한다.

실험에서는 SDAR‑8B‑Chat 모델(B=32, confidence=0.9)을 사용해 여러 벤치마크(GSM8K, Math500, HumanEval, MBPP 등)에서 평균 65‑80%의 토큰을 조기에 제거했다. 결과적으로 FLOP이 크게 감소했으며, LMDeploy 대비 최대 3.52배의 처리량 향상을 기록했다. 품질 측면에서는 기존 방법과 동등하거나 약간 개선된 점수를 보였으며, 이는 토큰 선택이 생성 품질을 해치지 않음을 입증한다.

FOCUS는 기존 DLLM 파이프라인에 최소한의 변경만으로 적용 가능하고, 오픈소스로 제공돼 실무 적용이 용이하다. 또한, 어텐션 기반 토큰 중요도와 디코딩 가능성 사이의 강한 상관관계를 발견함으로써, 향후 DLLM 최적화 연구에 새로운 방향을 제시한다.