AI-modeller og platforme

Supercharging Large Language Models med Multi-token Prediction

mm
Føj Unite.AI til dine foretrukne kilder på Google

Store sprogmodeller (LLM’er) som GPT, LLaMA og andre har taget verden med storm med deres bemærkelsesværdige evne til at forstå og generere menneske-lignende tekst. Men på trods af deres imponerende evner har den standardmetode til træning af disse modeller, kendt som “next-token prediction”, nogle indbyggede begrænsninger.

I next-token prediction trænes modellen til at forudsige det næste ord i en sekvens givet de foregående ord. Selvom denne tilgang har vist sig at være succesfuld, kan den føre til modeller, der kæmper med lange afhængigheder og komplekse resoneringstasks. Desuden kan mismatchet mellem den lærer- tvangstræning og den autoregressive generationsproces under inference føre til underoptimal præstation.

En ny forskningsartikel af Gloeckle et al. (2024) fra Meta AI introducerer en ny træningsparadigme kaldet “multi-token prediction“, der sigter mod at løse disse begrænsninger og give store sprogmodeller et boost. I denne blogpost dykker vi dybt ned i de centrale begreber, tekniske detaljer og potentielle implikationer af denne banebrydende forskning.

Single-token Prediction: Den konventionelle tilgang

Før vi dykker ned i detaljerne om multi-token prediction, er det essentiel at forstå den konventionelle tilgang, der har været hovedsagen i store sprogmodelltræning i år – single-token prediction, også kendt som next-token prediction.

The Next-token Prediction Paradigm

I next-token prediction-paradigmet trænes sprogmodeller til at forudsige det næste ord i en sekvens givet den foregående kontekst. Mere formelt er modellen opgaven at maksimere sandsynligheden for det næste token xt+1, givet de foregående tokens x1, x2, …, xt. Dette gøres typisk ved at minimere cross-entropy-taben:

L = -Σt log P(xt+1 | x1, x2, …, xt)

Dette simple, men kraftfulde træningsmål har været grundlaget for mange succesfulde store sprogmodeller, såsom GPT (Radford et al., 2018), BERT (Devlin et al., 2019) og deres varianter.

Teacher Forcing og Autoregressive Generation

Next-token prediction afhænger af en træningsteknik kaldet “teacher forcing“, hvor modellen får den korrekte grundsandhed for hvert fremtidigt token under træning. Dette giver modellen mulighed for at lære fra den korrekte kontekst og målsekvenser, hvilket faciliterer mere stabil og effektiv træning.

Men under inference eller generation opererer modellen på en autoregressiv måde, hvor den forudsiger ét token ad gangen baseret på de tidligere genererede tokens. Dette mismatch mellem træningsregimet (teacher forcing) og inferensregimet (autoregressiv generation) kan føre til potentielle diskrepanser og underoptimal præstation, især for længere sekvenser eller komplekse resoneringstasks.

Begrænsninger i Next-token Prediction

Selvom next-token prediction har været bemærkelsesværdigt succesfuld, har den også nogle indbyggede begrænsninger:

  1. Kort sigt: Ved kun at forudsige det næste token, kan modellen have svært ved at fange lange afhængigheder og den overordnede struktur og kohærens i teksten, hvilket potentielt kan føre til inkonsistenser eller inkohærente generationer.
  2. Lokal mønsterlåsning: Next-token prediction-modeller kan låse sig fast på lokale mønstre i træningsdata, hvilket gør det svært at generalisere til udenfor-distributionsscenarioer eller opgaver, der kræver mere abstrakt resonering.
  3. Resoneringsevner: For opgaver, der involverer multi-trin resonering, algoritmic tænkning eller komplekse logiske operationer, kan next-token prediction ikke give tilstrækkelige induktive bias eller repræsentationer til at støtte sådanne evner effektivt.
  4. Prøveineffektivitet: På grund af den lokale natur af next-token prediction, kan modellerne kræve større træningsdata for at erhverve den nødvendige viden og resoneringsevner, hvilket kan føre til potentielle prøveineffektiviteter.

Disse begrænsninger har motiveret forskere til at udforske alternative træningsparadigmer, såsom multi-token prediction, der sigter mod at løse nogle af disse svagheder og låse op for nye evner for store sprogmodeller.

Ved at kontrastere den konventionelle next-token prediction-tilgang med den nye multi-token prediction-teknik, kan læserne bedre værdsætte motivationen og de potentielle fordele ved sidstnævnte, hvilket sætter scenen for en dybere udforskning af denne banebrydende forskning.

Hvad er Multi-token Prediction?

Den centrale idé bag multi-token prediction er at træne sprogmodeller til at forudsige multiple fremtidige tokens samtidig, i stedet for kun det næste token. Specifikt under træning er modellen opgaven at forudsige de næste n tokens på hver position i træningskorpus, ved hjælp af n uafhængige output-lag, der opererer oven på en fælles model-stamme.

For eksempel med en 4-token prediction-opstilling, ville modellen blive trænet til at forudsige de næste 4 tokens på én gang, givet den foregående kontekst. Denne tilgang opmuntrer modellen til at fange længere afhængigheder og udvikle en bedre forståelse af den overordnede struktur og kohærens i teksten.

Et Legetøjseksempel

For at bedre forstå begrebet om multi-token prediction, lad os overveje et simpelt eksempel. Antag, at vi har følgende sætning:

“Den hurtige brune ræv hopper over den dovne hund.”

I den standard next-token prediction-tilgang ville modellen blive trænet til at forudsige det næste ord, givet den foregående kontekst. For eksempel, givet konteksten “Den hurtige brune ræv hopper over den”, ville modellen være opgaven at forudsige det næste ord, “dovne”.

Med multi-token prediction ville modellen derimod blive trænet til at forudsige multiple fremtidige ord på én gang. For eksempel, hvis vi sætter n=4, ville modellen blive trænet til at forudsige de næste 4 ord samtidig. Givet den samme kontekst “Den hurtige brune ræv hopper over den”, ville modellen være opgaven at forudsige sekvensen “dovne hund “. (Bemærk mellemrummet efter “hund” for at indikere slutningen af sætningen).

Ved at træne modellen til at forudsige multiple fremtidige tokens på én gang, opmuntrer den modellen til at fange lange afhængigheder og udvikle en bedre forståelse af den overordnede struktur og kohærens i teksten.

Tekniske detaljer

Forfatterne foreslår en simpel, men effektiv arkitektur til implementering af multi-token prediction. Modellen består af en fælles transformer-stamme, der producerer en latent repræsentation af input-konteksten, efterfulgt af n uafhængige transformer-lag (output-hoveder), der forudsiger de respektive fremtidige tokens.

Under træning er forward- og backward-passerne omhyggeligt orkestrerede for at minimere GPU-hukommelsesaftryk. Den fælles stamme beregner den latente repræsentation, og derefter udfører hver output-hoved sin forward- og backward-pas, akkumulerende gradienter på stamme-niveauet. Denne tilgang undgår at materialisere alle logit-vektorer og deres gradienter samtidig, reducerer den maksimale GPU-hukommelsesbrug fra O(nV + d) til O(V + d), hvor V er vokabular-størrelsen og d er dimensionen af den latente repræsentation.

Den hukommelses-effektive Implementering

En af udfordringerne ved træning af multi-token-prædiktorer er at reducere deres GPU-hukommelsesudnyttelse. Da vokabular-størrelsen (V) typisk er meget større end dimensionen af den latente repræsentation (d), bliver logit-vektorer hukommelses-udnyttelses-bottleneck.

For at løse denne udfordring foreslår forfatterne en hukommelses-effektiv implementering, der omhyggeligt tilpasser sekvensen af forward- og backward-operationer. I stedet for at materialisere alle logit-vektorer og deres gradienter samtidig, beregner implementeringen sekventielt forward- og backward-passer for hver uafhængig output-hoved, akkumulerende gradienter på stamme-niveauet.

Denne tilgang undgår at gemme alle logit-vektorer og deres gradienter i hukommelsen samtidig, reducerer den maksimale GPU-hukommelsesudnyttelse fra O(nV + d) til O(V + d), hvor n er antallet af fremtidige tokens, der forudses.

Fordele ved Multi-token Prediction

Forskningsartiklen præsenterer flere overbevisende fordele ved at bruge multi-token prediction til træning af store sprogmodeller:

  1. Forbedret Prøveeffektivitet: Ved at opmuntre modellen til at forudsige multiple fremtidige tokens på én gang, driver multi-token prediction modellen mod bedre prøveeffektivitet. Forfatterne demonstrerer betydelige forbedringer i præstation på kodeforståelse og generationsopgaver, med modeller op til 13B parametre, der løser omkring 15% flere problemer i gennemsnit.
  2. Hurtigere Inference: De ekstra output-hoveder, der trænes med multi-token prediction, kan udnyttes til selv-spekulativ decoding, en variant af spekulativ decoding, der tillader parallel token-forudsigelse. Dette resulterer i op til 3 gange hurtigere inferencetider over en bred vifte af batch-størrelser, selv for store modeller.
  3. Fremme af Lange-afhængigheder: Multi-token prediction opmuntrer modellen til at fange længere afhængigheder og mønstre i data, hvilket er særligt gavnligt for opgaver, der kræver forståelse og resonering over større kontekster.
  4. Algoritmic Resonering: Forfatterne præsenterer eksperimenter på syntetiske opgaver, der demonstrerer overlegenheden af multi-token prediction-modeller i udvikling af induktionshoveder og algoritmic resoneringsevner, især for mindre modellstørrelser.
  5. Kohærens og Konsistens: Ved at træne modellen til at forudsige multiple fremtidige tokens samtidig, opmuntrer multi-token prediction udviklingen af kohærente og konsistente repræsentationer. Dette er særligt gavnligt for opgaver, der kræver generering af længere, mere kohærent tekst, såsom fortælling, kreativ skrivning eller generering af instruktionsmanualer.
  6. Forbedret Generalisering: Forfatternes eksperimenter på syntetiske opgaver antyder, at multi-token prediction-modeller udviser bedre generaliserings-evner, især i udenfor-distributionssætninger. Dette kan være på grund af modellens evne til at fange længere afhængigheder og mønstre, hvilket kan hjælpe den til at extrapolere mere effektivt til usete scenarier.

Eksempler og Intuitioner

For at give mere intuition om, hvorfor multi-token prediction virker så godt, lad os overveje et par eksempler:

  1. Kodegenerering: I sammenhæng med kodegenerering kan forudsigelse af multiple tokens samtidig hjælpe modellen til at forstå og generere mere komplekse kode-strukturer. For eksempel, når der genereres en funktiondefinition, kan forudsigelse af kun det næste token ikke give nok kontekst for modellen til at generere hele funktionssignaturen korrekt. Men ved at forudsige multiple tokens på én gang, kan modellen bedre fange afhængighederne mellem funktionens navn, parametre og returtype, hvilket fører til mere præcis og kohærent kodegenerering.
  2. Naturlig Sprogresonering: Overvej en situation, hvor en sprogmodel er opgaven at besvare et spørgsmål, der kræver resonering over flere trin eller dele af information. Ved at forudsige multiple tokens samtidig, kan modellen bedre fange afhængighederne mellem de forskellige komponenter i resoneringen, hvilket fører til mere kohærente og præcise svar.
  3. Lang-form Textgenerering: Når der genereres lang-form tekst, såsom historier, artikler eller rapporter, kan det være svært for sprogmodeller, der er trænet med next-token prediction, at opretholde kohærens og konsistens over en længere periode. Multi-token prediction opmuntrer modellen til at udvikle repræsentationer, der fanger den overordnede struktur og flow i teksten, hvilket potentielt fører til mere kohærent og konsistent lang-form tekst.

Begrænsninger og Fremtidige Retninger

Selvom resultaterne i artiklen er imponerende, er der nogle begrænsninger og åbne spørgsmål, der kræver yderligere undersøgelse:

  1. Optimalt Antal Tokens: Artiklen udforsker forskellige værdier af n (antallet af fremtidige tokens, der skal forudses) og finder, at n=4 virker godt for mange opgaver. Men den optimale værdi af n kan afhænge af den specifikke opgave, datasæt og modellstørrelse. Udvikling af principper for bestemmelse af den optimale n kan føre til yderligere forbedringer i præstation.
  2. Vokabular-størrelse og Tokenisering: Forfatterne bemærker, at den optimale vokabular-størrelse og tokeniseringstrategi for multi-token prediction-modeller kan afvige fra dem, der bruges til next-token prediction-modeller. Udforskning af dette aspekt kan føre til bedre kompromiser mellem komprimeret sekvenslængde og beregnings-effektivitet.
  3. Biprediktions-tab: Forfatterne foreslår, at deres arbejde kan føre til interesse for udvikling af nye biprediktions-tab for store sprogmodeller, ud over den standard next-token prediction. Udforskning af alternative biprediktions-tab og deres kombinationer med multi-token prediction er en spændende forskningsretning.
  4. Teoretisk Forståelse: Selvom artiklen giver nogle intuitioner og empiriske beviser for effektiviteten af multi-token prediction, ville en dybere teoretisk forståelse af, hvorfor og hvordan denne tilgang virker så godt, være værdifuld.

Konklusion

Forskningsartiklen “Bedre og Hurtigere Store Sprogmodeller via Multi-token Prediction” af Gloeckle et al. introducerer en ny træningsparadigme, der har potentialet til at betydeligt forbedre præstationen og evnerne hos store sprogmodeller. Ved at træne modeller til at forudsige multiple fremtidige tokens samtidig, opmuntrer multi-token prediction udviklingen af lange afhængigheder, algoritmic resoneringsevner og bedre prøveeffektivitet.

Den tekniske implementering, der foreslås af forfatterne, er elegant og beregnings-effektiv, hvilket gør det muligt at anvende denne tilgang til stor-skala sprogmodelltræning. Desuden kan den ekstra output-hoveder, der trænes med multi-token prediction, udnyttes til selv-spekulativ decoding, hvilket resulterer i op til 3 gange hurtigere inferencetider.

Selvom der stadig er åbne spørgsmål og aspekter, der kræver yderligere undersøgelse, repræsenterer denne forskning et spændende skridt fremad i feltet store sprogmodeller. Da efterspørgslen efter mere kapable og effektive sprogmodeller fortsætter med at vokse, kan multi-token prediction blive en nøglekomponent i den næste generation af disse kraftfulde AI-systemer.

Jeg har brugt de sidste fem år på at dykke ned i den fascinerende verden af Machine Learning og Deep Learning. Min passion og ekspertise har ført mig til at bidrage til over 50 forskellige software-ingeniørprojekter, med en særlig fokus på AI/ML. Min fortsatte nysgerrighed har også ført mig mod Natural Language Processing, et felt jeg er ivrig efter at udforske yderligere.