Retrofitting Linear Attention into Diffusion Language Models

Stel je de wereld van kunstmatige intelligentie voor als een enorme, bruisende bibliotheek waar een superintelligente robot-bibliothecaris een verhaal probeert te schrijven. Jarenlang was de standaardmanier waarop deze bibliothecaris werkte "autoregressief": het schreef één woord, stopte, dacht na over alles wat het net had geschreven, schreef het volgende woord, stopte en dacht opnieuw na. Het was alsof je een toren bouwde steen voor steen, waarbij je wachtte tot de lijm droog was voordat je de volgende steen toevoegde. Hoewel deze methode betrouwbaar was, was het traag, vooral voor lange verhalen.

Onlangs arriveerde er een nieuwe stijl genaamd "Diffusion". In plaats van steen voor steen te bouwen, begint de Diffusion-bibliothecaris met een blanco pagina vol vraagtekens en probeert hij het hele verhaal in één keer te raden, waarbij hij de gissingen steeds opnieuw verfijnt totdat de woorden zin krijgen. Dit is alsof je een team van kunstenaars hebt dat tegelijkertijd een muurschildering schildert; ze kunnen veel sneller werken omdat ze niet hoeven te wachten tot de vorige steen droog is. Deze snelle methode heeft echter een probleem. Naarmate de muurschildering groter wordt, moeten de kunstenaars constant terugkijken naar elk enkel woord dat ze al hebben besloten om er zeker van te zijn dat de nieuwe woorden passen. Het is alsof je probeert een boek van 10.000 pagina's te onthouden terwijl je pagina 101 schrijft, wat uiteindelijk iedereen vertraagt en de geheugen van de bibliotheek vult.

Dit is het puzzelstukje waar een team onderzoekers een oplossing voor vond. Ze vroegen zich af: "Kunnen we deze snelle Diffusion-methode nog sneller maken door de manier waarop de bibliothecaris het verleden onthoudt te veranderen?" Hun antwoord is een slimme truc genaamd "Block-Hybrid Attention". Ze realiseerden zich dat hoewel de bibliothecaris de huidige paragraaf perfect moet onthouden (om de zinnen vloeiend te laten lopen), hij niet elke lettergreep van de vorige hoofdstukken in hoge definitie opnieuw hoeft te lezen. In plaats daarvan kan hij de oude hoofdstukken samenvatten in een kleine "samenvattingssheet". Dit stelt de bibliothecaris in staat om de nieuwe delen van het verhaal met bliksemsnelheid te schrijven zonder te worden afgeremd door het gewicht van de hele geschiedenis.

Het Grote Idee: Een Geheugensysteem met Twee Snelheden

Het artikel introduceert een methode om een bestaand, krachtig Diffusion Language Model (specifiek een model met 16 miljard parameters genaamd LLaDA 2.1) te upgraden (of te retrofitten) zodat het deze "samenvattingssheet"-strategie kan gebruiken zonder dat het vanaf nul opnieuw getraind hoeft te worden.

Beschouw het brein van het model als een brein met 20 verschillende lagen denken. In het originele model werkt elke laag als een super-attentieve bibliothecaris die elk vorig woord leest om het huidige woord te begrijpen. Dit is accuraat maar zwaar. De innovatie van de onderzoekers, Block-Hybrid Attention, splitst het werk op in twee modi:

1. De "Nu"-modus (Exacte Aandacht): Wanneer het model werkt aan het huidige blok woorden dat het actief genereert, behoudt het zijn "super-attentieve" modus. Het gebruikt het standaard, zware geheugen om elk woord in de huidige paragraaf te bekijken om ervoor te zorgen dat de zinnen perfect kloppen.
2. De "Verleden"-modus (Lineaire Aandacht): Wanneer het model de vorige blokken (de voltooide hoofdstukken) moet onthouden, schakelt het over naar "Lineaire Aandacht". In plaats van het hele boek opnieuw te lezen, raadpleegt het een samenvatting van een vaste grootte. Het is alsof je een index van een boek raadpleegt of een samenvatting van één pagina, in plaats van de hele tekst opnieuw te lezen. Deze samenvatting blijft dezelfde grootte, ongeacht hoe lang het boek wordt.

Hoe ze het deden: De Upgrade in Twee Fasen

De onderzoekers hebben de methode niet simpelweg vervangen en gehoopt op het beste; ze gebruikten een zorgvuldig tweefasig proces om ervoor te zorgen dat het model zijn intelligentie niet verloor tijdens de upgrade.

• Fase 1: De Schaduwtraining. Ze namen het originele, slimme model (de "Docent") en bevroren dit. Vervolgens vervingen ze 6 van de 20 aandachtslagen door hun nieuwe "Block-Hybrid"-versie (de "Leerling"). De Leerling werd getraind om de output van de Docent exact na te bootsen, gebruikmakend van een "gecorrumpeerde" versie van de tekst als input. Het was alsof een leerling een meesterkok volgde, probeerde de exacte bewegingen van de chef te kopiëren en precies dezelfde smaken proefde, maar dan alleen voor de specifieke gerechten die de leerling toegewezen kreeg. Deze fase duurde ongeveer 24 uur.
• Fase 2: De Fine-tuning. Zodra de Leerling leerde de Docent na te bootsen, lieten de onderzoekers het hele model weer samenwerken. Ze gebruikten een techniek genaamd LoRA (Low-Rank Adaptation) om minuscule, precieze aanpassingen te maken aan de verbindingen van het model. Dit was als het geven van een snelle "polijsting" aan het geüpgradede model om kleine imperfecties op te lossen die ontstonden doordat de nieuwe onderdelen nu met de oude onderdelen communiceerden. Deze fase duurde ongeveer 6 uur.

De Resultaten: Sneller, Lichter en Nog Steeds Slim

Het team testte hun geüpgradede model, dat ze LLaDA-HYBRID noemden, om te zien of het nog steeds goed was in schrijven en problemen oplossen, en of het daadwerkelijk sneller was.

Weet het nog steeds hoe het moet denken?
Verrassend genoeg wel. Ondanks dat ze 6 van de 20 lagen hadden vervangen door deze "samenvattingssheet"-methode, bleef de prestatie van het model zeer dicht bij het origineel.

• Op een coderingstest genaamd HumanEval scoorde het originele model 75,6%, en het hybride model scoende 72,0%.
• Op een wiskundetest genaamd CMATH kreeg het origineel 88,3%, en het hybride model 86,7%.
• Interessant genoeg verbeterde het hybride model op een andere coderingstest, MBPP+, van 57,7% naar 63,0%.
  Het artikel suggereert dat hoewel het model iets minder perfect is op de moeilijkste redeneertaken (zoals een moeilijke wetenschapstoets genaamd GPQA-Diamond, waar het daalde van 38,9% naar 30,8%), het grotendeels in staat is gebleven om code te schrijven en wiskundige problemen op te lossen.

Is het daadwerkelijk sneller?
Hier gebeurt de magie. Omdat het model niet langer de hele geschiedenis van het gesprek hoeft te herlezen voor elk nieuw woord, kan het meer verzoeken tegelijkertijd afhandelen.

• Bij het testen van hoeveel woorden het model per seconde kon genereren, was het hybride model 1,7 keer sneller dan het origineel bij zijn piekprestatie (het afhandelen van 128 gelijktijdige verzoeken).
• Het geheugengebruik van het originele model groeide naarmate het verhaal langer werd, waardoor het uiteindelijk tegen een muur aanliep waarbij het niet meer meer gebruikers kon afhandelen. Het hybride model kon, dankzij de "samenvattingssheet" van vaste grootte, meer gelijktijdige verzoeken afhandelen voordat het zonder geheugen kwam te zitten.

De Conclusie

Dit artikel suggereert dat we niet altijd een nieuwe, snellere AI vanaf nul hoeven te bouwen. In plaats daarvan kunnen we een krachtige, bestaande AI nemen en haar een "hybride geheugensysteem" geven. Door haar scherpe focus op het huidige moment te behouden terwijl ze het verleden samenvat in een compacte vorm, kunnen we deze modellen aanzienlijk sneller en efficiënter maken. De onderzoekers ontdekten dat deze upgrade kon worden uitgevoerd in ongeveer 60 uur training op publieke data, wat een veelbelovende weg biedt om AI-diensten sneller en goedkoper te maken zonder al te veel van hun intelligentie op te offeren. De auteurs merken echter op dat dit slechts het begin is; ze hebben slechts 6 lagen geüpgraded, en toekomstig werk zou manieren kunnen vinden om zelfs meer lagen te upgraden of nog langere verhalen te verwerken.