장문 맥락 확산 언어 모델 추론 가속화

마스크 확산 언어 모델(MDLM)은 모든 디노이징 단계에서 양방향 어텐션을 사용하며, 이는 단계 수 T와 시퀀스 길이 n에 대해 O(Tn^2)의 높은 연산 비용을 발생시킵니다. 특히, 80억 매개변수 모델에서 n=8192, T=128일 경우 KV 캐시만 40GB를 초과하여 단일 GPU로는 긴 문서 생성이 어렵습니다. 본 연구는 이러한 한계를 극복하기 위해 사전 학습된 MDLM의 추론 래퍼인 SW-SpeedDLM을 제안합니다. SW-SpeedDLM은 세 가지 핵심 구성 요소로 이루어져 있습니다. 첫째, 분할 슬라이딩 윈도우 디노이징(SSWD)은 각 디노이징 루프를 W 토큰의 윈도우로 제한하여 단계당 비용을 O(W^2n/S)로 줄입니다. 둘째, 교차 세그먼트 KV 압축(CSKV)은 완료된 각 윈도우를 C개의 요약 토큰으로 인코딩하여 이후 윈도우가 O(C) 비용으로 참조하도록 합니다. 셋째, 윈도우 수준 추측 수용(WLSA)은 작은 드래프트 모델이 k개의 디노이징 단계를 제안하고, 타겟 모델이 이를 단일 순방향 패스로 검증하여 윈도우당 최대 2.4배의 속도 향상을 제공합니다. WLSA는 타겟 모델의 정확한 주변 분포를 보존함을 증명했습니다. MDLM-860M 및 LLADA-8B 모델을 PG-19, LongBench, WritingPrompts 데이터셋에 적용한 결과, SW-SpeedDLM은 n=8192에서 기존 풀 어텐션 방식(최대 n=2048) 대비 3.7배 높은 처리량을 달성했습니다. 또한, 풀 어텐션이 GPU 메모리 한계에 도달하는 n=4096 대비 피크 메모리를 2.5배 절감했습니다. 이러한 개선에도 불구하고 PG-19의 문자당 비트(bpc)는 0.18만 증가하여 품질 저하를 최소화했습니다. SW-SpeedDLM은 단일 A100 GPU(40GB)에서 최대 16,384 토큰의 시퀀스 생성을 가능하게 하여, 장문 맥락 확산 언어 모델의 실용적인 적용 가능성을 크게 확장합니다.