🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

Dit artikel introduceert Learned Relay Representations (Relay), een methode die Masked Diffusion Models in staat stelt om latente informatie via een differentieerbare kanaal per token over denoisingsstappen te propageren, waardoor de prestaties op plannings- en coderingstaken worden verbeterd en de inferentielatentie aanzienlijk wordt verlaagd.

Oorspronkelijke auteurs: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Gepubliceerd 2026-05-25

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

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

Het Probleem: De "Harde Reset" Amnesie

Stel je voor dat je een complex puzzel probeert op te lossen, zoals een Sudoku of het schrijven van een stuk code. Je gebruikt een zeer slimme assistent (een AI-model) om je te helpen.

Op dit moment werkt deze assistent als een persoon met amnesie die alles vergeet op het moment dat ze een enkele zin hebben voltooid.

De assistent kijkt naar een blanco pagina vol vraagtekens.
Het denkt hard na, doet wat complexe mentale berekeningen en besluit om één of een paar letters in te vullen.
De Harde Reset: Zodra het die letters heeft opgeschreven, gooit het direct alle complexe gedachten, berekeningen en het "gevoel" weg dat het had over de rest van de pagina.
Voor de volgende stap begint het helemaal opnieuw, kijkt het alleen naar de letters die het net heeft geschreven en vergeet het volledig het rijke interne werk dat het zojuist heeft verricht.

Dit wordt het "Harde Reset"-probleem genoemd. Het paper stelt dat dit inefficiënt is. Het is alsof een kok een soep proeft, "zout nodig" opschrijft en vervolgens direct de smaak van de bouillon vergeet voordat het zout wordt toegevoegd. Ze moeten elke keer dat ze een ingrediënt toevoegen de hele soep opnieuw van scratch proeven.

De Oplossing: De "Relay" Slang

De auteurs stellen een nieuwe methode voor die Relay heet.

Stel je een estafetteloop voor. In plaats dat de loper stopt, de wedstrijd vergeet en opnieuw begint, geeft hij een slang door aan de volgende loper. Deze slang draagt de momentum, de strategie en de vermoeidheid van de vorige loper, zodat de volgende precies kan oppakken waar de vorige is gebleven.

In de AI-wereld is de "slang" een continue stroom van informatie (een verborgen toestand) die het model meeneemt.

Vóór Relay: Het model berekent, schrijft een token en verwijdert de berekening.
Met Relay: Het model berekent, schrijft een token, en geeft een "notitie" door aan zichzelf voor de volgende stap. Deze notitie zegt: "Hé, ik dacht aan dit specifieke patroon, en ik ben 80% zeker van dit volgende deel."

Hoe Het Werkt: Leren de Slang Doorgeven

Het paper introduceert een paar belangrijke trucs om dit werkend te maken:

De "Vooruitdenkende" Training:
Meestal worden AI-modellen getraind om alleen het huidige antwoord goed te krijgen. De auteurs hebben het model getraind om "vooruitdenkend" te zijn. Ze leerden het model: "Schrijf niet alleen vandaag het juiste letter, maar schrijf ook een notitie (de relay) die je zal helpen morgen het juiste letter te schrijven."
Afgeknotte Backpropagation (BPTT):
Dit is een chique wiskundige term voor "vooruitkijken om te leren". Stel je voor dat je een dansroutine oefent. In plaats van alleen maar één beweging te oefenen, oefen je een reeks van drie bewegingen. Als je struikelt op de derde beweging, realiseer je je: "Oh, ik heb de eerste beweging verpest omdat ik me niet had voorbereid op de derde."
Het paper gebruikt een methode genaamd afgeknotte BPTT om het model tijdens de training een paar stappen vooruit te laten kijken. Het leert informatie specifiek door te geven om toekomstige stappen makkelijker te maken.
Het "Zachte" Kanaal:
De informatie die wordt doorgegeven is niet zomaar een woord of een letter; het is een "zacht" getal (een continue waarde). Denk hierbij aan een dimmer in plaats van een aan/uit-lamp. Het stelt het model in staat subtiele hints en kansen mee te nemen die nog niet zijn omgezet in een definitieve beslissing.

De Resultaten: Sneller en Slimmer

De auteurs hebben dit getest op twee dingen:

Sudoku-puzzels: Ze behandelden Sudoku als een planningsopdracht. Het "Relay"-model loste puzzels op met minder pogingen (minder "forward passes") en maakte minder fouten dan het standaardmodel. Het was beter in het "grote plaatje" in gedachten te houden terwijl het het rooster invulde.
Coderen (Software schrijven): Ze namen een state-of-the-art coderende AI (Fast-dLLM v2) en voegden het Relay-systeem toe.
- Betere Nauwkeurigheid: Het schreef betere code.
- Snellere Snelheid: Het moest zijn "denkmotor" 32% minder vaak draaien om hetzelfde resultaat te krijgen.

Waarom Dit Belangrijk Is

Het paper beweert dat we door de "Harde Reset" te stoppen en het model zijn eigen "gedachten" als een relay-slang mee te laten dragen, AI-modellen kunnen maken die:

Slummer: Ze kunnen beter redeneren over lange reeksen (zoals complexe code of logische puzzels).
Sneller: Ze hoeven niet elke keer alles opnieuw van scratch te berekenen, wat tijd en energie bespaart.

Kortom, het paper leert AI-modellen om te stoppen met het vergeten van hun eigen werk en te beginnen met daarop voortbouwen, stap voor stap, net zoals een mens dat doet bij het oplossen van een moeilijk probleem.

Technische Samenvatting: Gelernde Relaisrepresentaties voor Vooruitdenkende Discrete Diffusiemodellen

1. Probleemstelling: De "Harde Reset" in Gemaskerde Diffusiemodellen

Gemaskerde Diffusiemodellen (MDM's) genereren discrete sequenties via iteratieve denoising. Beginnend bij een volledig gemaskerde sequentie, onmaskert het model iteratief een fractie van posities. Er bestaat echter een kritieke inefficiëntie in de standaard MDM-inferentie: na elke forward pass verwerpt het model de rijke interne continue berekeningen (verborgen toestanden) die zijn gegenereerd voor gemaskerde posities. De volgende stap begint uitsluitend vanuit de nieuw vastgelegde discrete tokens.

De auteurs noemen dit het harde reset-probleem. Hoewel MDM's meerdere forward passes uitvoeren per generatie, verhindert het onvermogen om tussentijdse continue berekeningen te hergebruiken dat het model redenering over stappen effectief kan amortiseren. Dit staat in contrast met recurrente berekening, waarbij het uitrollen van een model met vaste parameters over stappen heeft aangetoond dat het de functieklasse uitbreidt die het model kan benaderen, wat de prestaties op complexe redeneertaken ten goede komt. De centrale vraag die wordt behandeld, is hoe MDM's in staat kunnen worden gesteld om recurrente berekening te ondersteunen die rijkere informatie draagt over opeenvolgende onmaskeringsstappen heen, zonder tussentijdse latente toestanden te verwerpen.

2. Methodologie: Gelernde Relaisrepresentaties (Relay)

De voorgestelde oplossing, Relay, introduceert een continue, differentieerbare latente toestand die wordt doorgegeven over decodestappen heen, waardoor het model "vooruitdenkend" wordt.

Kernmechanisme

Geaugmenteerde ToestandsTraject: Bij elke inferentiestap $k$ onderhoudt het model een geaugmenteerde toestand $s_k = (x_{t_k}, h_k)$ , waarbij $x_{t_k}$ de huidige gedeeltelijk gemaskerde sequentie is en $h_k$ een geleerde relais-toestand (continue latente vector) is.
Architectuur: Het model bestaat uit een backbone $f_\theta$ $f_{θ}$ , token-embeddings $Emb_\theta$ $E m b_{θ}$ , een unembedding-head $UnEmb_\theta$ $U n E m b_{θ}$ en een specifiek relaismodule $R_\theta$ $R_{θ}$ .
- De volgende verborgen toestand wordt berekend als: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ .
- De relais-toestand $h_k$ wordt geprojecteerd en opgeteld bij de ingebedde tokens voordat deze de backbone binnenkomen.
- De output-logits $\ell_k$ worden afgeleid van $h_{k+1}$ via $UnEmb_\theta$ .
Training via Afgeknotte BPTT: In tegenstelling tot standaard MDM's die worden getraind met onafhankelijke stappen, wordt Relay getraind met Afgeknotte Backpropagation Through Time (BPTT).
- Het model voert een rollout uit van $K$ stappen (bijv. $K=2$ ) beginnend bij een volledig gemaskerde sequentie.
- Tijdens de rollout worden posities ongemaskerd op basis van een beleid (bijv. betrouwbaarheidsdrempels), en worden de vastgelegde tokens "teacher-forced" naar ground truth-waarden om foutopstapeling tijdens de rollout te voorkomen.
- De loss is de som van cross-entropy losses over de $K$ stappen.
- Gradienten worden teruggepropageerd via het relaispad ( $h_k \to h_{k+1}$ ), waardoor het model niet alleen leert de huidige stap te voorspellen, maar ook informatie in $h_k$ te coderen die maximaal nuttig is voor toekomstige stappen.

Inferentie

Tijdens de inferentie blijft het proces grotendeels ongewijzigd wat betreft het onmaskeringsschema of het bemonsteringsbeleid. De enige toevoeging is dat het model de relais-toestand $h_k$ meedraagt naast de vastgelegde tokens. De relais-toestand dient als een differentieerbaar geheugenkanaal voor informatie die nog niet is vastgelegd in discrete tokens.

3. Belangrijkste Bijdragen

Relay Framework: Een algemene methode voor het integreren van recurrente berekening in MDM's door het model te trainen om een geleerde latente relais door te geven over decodestappen heen via afgeknotte BPTT. Het is architectuur-onafhankelijk en compatibel met bestaande MDM-versnellingstechnieken zoals block diffusion en KV-caching.
Schaalbaarheid en Adaptatie: Validatie van Relay op schaal van Large Language Models (LLM's) door volledige parameteradaptatie van Fast-dLLM v2 (1.5B). De methode slaagt erin een state-of-the-art Discrete Language Model (DLM) aan te passen met minimale extra training-FLOP's.
Prestatie-Latentie Verbeteringen: Empirisch bewijs dat Relay de prestatie-latentie Pareto-grens verbetert, waarbij het standaard supervised fine-tuning (SFT) overtreft op coderingstaken terwijl het de inferentielatentie aanzienlijk verlaagt.
Ontwerpruimte Analyse: Uitgebreide ablatiestudies op een uitdagende Sudoku-gebaseerde planningstaak die de bijdragen van de rollout-procedure, het relaiskanaal en het BPTT-trainingdoel isoleren.

4. Experimentele Resultaten

Sudoku Planningstaak

Met behulp van de Sudoku-Extreme dataset vergeleken de auteurs vier trainingdoelen:

MLM: Standaard uniforme gemaskerde diffusie (enkele stap, geen relais).
Rollout: Meerdere-staps rollout met teacher forcing maar zonder relaiskanaal ( $R_\theta \equiv 0$ ).
Relay (sg): Omvat het relaispad maar gebruikt stop-gradients (geen BPTT door het relais).
Relay: Volledige methode met $K=2$ BPTT door het relais.

Vondsten:

Nauwkeurigheids-NFE Grens: Relay bereikte de beste afweging tussen nauwkeurigheid en het Aantal Forward Evaluaties (NFE).
Invloed van Componenten: De overstap van MLM naar Rollout verbeterde de prestaties. Het toevoegen van het relaiskanaal (Relay sg) zorgde voor een aanzienlijke sprong. Tot slot leverde het inschakelen van BPTT (Relay) de beste resultaten op, wat aantoont dat het model leert het deels ingevulde bord zelfconsistent te houden onder agressievere onmaskering.
Efficiëntie: Bij een vaste betrouwbaarheidsdrempel legde Relay meer cellen vast per forward pass terwijl het de legaliteit behield (geen rij/kolom/vak-overschrijdingen), en bereikte het dezelfde nauwkeurigheid in minder totale forward passes in vergelijking met baselines.

Adaptatie van Voorgetraind Model (Fast-dLLM v2)

De auteurs pasten Fast-dLLM v2 (1.5B) aan met Relay op een mengsel van code- en wiskundedatasets.

Coderingsprestaties: Op HumanEval- en MBPP-benchmarks overtrof Relay standaard SFT.
- HumanEval: Relay behaalde 42,1% nauwkeurigheid (vs. 38,4% voor SFT) met 88,3 NFE (vs. 130,7 voor SFT).
- MBPP: Relay behaalde 46,6% nauwkeurigheid (vs. 43,9% voor SFT) met 78,8 NFE (vs. 84,8 voor SFT).
Latentie: Relay verlaagde de inferentielatentie met tot 32% in vergelijking met standaard SFT, terwijl de nauwkeurigheid verbeterde.
Geheugenoverhead: Profileren op een A100 GPU toonde aan dat hoewel Relay de live-activatiegeheugen tijdens de tweede forward pass verhoogt, het piekgeheugenverbruik bijna identiek blijft aan vanilla SFT (ongeveer 20,1 GiB vs. 21,2 GiB), omdat de bottleneck de terugwaartse pass van de vocabulaire-projectie is en niet de opslag van de relais-toestand.

5. Betekenis en Claims

Het artikel claimt dat Relay fundamenteel de aard van iteratieve niet-autoregressieve generatie verandert door MDM's in staat te stellen tussentijdse continue berekeningen op te hopen.

Vooruitdenkend: De methode stelt modellen in staat om expliciet tussentijdse representaties te optimaliseren voor toekomstige onmaskeringsbeslissingen, in plaats van elke stap te behandelen als een fris voorspellingsprobleem.
Verbetering van de Pareto-grens: Door het aantal forward passes te verminderen dat nodig is om een doelnauwkeurigheid te bereiken, verbetert Relay de prestatie-latentie-grens, waardoor discrete diffusiemodellen efficiënter worden voor complexe redeneer- en generatietaken.
Compatibiliteit: De aanpak wordt gepresenteerd als een plug-and-play-uitbreiding die werkt met bestaande state-of-the-art DLM's en versnellingstechnieken (bijv. block diffusion, KV-caching) zonder dat er ingrijpende architecturale wijzigingen nodig zijn.

De auteurs concluderen dat Relay een betekenisvolle stap voorwaarts is richting niet-gierige, vooruitdenkende generatie, en biedt een continu substraat voor latente redenering dat verder kan worden verkend voor interpretatie en toepassing in andere modaliteiten (bijv. afbeeldingen, moleculen).