💬 NLP

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

Het artikel introduceert Prefilling-dLLM, een training-vrij framework dat de lange-context inferentie in diffusion language models versnelt door relevante prefix-chunks te cachen en er spaars uit te selecteren om de computationele complexiteit te verminderen van kwadratisch ten opzichte van de volledige sequentielengte naar kwadratisch ten opzichte van de decode-lengte, waardoor het state-of-the-art versnellingen bereikt en het lost-in-the-middle fenomeen mitigeert.

Oorspronkelijke auteurs: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Gepubliceerd 2026-06-10

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorme legpuzzel probeert op te lossen, maar in plaats van naar de hele afbeelding in één keer te kijken, moet je de gehele doos met puzzelstukjes opnieuw onderzoeken elke keer dat je één nieuw stukje op tafel legt.

Dat is in essentie hoe Diffusion Large Language Models (dLLMs) momenteel werken wanneer ze lange verhalen of documenten verwerken. Elke keer dat het model een nieuw woord probeert te genereren, leest en verwerkt het de volledige geschiedenis van het gesprek opnieuw vanaf het begin. Naarmate het verhaal langer wordt, wordt dit "opnieuw lezen" zo traag en kostbaar dat het voelt alsof je een marathon probeert te lopen terwijl je een zware rugzak draagt die met elke stap zwaarder wordt.

Het artikel introduceert een nieuwe methode genaamd Prefilling-dLLM om dit op te lossen. Hier is hoe het werkt, met behulp van enkele alledaagse analogieën:

1. Het Probleen: De "Opnieuw Lees"-valstrik

In traditionele AI-modellen onthoud je een pagina uit een boek zodra je deze hebt gelezen en focus je je alleen op de nieuwe pagina die je nu leest. Maar bij deze diffusiemodellen vergeet de computer het boek en leest hij het hele boek opnieuw, van pagina 1 tot pagina 1.000, elke keer dat hij één nieuw woord schrijft.

Het resultaat: Als het boek kort is, is het prima. Als het boek 32.000 pagina's lang is, besteedt de computer 99% van zijn tijd aan het opnieuw lezen van de oude pagina's en besteedt hij slechts 1% aan het daadwerkelijk schrijven van nieuwe woorden.

2. De Oplossing: Het "Slimme Bibliothecaris"-systeem

De auteurs stellen een systeem voor genaamd Prefilling-dLLM dat werkt als een super-efficiënte bibliothecaris. In plaats van de hele bibliotheek telkens opnieuw te lezen, doet de bibliothecaris twee dingen:

Stap A: De "Prefill" (De bibliotheek organiseren)

Voordat het schrijven begint, neemt de bibliothecaris de enorme invoertekst (de "prefix") en hakt deze op in stukken (zoals hoofdstukken in een boek).

De truc: De bibliothecaris leest elk hoofdstuk één keer, maakt een "samenvattingskaart" (een KV-cache genoemd) voor dat hoofdstuk en zet deze op een plank.
De innovatie: De bibliothecaris leest niet elk afzonderlijk woord in elk hoofdstuk. Ze gebruiken een speciale truc om alleen de belangrijkste zinnen uit elk hoofdstuk te bewaren en de overbodige informatie weg te gooien. Dit maakt de samenvattingskaarten veel kleiner en sneller te verwerken.

Stap B: De "Decode" (Het verhaal schrijven)

Wanneer het model nu een volgend woord moet schrijven, leest het niet de hele bibliotheek opnieuw.

De selectie: Het model vraelt: "Welke hoofdstukken zijn op dit moment echt relevant voor wat ik aan het schrijven ben?" Het gebruikt een slim scoresysteem om alleen de top paar hoofdstukken (de meest relevante stukken) van de plank te kiezen.
De efficiëntie: Het negeert de andere 90% van de bibliotheek die niet nodig is voor deze specifieke zin. Het kijkt alleen naar de relevante "samenvattingskaarten" die eerder zijn gemaakt.

3. Waarom dit een Game-Changer is

Het artikel beweert dat deze aanpak een enorme versnelling is omdat:

Geen meer opnieuw lezen: Het zware werk van het lezen van de lange tekst gebeurt één keer aan het begin.
Slim overslaan: Tijdens de eigenlijke schrijffase kijkt het model slechts naar een fractie van de tekst (de meest relevante stukken), in plaats van naar het geheel.
De snelheid: Bij lange contexten (8.000 tot 32.000 woorden) is deze methode 9 tot 28 keer sneller dan eerdere methoden, terwijl de antwoorden nog steeds correct blijven.

4. Het oplossen van het "Lost in the Middle"-mysterie

Er is een bekend probleem bij AI-modellen met lange contexten genaamd "Lost in the Middle". Stel je een persoon voor die een lange lijst met feiten leest; die persoon onthoudt de eerste paar en de laatste paar perfect, maar vergeet de zaken in het midden volledig.

De oplossing uit het artikel: De auteurs ontdekten dat door de tekst in stukken te verdelen en een speciale "anker"-token (zoals een hoofdstuktitel) aan het begin van elk stuk toe te voegen, het model informatie overal in de tekst kan vinden, zelfs in het midden. Het is also�f er een inhoudsopgave is die perfect werkt, zodat het model nooit de weg kwijtraakt, ongeacht hoe lang het verhaal ook is.

5. De Balans van de "Chunk Size"

Het artikel heeft ook getest hoe groot deze "hoofdstukken" moeten zijn.

Te groot: Als de stukken enorm groot zijn, kan het model belangrijke details in het midden van de tekst missen.
Te klein: Als de stukken heel klein zijn, besteedt de computer te veel tijd aan het beheren van de lijst met stukken.
Het ideale punt: Ze ontdekten dat voor zeer lange teksten het gebruik van veel kleinere stukken het beste werkt om de nauwkeurigheid hoog te houden.

Samenvatting

Prefilling-dLLM is als de overstap van een student die een tekstboek van 500 pagina's telkens opnieuw leest voordat hij een vraag beantwoordt, naar een student die één keer een slimme, beknopte studiegids maakt en vervolgens alleen naar de specifieke pagina's bladert die hij nodig heeft bij het beantwoorden van de vragen.

Het resultaat? De AI kan veel langere gesprekken en documenten aan zonder dat het proces tot stilstand komt, en het vergeet de belangrijke details die verborgen zitten in het midden van de tekst niet.

Technische Samenvatting: Prefilling-dLLM

Probleemstelling

Diffusion Large Language Models (dLLMs) bieden een veelbelovend alternatief voor autoregressieve (AR) modellen door parallelle token-generatie mogelijk te maken via iteratieve denoising. Echter, dLLMs kampen met een kritieke inefficiëntie in scenario's met lange context: bij elke denoising-stap moet het model de gehele input-prefix opnieuw encoderen om self-attention te berekenen. Dit resulteert in een computationele complexiteit die kwadratisch schaalt met de totale sequentielengte ( $O((L_p + L_d)^2 \cdot T)$ , waarbij $L_p$ de prefix-lengte is, $L_d$ de decode-lengte en $T$ het aantal denoising-stappen). In tegen tegenstelling tot AR-modellen, waarbij de prefill-fase compute-bound is en de decode-fase memory-bound, blijft dLLM-inferentie gedurende alle iteraties compute-bound omdat de volledige sequentie (prefix + respons) bij elke stap gezamenlijk wordt verwerkt. Deze herhaalde herberekening van de prefix wordt problematisch voor toepassingen met een lange context (bijv. 8K–32K+ tokens).

Methodologie: Prefilling-dLLM

De auteurs stellen PREFILLING-DLLM voor, een training-vrij, prefill-decode disaggregatie-framework dat ontworpen is om prefix-herberekening te elimineren en sparsity te introduceren. De methode werkt in drie stadia:

1. Chunked Prefill en KV Caching

In plaats van de gehele prefix als één enkel blok te verwerken, wordt de input-prefix opgedeeld in $N$ niet-overlappende chunks van een vaste grootte $C$ .

Onafhankelijke Verwerking: Elke chunk wordt onafhankelijk verwerkt met intra-chunk bidirectionele attention. Een speciale Beginning-of-Sequence (BOS) token wordt aan elke chunk toegevoegd om als delimiter te fungeren.
Predictieve Scoring: Om te bepalen welke chunks relevant zijn zonder grondwaarheid-targets, genereert het framework pseudo-labels ( $m$ $m$ ) via een voorlopige diffusion-pass. Het scoort vervolgens elke chunk met behulp van twee signalen:
- Self-Information Score ( $s_I$ ): De negatieve log-likelihood van de query-window geconditioneerd op de chunk.
- Pseudo-Label Score ( $s_P$ ): De likelihood van de pseudo-labels geconditioneerd op de chunk en de query.
Selectie: De top- $K$ meest relevante chunks worden geselecteerd op basis van de gecombineerde score ( $s = s_I + s_P$ ). De KV-representaties van deze geselecteerde chunks worden gecached.

2. Two-Level Sparsity

Het framework introduceert sparsity op zowel chunk- als tokenniveau om de complexiteit te verminderen:

Inter-chunk Sparsity: Alleen de top- $K$ chunks nemen deel aan het attention-mechanisme tijdens het decoderen, in plaats van de volledige set van $N$ chunks.
Intra-chunk Sparsity: Binnen elke geselecteerde chunk worden tokens gerangschikt op basis van hun cumulatieve bidirectionele attention-gewicht met de query-tokens. Alleen de top- $B$ meest relevante tokens per chunk worden behouden in de KV-cache, waarbij irrelevante tokens worden verwijderd.

3. Decoding met Statische KV-Hergebruik

Tijdens het iteratieve denoising-proces:

De gecachte KV-representaties van de geselecteerde $K$ chunks blijven statisch en worden hergebruikt over alle $T$ denoising-stappen.
Het model berekent alleen de KV-representaties voor de query en de evoluerende respons-tokens opnieuw.
Dit vermindert de per-stap complexiteit van $O((L_p + L_d)^2 \cdot T)$ naar $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T$ , wat de kosten van het decoderen effectief ontkoppelt van de volledige prefix-lengte.

Attention Kernel Optimalisatie

Om de geheugenbottlenecks aan te pakken die worden veroorzaakt door het laden van niet-contigue chunk KV-caches, implementeren de auteurs een Split-S FlexAttention kernel. Deze kernel paralleliseert het decoderen over de niet-contigue gecachte chunk KV's door de KV-dimensie te partitioneren en chunks onafhankelijk te verwerken voordat de resultaten worden samengevoegd via log-sum-exp reductie, waardoor het verzamelen van data in contigue geheugens niet nodig is.

Belangrijkste Bevindingen en Resultaten

Prestaties en Efficiëntie

Versnelling: Op LongBench en InfiniteBench bereikt PREFILLING-DLLM een versnelling van 9,1× tot 28,0× bij contextlengtes van 8K tot 32K vergeleken met standaard Transformer-baselines.
Kwaliteit: De methode bereikt een state-of-the-art kwaliteit onder de dLLM-acceleratiemethoden. Op LongBench behaalt het een gemiddelde score van 34,59 (Dream-7B) en 37,02 (UltraLLaDA), waarmee het baselines zoals Fast-dLLM, Sparse-dLLM en dKV-Cache overtreft.
Long-Context Schaling: Op InfiniteBench (contexten >128K) behaalt de methode een gemiddelde nauwkeurigheid van 43,62% op Dream-7B, wat de sterkste baseline (Fast-dLLM v2 met 30,32%) aanzienlijk overtreft. Opvallend genoeg behaalt het 95,42% nauwkeurigheid op de Passkey retrieval taak en 70,00% op de Number retrieval taak.

Aanpak van "Lost-in-the-Middle"

Fenomeen: Het paper observeert dat dLLMs een "lost-in-the-middle" fenomeen vertonen onder context-extrapolatie, waarbij de nauwkeurigheid afneemt voor informatie die zich in het midden van de context bevindt, wat een monotone decay-bias naar de respons vertoont in plaats van de U-vormige curve die bij AR-modellen wordt gezien.
Mitigatie: Door chunks te selecteren op basis van predictieve relevantie in plaats van positionele nabijheid, elimineert PREFILLING-DLLM deze bias. De methode behoudt een hoge exact-match nauwkeurigheid over alle posities in de context, zelfs bij lengtes van 32K en 128K.

Mechanisme Inzichten

Attention Anchors: De aan elke chunk toegevoegde BOS-tokens fungeren als periodieke attention anchors. Analyse toont aan dat deze tokens de attention-massa uniform over de context verdelen (het vormen van stabiele, laag-magnitude ridges), in plaats van te degenereren naar "attention sinks" die een onevenredige hoeveelheid attention-massa absorberen. Dit mechanisme maakt positie-invariante retrieval mogelijk.
Chunk Grootte Trade-off: Er is een trade-off tussen chunk-grootte en retrieval-nauwkeurigheid. Kleinere chunks (bijv. 256 tokens) behouden een hoge nauwkeurigheid zelfs bij 128K contexten, terwijl grotere chunks (bijv. 4096 tokens) scherp degraderen voorbij 32K. Multi-chunk selectie is essentieel voor long-context inferentie.

Betekenis en Claims

Het paper claimt dat PREFILLING-dLLM de kloof overbrugt tussen de compute-bound natuur van dLLM-inferentie en de efficiëntie van prefill-decode disaggregatie-architecturen die worden gebruikt in AR-modellen. De primaire bijdragen zijn:

Training-vrije Disaggregatie: Het introduceert een framework dat prefix KV-representaties één keer cachet en hergebruikt, waardoor de kwadratische herberekeningskosten van de prefix bij elke denoising-stap worden geëlimineerd zonder dat hertraining van het model vereist is.
Sparse Prefilling: Het demonstreert dat sparse prefilling (het selecteren van relevante chunks en tokens) de dense attention in dLLMs kan overtreffen, wat de complexiteit vermindert terwijl de generatiekwaliteit gelijk blijft of verbetert.
Systeem-niveau Optimalisatie: Het biedt een gespecialiseerde attention kernel (Split-S FlexAttention) die de chunked structuur van de prefix gebruikt om aanzienlijke latentiereducties te bereiken (tot 10,2× over vanilla FlexAttention) tijdens de decode-fase.
Robuustheid bij Lange Contexten: Het laat zien dat deze aanpak het "lost-in-the-middle" fenomeen mitigeert en dLLMs in staat stelt effectief te schalen naar 128K+ tokens met een hoge retrieval-nauwkeurigheid, een capaciteit die voorheen onverkend was bij dLLM-acceleratie.

De auteurs merken beperkingen op, waaronder de statische aard van chunk-selectie (geen dynamische herselectie tijdens het decoderen), de noodzaak van taakspecifieke afstemming van chunk-grootte en $K$ , en het feit dat de huidige evaluatie beperkt is tot Engelse benchmarks op specifieke modelarchitecturen (Dream-7B en UltraLLaDA).