Prefilling-dLLM: Diffusion Language Models의 긴 문맥 추론을 위한 예측적 프리필링 (Predictive

원문 발행 2026. 06. 09. 17:06원문 언어 영어AI 한국어 번역arXiv cs.CL (NLP) 원문 보기

요약

Diffusion Language Models(dLLM)의 긴 문맥 처리 시 발생하는 높은 연산 비용을 해결하기 위한 Prefilling-dLLM 프레임워크를 제안합니다. 별도의 학습 없이 KV 캐싱과 희소 프리필링을 통해 연산 복잡도를 획기적으로 낮추고 추론 속도를 대폭 향상시켰습니다.

핵심 포인트

dLLM의 문맥 길이에 따른 제곱 복잡도 문제를 해결
프리필-디코드 분리 프레임워크를 통한 연산 효율화
8K~32K 문맥에서 최대 28배의 속도 향상 달성
Attention Anchor를 통해 'lost-in-the-middle' 현상 제거

Diffusion large language models (dLLMs)는 매 디노이징 (denoising) 단계마다 전체 접두사 (prefix)를 다시 인코딩하며, 이는 문맥 길이 (context length)에 따라 제곱으로 증가하는 재계산을 유발하여 긴 문맥 (long-context) 시나리오에서 비용이 매우 높아지게 만듭니다. 우리는 dLLM을 위한 별도의 학습이 필요 없는 프리필-디코드 분리 (prefill-decode disaggregation) 프레임워크인 Prefilling-dLLM을 제안합니다. 이 프레임워크는 접두사를 N개의 청크 (chunks)로 분할하고, 이들의 KV 표현 (KV representations)을 한 번만 캐싱하며, 디코딩을 위해 청크 내 토큰 희소성 (intra-chunk token sparsity)을 활용하여 가장 관련성이 높은 상위 K개의 청크를 선택합니다. 이를 통해 희소 프리필링 (sparse prefilling)이 밀집 어텐션 (dense attention)보다 뛰어난 성능을 보일 수 있음을 입증하였으며, 단계별 복잡도를 전체 시퀀스 길이에 대한 제곱에서 디코드 길이 (decode length)에 대한 제곱으로 감소시켰습니다. LongBench 및 InfiniteBench에서 Prefilling-dLLM은 dLLM 가속 방법들 중 최첨단 (state-of-the-art) 품질을 달성하였으며, 비연속적으로 캐싱된 청크 KV에 대해 디코딩을 병렬화하는 어텐션 커널 (attention kernel)을 통해 8K~~32K 문맥에서 9.1~~28.0배의 속도 향상을 구현했습니다. 나아가, 각 청크 앞에 추가된 문장 시작 (beginning-of-sequence) 토큰들이 주기적인 어텐션 앵커 (attention anchors) 역할을 하여 'lost-in-the-middle' 현상을 제거함을 보여줍니다. 코드는 https://github.com/menik1126/Prefilling-dLLM 에서 확인할 수 있습니다.

AI 자동 생성 콘텐츠

원문 바로가기

Prefilling-dLLM: Diffusion Language Models의 긴 문맥 추론을 위한 예측적 프리필링 (Predictive

요약

핵심 포인트

댓글