디퓨전 LLM을 위한 의존성 인식 고속 추론 DAWN

초록

DAWN은 확산 기반 대형 언어 모델(dLLM)의 병렬 디코딩 과정에서 토큰 간 의존성을 고려해, 기존 보수적 기준보다 더 많은 토큰을 동시에 복원하면서도 품질 저하를 최소화한다. 어텐션 맵을 활용해 의존성 그래프를 구축하고, 고신뢰 토큰을 앵커로 삼아 주변 토큰의 신뢰도를 완화시키며, 충돌 관계를 탐지해 비충돌 집합을 선택한다. 실험 결과 1.8배~8.0배의 속도 향상을 달성하면서 BLEU, ROUGE 등 품질 지표를 유지한다.

상세 분석

DAWN은 확산 기반 LLM(dLLM)의 핵심 한계인 “비독립적인 위치 예측” 문제를 해결하기 위해 세 가지 모듈을 파이프라인 형태로 결합한다. 첫 번째 모듈인 Dependency Graph Construction은 각 디노이징 단계에서 모델이 출력하는 어텐션 맵을 평균화하고, “어텐션 싱크”(특정 토큰이 비정상적으로 높은 어텐션을 끌어모으는 현상)를 탐지해 필터링한다. 이후, 정규화된 어텐션 스코어가 사전 정의된 임계값 τ_edge 를 초과하는 경우에만 방향성 엣지를 추가함으로써, 토큰 간의 유의미한 의존 관계만을 포함하는 희소 그래프를 만든다. 이 그래프는 토큰 i 가 토큰 j 로부터 얼마나 영향을 받는지를 정량화한 근사치이며, 이후 스케줄링 단계의 의사결정 근거가 된다.

두 번째 모듈인 Anchor‑Guided Decoding은 고신뢰( confidence ≥ τ_high ) 토큰을 “앵커”로 지정한다. 실험적으로 앵커 토큰에 강하게 의존하는 마스크된 위치는 현재 confidence 가 낮더라도 최종 출력과 높은 일치율을 보인다. 따라서 DAWN은 앵커와 연결된 토큰에 대해 confidence 임계값을 낮추어(τ_anchor) 더 많은 토큰을 동시에 복원한다. 이는 기존 방법이 “높은 confidence → 독립”이라는 단순 가정에 머물러 놓친 부분을 보완한다.

세 번째 모듈인 Conflict‑Based Scheduling은 남은 후보 토큰 집합에서 그래프 상의 충돌 관계(양방향 의존 혹은 강한 단방향 의존)를 탐지한다. 충돌이 없는 최대 독립 집합을 탐욕적으로 선택해(τ_conflict) 추가적인 병렬 복원을 수행한다. 이렇게 두 단계에서 선택된 토큰 집합 U_anchor ∪ U_conflict 은 서로 간에 의미적·통계적 충돌이 최소화된 상태이므로, 한 번에 다수의 토큰을 언마스크해도 품질 저하가 억제된다.

핵심 인사이트는 (1) 어텐션 싱크는 의존성 추정에 잡음이 되므로 사전 필터링이 필요하고, (2) 고신뢰 앵커 토큰은 주변 토큰의 예측을 안정화시켜 confidence 기준을 완화할 수 있다는 점이다. 이 두 가지를 그래프 기반 스케줄링에 통합함으로써, DAWN은 기존의 “confidence‑only” 혹은 “entropy‑only” 기반 선택보다 훨씬 넓은 병렬화 폭을 확보한다. 실험에서는 LLaDA‑8B‑Instruct, LLaMA‑7B 등 다양한 dLLM에 적용했으며, GSM8K, HumanEval, WikiText 등 여러 베치에서 1.80×~8.06×의 속도 향상을 기록하면서 BLEU, ROUGE‑L, Exact Match 등 품질 지표는 0.1% 이내의 손실만 보였다. 또한, 각 모듈을 제거한 Ablation 실험을 통해 그래프 필터링, 앵커 활용, 충돌 회피가 각각 속도·품질 트레이드오프에 기여함을 확인했다.

전반적으로 DAWN은 훈련 없이 기존 dLLM에 바로 적용 가능한 “플러그인” 방식이며, 어텐션 기반 의존성 추정과 그래프 스케줄링을 결합한 새로운 병렬 디코딩 패러다임을 제시한다.