💬 NLP

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

이 논문은 관련 있는 접두사 청크(prefix chunks)를 캐싱하고 희소하게 선택함으로써 계산 복잡도를 전체 시퀀스 길이에 대한 이차 복잡도에서 디코드 길이에 대한 이차 복잡도로 줄여, 디퓨전 언어 모델의 긴 문맥 추론을 가속화하고 'lost-in-the-middle' 현상을 완화하는 학습이 필요 없는 프레임워크인 Prefilling-dLLM을 소개한다.

원저자: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

게시일 2026-06-10

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 직소 퍼즐을 맞추려 한다고 상상해 보세요. 하지만 전체 그림을 한 번에 보는 대신, 조각을 단 하나 놓을 때마다 매번 퍼즐 상자 전체를 다시 확인해야 합니다.

이것이 현재 디퓨전 대규모 언어 모델(dLLM)이 긴 이야기나 문서를 처리할 때 작동하는 방식의 본질입니다. 모델이 새로운 단어를 생성하려고 할 시도할 때마다, 대화의 처음부터 끝까지 모든 히스토리를 다시 읽고 재처리합니다. 이야기가 길어질수록, 이 "다시 읽기" 과정은 너무 느리고 비용이 많이 들어서, 마치 발걸음을 옮딜 때마다 점점 더 무거워지는 배낭을 메고 마라톤을 하는 것과 같습니다.

이 논문은 이를 해결하기 위해 Prefilling-dLLM이라는 새로운 방법을 소개합니다. 일상적인 비유를 통해 그 작동 원리를 설명하겠습니다.

1. 문제점: "다시 읽기"의 함정

전통적인 AI 모델의 경우, 책의 한 페이지를 읽고 나면 그것을 기억하고 다음에 읽을 새 페이지에만 집중합니다. 하지만 이러한 디퓨전 모델에서는 컴퓨터가 책을 잊어버리고, 단 하나의 새로운 단어를 쓸 때마다 1페이지부터 1,000페이지까지 책 전체를 다시 읽습니다.

결과: 책이 짧다면 괜찮습니다. 하지만 책이 32,000페이지라면, 컴퓨터는 새로운 글을 쓰는 데 1%의 시간만 쓰고 나머지 99%의 시간을 과거의 페이지들을 다시 읽는 데 소비하게 됩니다.

2. 해결책: "스마트한 사서" 시스템

저자들은 Prefilling-dLLM이라는, 매우 효율적인 사서처럼 작동하는 시스템을 제안합니다. 매번 도서관 전체를 다시 읽는 대신, 이 사서는 두 가지 일을 수행합니다.

단계 A: "프리필링(Prefilling)" (도서관 정리하기)

글쓰기가 시작되기 전, 사서는 방대한 입력 텍스트(접두사, prefix)를 가져와서 덩어리(chunk) 단위로 나눕니다(마치 책의 장(chapter)처럼).

비법: 사서는 각 장을 단 한 번 읽고, 그에 대한 "요약 카드"(KV 캐시라고 불림)를 만들어 선반에 올려둡니다.
혁신: 사서는 각 장의 모든 단어를 다 읽지 않습니다. 각 장에서 가장 중요한 문장들만 남기는 특별한 기술을 사용하여 불필요한 부분(fluff)을 버립니다. 이를 통해 요약 카드를 훨씬 작고 다루기 쉽게 만듭니다.

단계 B: "디코딩(Decoding)" (이야기 쓰기)

이제 모델이 다음 단어를 써야 할 때, 도서관 전체를 다시 읽지 않습니다.

선택: 모델은 "지금 내가 쓰고 있는 내용과 실제로 관련이 있는 장(chapter)은 무엇인가?"라고 묻습니다. 모델은 스마트한 점수 산정 시스템을 사용하여 선반 위에 있는 상위 몇 개의 장(가장 관련성이 높은 덩어리들)만을 골라냅니다.
효율성: 모델은 현재 문장에 필요하지 않은 나머지 90%의 도서관 내용은 무시합니다. 오직 이전에 만들어둔 관련 있는 "요약 카드"들만 살펴봅니다.

3. 이것이 왜 게임 체인저인가

이 논문은 이 접근 방식이 엄청난 속도 향상을 가져온다고 주장합니다. 그 이유는 다음과 같습니다:

더 이상의 다시 읽기 없음: 긴 텍스트를 읽는 힘든 작업은 처음에 단 한 번 이루어집니다.
스마트한 건너뛰기: 실제 글쓰기 단계에서 모델은 전체 텍스트가 아닌, 아주 적은 부분(상위 관련 덩어리들)만을 살펴봅니다.
속도: 긴 문맥(8,000~32,000 단어)에서 이 방법은 기존 방식보다 9배에서 28배 더 빠르며, 정확도 또한 유지합니다.

4. "중간에서 길을 잃는(Lost in the Middle)" 미스터리 해결

긴 문맥을 가진 AI 모델에는 "Lost in the Middle"이라는 알려진 문제가 있습니다. 긴 사실 목록을 읽는 사람을 상상해 보세요. 처음 몇 개와 마지막 몇 개는 완벽하게 기억하지만, 중간에 있는 내용은 완전히 잊어버립니다.

논문의 해결책: 저자들은 텍스트를 덩어리로 나누고 각 덩어리의 시작 부분에 특별한 "앵커(anchor)" 토큰(마치 책의 장 제목과 같은 역할)을 추가함으로써, 모델이 텍스트 어디에서든 정보를 찾을 수 있다는 것을 발견했습니다. 이는 완벽하게 작동하는 목차를 갖는 것과 같아서, 이야기가 아무리 길어져도 모델이 길을 잃지 않게 해줍니다.

5. "덩어리 크기(Chunk Size)"의 균형

논문은 또한 이 "장(chapter)"의 크기가 어느 정도여야 하는지도 테스트했습니다.

너무 클 경우: 덩어리가 너무 크면 모델이 텍스트 중간의 중요한 세부 사항을 놓칠 수 있습니다.
너무 작을 경우: 덩어리가 너무 작으면 컴퓨터가 덩어리 목록을 관리하는 데 너무 많은 시간을 소비하게 됩니다.
최적의 지점: 저자들은 매우 긴 텍스트의 경우, 정확도를 높게 유지하기 위해 더 많은 작은 덩어리를 사용하는 것이 가장 효과적이라는 것을 발견했습니다.

요요약

Prefilling-dLLM은 퀴즈 문제 하나를 풀 때마다 500페이지짜리 교과서를 다시 읽는 학생에서, 처음에 스마트하고 압축된 학습 가이드를 한 번 만든 뒤, 답을 할 때 필요한 특정 페이지만 펼쳐보는 학생으로 변하는 것과 같습니다.

그 결과는 어떨까요? AI는 속도가 느려지지 않으면서도 훨씬 더 긴 대화와 문서를 처리할 수 있으며, 텍스트 중간에 숨겨진 중요한 세부 사항도 잊어버리지 않습니다.

기술 요약: Prefilling-dLLM

문제 정의

확산 대규모 언어 모델(Diffusion Large Language Models, dLLMs)은 반복적인 디노이징(denoising)을 통해 병렬 토큰 생성을 가능하게 함으로써 자기회귀(autoregressive, AR) 모델의 유망한 대안을 제시합니다. 그러나 dLLM은 긴 문맥 시나리오에서 심각한 비효율성에 직면합니다. 매 디노이징 단계마다 셀프 어텐션(self-attention)을 계산하기 위해 전체 입력 접두사(prefix)를 다시 인코딩해야 하기 때문입니다. 이는 계산 복잡도가 전체 시퀀스 길이에 대해 제곱으로 스케일링되는 결과( $O((L_p + L_d)^2 \cdot T)$ , 여기서 $L_p$ 는 접두사 길이, $L_d$ 는 디코드 길이, $T$ 는 디노이징 단계)를 초래합니다. AR 모델의 경우 프리필(prefill) 단계는 연산 집약적(compute-bound)이고 디코드 단계는 메모리 집약적(memory-bound)인 것과 달리, dLLM 추론은 전체 시퀀스(접두사 + 응답)가 모든 단계에서 공동으로 처리되기 때문에 모든 반복 과정 동안 연산 집약적인 상태를 유지합니다. 이러한 반복적인 접두사 재계산은 긴 문맥 애플리케이션(예: 8K–32K+ 토큰)에서 매우 부담스러운 작업이 됩니다.

방법론: Prefilling-dLLM

저자들은 접두사 재계산을 제거하고 희소성(sparsity)을 도입하도록 설계된 학습이 필요 없는(training-free) 프리필-디코드 분리(prefill-decode disaggregation) 프레임워크인 PREFILLING-DLLM을 제안합니다. 이 방법은 세 가지 단계로 작동합니다.

1. 청크 단위 프리필 및 KV 캐싱

입력 접두사를 단일 블록으로 처리하는 대신, 입력 접두사를 고정된 크기 $C$ 를 가진 $N$ 개의 중첩되지 않는 청크(chunk)로 분할합니다.

독립적 처리: 각 청크는 청크 내부의 양방향 어텐션을 사용하여 독립적으로 처리됩니다. 구분자 역할을 하기 위해 각 청크 앞에 특수한 문장 시작(Beginning-of-Sequence, BOS) 토큰이 추가됩니다.
예측적 스코어링: 정답 타겟 없이도 어떤 청크가 관련이 있는지 결정하기 위해, 프레임워크는 예비 확산 패스(preliminary diffusion pass)를 통해 의사 라벨(pseudo-labels, $m$ $m$ )을 생성합니다. 그 후 두 가지 신호를 사용하여 각 청키의 점수를 매깁니다:
- 자기 정보량 점수 (Self-Information Score, $s_I$ ): 청크에 조건화된 쿼리 윈도우의 음의 로그 가능도(negative log-likelihood).
- 의사 라벨 점수 (Pseudo-Label Score, $s_P$ ): 청크 및 쿼리에 조건화된 의사 라벨의 가능도.
선택: 결합 점수( $s = s_I + s_P$ )를 기준으로 가장 관련성이 높은 상위- $K$ 개 청크를 선택합니다. 선택된 청크들의 KV 표현(representation)이 캐싱됩니다.

2. 2단계 희소성 (Two-Level Sparsity)

프레임워크는 복잡도를 줄이기 위해 청크 레벨과 토큰 레벨 모두에서 희소성을 도입합니다.

청크 간 희소성 (Inter-chunk Sparsity): 디코딩 중에 전체 $N$ 개 청크가 아닌 상위- $K$ 개 청크만이 어텐션 메커니즘에 참여합니다.
청크 내 희소성 (Intra-chunk Sparsity): 선택된 각 청크 내에서 토큰들은 쿼리 토큰과의 누적 양방향 어텐션 가중치에 따라 순위가 매겨집니다. 각 청크에서 가장 관련성이 높은 상위- $B$ 개 토큰만 KV 캐시에 유지하고 나머지 불필요한 토큰은 제거합니다.

3. 정적 KV 재사용을 통한 디코딩

반복적인 디노이징 과정 동안:

선택된 $K$ 개 청크의 캐싱된 KV 표현은 정적(static) 상태를 유지하며 모든 $T$ 디노이징 단계에서 재사용됩니다.
모델은 쿼리 및 진화하는 응답 토큰에 대한 KV 표현만을 다시 계산합니다.
이를 통해 단계당 복잡도를 $O((L_p + L_d)^2 \cdot T)$ 에서 $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ 로 줄여, 디코딩 비용을 전체 접두사 길이로부터 효과적으로 분리합니다.

어텐션 커널 최적화

비연속적으로 캐싱된 청크 KV 캐시를 로드할 때 발생하는 메모리 병목 현상을 해결하기 위해, 저자들은 Split-S FlexAttention 커널을 구현했습니다. 이 커널은 KV 차원을 분할하고 결과를 병합하기 전에 청크를 독립적으로 처리한 뒤 log-sum-exp 리덕션을 통해 결과를 합병함으로써, 데이터를 연속적인 메모리로 모으는 과정 없이 비연속적으로 캐싱된 청크 KV들에 대해 디코딩을 병렬화합니다.

주요 결과 및 발견

성능 및 효율성

속도 향상: LongBench와 InfiniteBench에서 PREFILLING-DLLM은 8K에서 32K의 문맥 길이에서 표준 Transformer 베이스라인 대비 9.1배에서 28.0배의 속도 향상을 달야냅니다.
품질: 이 방법은 dLLM 가속 방법들 중 최고 수준의 품질을 달성합니다. LongBench에서 Dream-7B 기준 평균 34.59, UltraLLaDA 기준 37.02를 기록하여 Fast-dLLM, Sparse-dLLM, dKV-Cache와 같은 베이스라인을 능가했습니다.
긴 문맥 스케일링: InfiniteBench(128K 이상의 문맥)에서 이 방법은 Dream-7B 기준으로 **43.62%**의 평균 정확도를 달성하여, 가장 강력한 베이스라인인 Fast-dLLM v2(30.32%)를 크게 앞질렀습니다. 특히, Passkey retrieval 작업에서 95.42%, Number retrieval에서 **70.00%**의 정확도를 기록했습니다.

"Lost-in-the-Middle" 현상 해결

현상: 논문은 dLLM이 문맥 외삽(context extrapolation) 상황에서 "lost-in-the-middle" 현상을 보인다는 점을 관찰했습니다. 즉, 문맥 중간에 위치한 정보에 대한 정확도가 저하되는데, 이는 AR 모델에서 보이는 U자형 곡선이 아니라 응답 쪽으로 치우친 단조 감소 편향(monotonic decay bias)을 보입니다.
완화: PREFILLING-DLLM은 위치적 근접성이 아닌 예측적 관련성에 기반하여 청크를 선택함으로써 이 편향을 제거합니다. 이 방법은 32K 및 128K 길이에서도 문맥의 모든 위치에서 높은 일치 정확도를 유지합니다.

메커니즘 통찰

어텐션 앵커 (Attention Anchors): 각 청크 앞에 추가된 BOS 토큰은 주기적인 어텐션 앵커 역할을 합니다. 분석 결과, 이 토큰들은 어텐션 질량을 불균형하게 흡수하는 "어텐션 싱크(attention sinks)"로 퇴화하는 대신, 문맥 전체에 어텐션 질량을 균일하게 분산시키는(안정적이고 낮은 강도의 능선 형성) 역할을 합니다. 이 메커니즘은 위치 불변적(position-invariant) 검색을 가능하게 합니다.
청크 크기 트레이드오프: 청크 크기와 검색 정확도 사이에는 트레이드오프가 존재합니다. 작은 청크(예: 256 토큰)는 128K 문맥에서도 높은 정확도를 유지하지만, 큰 청크(예: 4096 토큰)는 32K를 넘어서면 급격히 저하됩니다. 긴 문맥 추론을 위해서는 멀티 청크 선택이 필수적입니다.

의의 및 주장

본 논문은 PREFILLING-dLLM이 dLLM 추론의 연산 집약적 특성과 AR 모델에서 사용되는 프리필-디코드 분리 아키텍처의 효율성 사이의 간극을 메운다고 주장합니다. 주요 기여는 다음과 같습니다:

학습이 필요 없는 분리 (Training-Free Disaggregation): 접두사 KV 표현을 한 번 캐싱하고 재사용하여, 모델 재학습 없이 매 디노이징 단계마다 발생하는 이차적인 접두사 재계산 비용을 제거하는 프레임워크를 도입했습니다.
희소 프리필 (Sparse Prefilling): 희소 프리필(관련 청크 및 토큰 선택)이 밀집 어텐션(dense attention)보다 우수할 수 있음을 입증하여, 생성 품질을 유지하거나 개선하면서 복잡도를 줄였습니다.
시스템 레벨 최적화: 접두사의 청크 구조를 활용하여 디코드 단계에서 상당한 레이턴시 감소(vanilla FlexAttention 대비 최대 10.2배)를 달성하는 특화된 어텐션 커널(Split-S FlexAttention)을 제공합니다.
긴 문맥에 대한 견고성: 이 접근 방식이 "lost-in-the-middle" 현상을 완화하고, dLLM이 높은 검색 정확도로 128K 이상의 토큰까지 효과적으로 확장할 수 있도록 함을 보여주며, 이는 기존에 탐구되지 않았던 dLLM 가속 영역입니다.

저자들은 한계점으로 정적 청크 선택 방식(디코딩 중 동적 재선택 불가), 청크 크기 및 $K$ 값에 대한 태스크별 튜닝 필요성, 그리고 현재 평가가 특정 아키텍처(Dream-7B 및 UltraLLaDA)의 영어 벤치마크에 국한되어 있다는 점을 언급했습니다.