Joint Optimization of Retrieval-Augmented KV Caching and Adaptive Speculative Decoding for Autoregressive Generation

Wanneer een groot taalmodel een zin schrijft, woord voor woord, wordt het geconfronteerd met een terugkerende flessenhals: om het volgende woord te voorspellen, moet het elk woord dat eraan voorafging onthouden. Dit geheugen, opgeslagen als een enorme verzameling getallen die een key-value cache wordt genoemd, groeit bij elk nieuw token. In de moderne computertechnologie is het ophalen van dit geheugen vaak langzamer dan de eigenlijke berekening van het volgende woord. Om dit te versnellen, hebben ingenieurs twee belangrijke trucjes ontwikkeld. De eerste is het hergebruiken van geheugen van eerdere verzoeken als de beginwoorden identiek zijn, vergelijkbaar met het overslaan van de eerste paar pagina's van een boek dat je al hebt gelezen. De tweede is speculatieve decodering, waarbij een kleiner, sneller model enkele woorden vooruit gokt, en het hoofdmodel snel controleert of die gokken correct zijn. Als de gokken juist zijn, slaat het systeem het trage werk over om ze één voor één te genereren.

De vraag die onderzoekers stelden, was of het combineren van deze twee trucjes een supersnel systeem zou creëren. Intuïtief leek het logisch: hergebruik het geheugen dat je al hebt, en gok vervolgens een stukje vooruit om tijd te besparen. Echter, de realiteit van computerhardware is vaak minder vergevingsgezind dan de intuïtie suggereert. Een nieuwe studie door Chenyi Huang en collega's van Johns Hopkins University, Carnegie Mellon University, New York University en de University of Southern California onderzoekt precies deze combinatie. Zij bouwden een systeem dat deze twee trucs niet blindelings toepast. In plaats daarvan fungeert het als een zorgvuldige manager die constant de kosten van elke mogelijke actie meet en beslist of het voor elk specifiek verzoek het geheugen moet hergebruiken, vooruit moet gokken, of dat het beide niet moet doen.

De onderzoekers testten hun systeem, dat zij RAKV-ASD noemen, op een standaard grafische kaart van een laptop met behulp van twee paren taalmodellen. Ze voerden duizenden getimede experimenten uit, waarbij ze het systeem diverse prompts met verschillende lengtes voedden en het tekst van variërende groottes liet genereren. De kern van hun werk was een controller die naar de huidige situatie keek en de beste route koos uit een menu van opties. Deze opties omvatten het hergebruiken van geheugen zonder te gokken, het hergebruiken van geheugen terwijl er één, twee of tot wel acht woorden vooruit wordt gegokt, of het volledig negeren van de geheugencache en het vanaf nul beginnen. Het systeem mat hoe lang elke stap daadwerkelijk op de hardware duurde, in plaats van te vertrouwen op theoretische schattingen.

De resultaten waren verrassend voor iedereen die een universele versnelling verwachtte. De studie toonde aan dat het simpelweg combineren van geheugenhergebruik met een vaste hoeveelheid gokken het systeem vaak juist langzamer maakte, in plaats van sneller. In veel gevallen kostte de tijd die werd besteed aan het genereren en vervolgens verifiëren van de gokken meer tijd dan de tijd die werd bespaard door het oorspronkelijke werk over te slaan. De onderzoekers ontdekten dat het "gokken"-gedeelte van het proces alleen de moeite waard is wanneer het systeem bijna zeker weet dat de gokken correct zullen zijn. In hun experimenten koos de controller er regelmatig voor om helemaal niets te doen met betrekking tot het gokken, waarbij een instelling werd gekozen waarbij het geheugen werd hergebruikt maar de tekst woord voor woord werd gegenereerd. Deze beslissing was geen falen van het systeem, maar een succes van de logica ervan: het herkende dat de overhead van het gokken te hoog was voor de specifieke omstandigheden van de test.

Toen de onderzoekers hun slimme, adaptieve systeem vergeleken met een versie die het systeem dwong om altijd vier woorden vooruit te gokken, was het verschil treffend. Het adaptieve systeem verminderde de tijd die nodig was om tekst te genereren met 64,0 procent en verhoogde het aantal gegenereerde woorden per seconde met 170,8 procent vergeleken met de gedwongen gokbenadering. Dit komt omdat het adaptieve systeem de zware straf van mislukte gokken vermeed. Echter, wanneer het werd vergeleken met een systeem dat alleen het geheugen hergebruikte zonder enig gokken, was het adaptieve systeem slechts 0,45 procent langzamer. Dit minieme verschil suggereert dat het voor de specifieke hardware en modellen die werden getest, het vooruit gokken geen echt voordeel bood. Het systeem was zo goed in het weten wanneer het niet moest gokken, dat het in essentie fungeerde als een systeem dat nooit gokt.

De studie stelde ook een duidelijke grens vast voor wanneer gokken nuttig zou worden. De onderzoekers berekenden dat het systeem pas zou profiteren van gokken als de waarschijnlijkheid dat de gokken correct zouden zijn een zeer hoge drempel bereikte, tussen de 89,4 procent en 97,3 procent. In hun tests lag het werkelijke acceptatiepercentage van de gokken lager dan dit, wat betekende dat de omstandigheden niet gunstig waren voor de extra snelheid. De controller van het systeem was opmerkelijk nauwkeurig in het maken van dit onderscheid. Hoewel het niet perfect kon voorspellen hoe lang een taak exact zou duren, was het 100 procent nauwkeurig in het beslissen of het wel of niet zou gokken. Het identificeerde correct dat de kosten van het gokken zwaarder wogen dan de voordelen in bijna elk scenario dat ze testten.

Dit werk benadrukt een cruciale les voor de toekomst van kunstmatige intelligentiesystemen: complexere strategieën zijn niet altijd beter. Het idee om meerdere optimalisatietechnieken op elkaar te stapelen klinkt efficiënt, maar als de individuele onderdelen niet perfect aansluiten op de capaciteiten van de hardware, kunnen ze elkaar tenietdoen. De onderzoekers toonden aan dat een systeem dat weet wanneer het moet stoppen, vaak sneller is dan een systeem dat probeert alles te doen. Door de beslissing om te gokken als een variabele keuze te behandelen in plaats van als een vaste regel, creëerden ze een methode die zich aanpast aan de realiteit van de machine. De bevindingen suggereren dat voor de huidige hardware de meest efficiënte weg vaak inhoudt dat je hergebruikt wat je weet en de rest zorgvuldig genereert, in plaats van vooruit te haasten met gokken die mogelijk gecorrigeerd moeten worden.

De studie werd uitgevoerd met strikte aandacht voor detail, waarbij 3.600 getimede uitvoeringen werden gedraaid om te waarborgen dat de resultaten niet slechts een toevalstreffer waren. Ze verifieerden dat de tekst die door hun adaptieve systeem werd geproduceerd identiek was aan de tekst die door de standaard, langzamere methode werd geproduceerd, wat bewees dat de snelheidswinsten niet ten koste gingen van de nauwkeurigheid. Hoewel de resultaten specifiek zijn voor de modellen en hardware die in het experiment werden gebruikt, houdt het principe een bredere waarheid in stand: in de complexe wereld van grote taalmodellen is de slimste zet vaak degene die herkent wanneer een kortere route eigenlijk een omweg is. De onderzoekers vonden geen wondermiddel dat alle AI-generatie versnelt, maar ze vonden wel een betrouwbare manier om de vertragingen te vermijden die voortkomen uit het proberen te slim zijn. Hun werk biedt een gemeten, op data gebaseerde grens voor wanneer speculatie helpt en wanneer het hindert, en biedt een duidelijke gids voor het bouwen van efficiëntere systemen in de toekomst.