TokMem: One-Token Procedural Memory for Large Language Models

Het artikel introduceert TokMem, een procedureel geheugenframework voor grote taalmodellen dat herbruikbare taken compilaat tot één trainbaar geheugentoken, waardoor doelgerichte gedragingen met constante overhead worden bereikt zonder het basismodel aan te passen of contextoverhead te veroorzaken.

De kern

TokMem: De "Magische Sleutel" voor Slimme Taalmodellen

Stel je voor dat je een enorm slimme robot hebt (een Large Language Model of LLM) die alles kan, maar die een beetje vergeetachtig is als het gaat om specifieke taken. Om de robot iets te laten doen, moet je hem elke keer een heel lang, gedetailleerd instructieboekje geven.

Het probleem? Dat boekje wordt steeds dikker. Als je de robot 1000 verschillende taken wilt laten doen, moet je 1000 boekjes meenemen. De robot moet dan elke keer die hele stapel lezen voordat hij begint, wat traag is, veel energie kost en soms zelfs leidt tot verwarring.

TokMem is een nieuwe uitvinding die dit probleem oplost. Het is als een slimme sleutelkast voor de robot.

Hoe werkt het? (De Analogie)

In plaats van de robot elke keer het hele recept te laten lezen, leert TokMem de robot om één klein sleuteltje te gebruiken voor elke taak.

1. De Sleutels (Memory Tokens):
   Stel je voor dat je een taak hebt zoals "maak een gezond boodschappenlijstje". In plaats van de robot te vertellen: "Kijk naar de koelkast, zoek groenten, bereken calorieën, maak een lijst...", trainen we de robot om één speciaal sleuteltje te herkennen. Laten we dit sleuteltje [gezond_diner] noemen.
   Dit sleuteltje is geen tekst die je leest, maar een klein, digitaal signaal (een "token") dat direct in het geheugen van de robot zit.

2. De Sleutelkast (Memory Bank):
   De robot heeft nu een kastje met duizenden van deze sleuteltjes. Elke sleutel staat voor een specifieke vaardigheid:

   • Sleutel A: "Zoek recepten"
   • Sleutel B: "Bereken calorieën"
   • Sleutel C: "Schrijf een lijstje"

3. Het Gebruik (Inference):
   Als je de robot vraagt: "Maak een gezond dinerlijstje", hoeft de robot niet meer te lezen wat "gezond" betekent. Hij pakt direct de sleutel [gezond_diner] uit zijn kastje.
   Maar het mooie is: voor complexe taken kan de robot meerdere sleutels achter elkaar gebruiken.

   • Eerst pakt hij de sleutel [recept_zoeken].
   • Dan pakt hij de sleutel [calorie_berekenen].
   • Tot slot pakt hij de sleutel [lijstje_maken].

De robot voert deze taken uit alsof hij een dansstapje doet: hij gebruikt de sleutel, doet de actie, en pakt de volgende sleutel. Er is geen zware tekst meer om te lezen, alleen maar het trekken van de juiste sleutel.

Waarom is dit zo cool?

• Snelheid en Efficiëntie: Omdat de robot niet meer hoeft te lezen wat hij moet doen (geen lange teksten), is hij veel sneller. Het is alsof je van het lezen van een handleiding overschakelt op het indrukken van een knop.
• Geen Vergeten: Normaal gesproken, als je een robot leert om 1000 nieuwe dingen te doen, vergeet hij soms de oude dingen (zoals een mens die een nieuwe taal leert en de oude vergeten raakt). Bij TokMem blijft de robot "vrij" van zijn oude kennis. De nieuwe vaardigheden zitten in de extra sleuteltjes, niet in de hersenen van de robot zelf. Hij vergeet dus niets.
• Klein en Krachtig: Je hebt maar heel weinig extra geheugen nodig voor al die sleuteltjes. Het is alsof je een hele bibliotheek in één klein zakje stopt.

De Vergelijking met Andere Methoden

• De Oude Manier (Prompting): Dit is alsof je elke keer een heel boekje meeneemt naar de robot. "Hier is hoe je een lijstje maakt, hier is hoe je zoekt..." Dit wordt snel te zwaar.
• De Zoekmachine (RAG): Dit is alsof je een bibliotheek hebt en elke keer een boek moet zoeken en lezen voordat je iets doet. Het werkt, maar het kost tijd om het boek te vinden en te lezen.
• TokMem: Dit is alsof je een magische sleutel hebt. Je steekt hem in het slot, en de deur (de taak) gaat direct open. Geen zoeken, geen lezen, gewoon doen.

Conclusie

TokMem is een slimme manier om kunstmatige intelligentie te leren om taken uit te voeren zonder dat de robot zelf zwaarder of trager wordt. Het maakt de AI flexibeler, sneller en zorgt ervoor dat hij nooit zijn oude vaardigheden vergeet terwijl hij nieuwe leert. Het is de toekomst van hoe we computers laten "onthouden" hoe ze dingen moeten doen, in plaats van wat ze moeten doen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "TOKMEM: ONE-TOKEN PROCEDURAL MEMORY FOR LARGE LANGUAGE MODELS", geschreven in het Nederlands.

Probleemstelling

Grote Taalmodellen (LLM's) worden momenteel voornamelijk gestuurd via prompts. Hoewel dit effectief is, kent deze aanpak aanzienlijke nadelen:

1. Inefficiëntie: Prompts moeten bij elke nieuwe query opnieuw worden verwerkt, wat leidt tot hoge rekenkosten (kwadratische schaling van self-attention).
2. Contextbeperking: Lange prompts verbruiken het beperkte contextvenster, wat kan leiden tot truncatie en informatieverlies.
3. Schaalbaarheid: Het handmatig onderhouden en hergebruiken van complexe instructies voor verschillende taken is arbeidsintensief en schaalbaarheidsproblemen ontstaan naarmate het aantal taken groeit.
4. Bestaande oplossingen: Retrieval-Augmented Generation (RAG) lost dit deels op door externe documenten op te halen, maar deze blijven expliciete tekst die opnieuw geïnterpreteerd moet worden. Dit behoudt de overhead van lange contexten en behandelt procedurele kennis (hoe iets te doen) als declaratieve kennis (feiten), wat minder efficiënt is.

Methodologie: TokMem

De auteurs introduceren TokMem (One-Token Procedural Memory), een modulair framework dat procedurele kennis compliceert tot één trainbaar geheugentoken per taak.

Kernconcepten:

• Procedurele Tokens: In plaats van lange tekstuele instructies, wordt elke herbruikbare taakprocedure (bijv. "parse dieet", "zoek voedsel") gecodeerd als een uniek, trainbaar vector-embeddings (een speciaal token).
• Frozen Backbone: De basis-LLM blijft volledig bevroren (frozen). Alleen de embeddings van de geheugentokens worden getraind. Dit zorgt voor parameter-efficiëntie en stabiliteit.
• Index en Sturing: Elk geheugentoken fungeert dubbel:
  1. Als index om de juiste procedure te selecteren.
  2. Als sturingssignaal (steering signal) dat de generatie van het model direct beïnvloedt.
• Inference (Geheugenherroeping): Tijdens inferentie voert het model een query uit, voorspelt het de meest waarschijnlijke geheugentoken(s) en voegt deze toe aan de inputsequentie. Voor meervoudige taken (compositional recall) kan het model sequentieel meerdere tokens oproepen (bijv. eerst parsen, dan zoeken, dan formatteren) zonder lange prompts.

Technische Details:

• Training: Het trainen gebeurt via next-token prediction. De loss wordt alleen berekend op de posities van de geheugentokens en de bijbehorende responsen.
• Stabilisatie (Renormalization): Om te voorkomen dat nieuwe tokens de routing-dynamiek verstoren (norm-inflatie), wordt na het toevoegen van nieuwe tokens een renormalisatiestap uitgevoerd. Hierbij worden de normen van nieuwe embeddings geschaald om overeen te komen met de bestaande geheugenbank, zonder de richting te veranderen.
• Decoupled Variant (TokMem+DC): Een variant waarbij indexering en sturing worden gescheiden in twee matrices. De paper concludeert echter dat de eenvoudige, gekoppelde versie voor de meeste modellen voldoende is.

Belangrijkste Bijdragen

1. Compacte Procedurele Representatie: Het bewijs dat complexe taken kunnen worden gecomprimeerd tot één token, wat de contextoverhead drastisch verlaagt.
2. Continu Leren zonder Catastrophic Forgetting: Omdat de backbone bevroren blijft en procedures in geïsoleerde tokens worden opgeslagen, kan het model nieuwe vaardigheden toevoegen zonder bestaande kennis te overschrijven of te degraderen.
3. Superieure Routing: Het model leert niet alleen de taak uit te voeren, maar ook welk token er moet worden opgeroepen, wat leidt tot zeer nauwkeurige taakselectie.
4. Compositional Generalization: Het vermogen om losse procedures (tokens) te combineren om complexe, meervoudige stappen uit te voeren die niet expliciet tijdens het trainen zijn gezien.

Resultaten

De auteurs evalueren TokMem in twee scenario's:

1. Atomische Herroeping (Super-Natural Instructions):

• Opzet: 1.000 verschillende NLP-taken worden als aparte procedures getraind.
• Resultaat: TokMem presteert consistent beter dan RAG en fine-tuning (LoRA). Terwijl RAG snel degradeert naarmate het aantal taken toeneemt (door ruis in de retrieval), behoudt TokMem een hoge nauwkeurigheid (>94% routing-nauwkeurigheid zelfs bij 1.000 taken).
• Efficiëntie: TokMem vereist aanzienlijk minder trainingsdata om nieuwe taken te leren vergeleken met fine-tuning.

2. Compositional Herroeping (Function Calling / APIGen):

• Opzet: Het samenvoegen van meerdere tool-aanroepen (bijv. eerst data ophalen, dan berekenen, dan formatteren).
• Resultaat: TokMem presteert gelijkwaardig aan of beter dan LoRA-fine-tuning, maar met orde van grootte minder trainbare parameters (bijv. 0.2M parameters voor TokMem vs 3.41M voor LoRA op Llama 8B).
• Generalisatie: TokMem toont sterkere generalisatie naar langere ketens van tool-aanroepen dan fine-tuning, zelfs wanneer het model alleen is getraind op kortere ketens.

Vergelijking met Prefix Tuning:
In experimenten met tekstmemorization (Fanfics-dataset) bleek TokMem (met infix-plaatsing: query-token-respons) sneller te convergeren en een lagere perplexiteit te bereiken dan traditionele prefix tuning, vooral bij een klein aantal tokens.

Betekenis en Impact

TokMem biedt een paradigmaverschuiving in hoe LLM's worden gestuurd:

• Van Declaratief naar Procedureel: Het verschuift de focus van het opslaan van feiten (tekst) naar het opslaan van vaardigheden (tokens), wat dichter bij menselijk procedurale geheugen ligt.
• Schaalbare Agenten: Het maakt het mogelijk om agents te bouwen met enorme bibliotheken aan vaardigheden zonder dat de contextvenster of de rekenkosten exponentieel toenemen.
• Personalisatie: Omdat de backbone bevroren blijft, kunnen gebruikers of organisaties hun eigen "memory banks" met aangepaste procedures toevoegen zonder het basismodel opnieuw te hoeven trainen.

Kortom, TokMem lost het probleem van inefficiënte prompts op door procedurele kennis te distilleren tot compacte, trainbare eenheden, waardoor LLM's flexibeler, efficiënter en beter in staat zijn tot continu leren worden.