Eliciting Numerical Predictive Distributions of LLMs Without Autoregression

Deze studie toont aan dat regressieprobes, getraind op interne representaties van grote taalmodellen, statistische eigenschappen van numerieke voorspellingen direct kunnen afleiden zonder kostbare autoregressieve generatie.

De kern

Titel: De LLM's 'Gedachten' lezen zonder te praten

Stel je voor dat een Grote Taalmodel (LLM) een superintelligente voorspeller is. Als je hem een reeks getallen geeft (bijvoorbeeld de temperatuur van de afgelopen week), kan hij een voorschatting maken voor morgen. Maar er is een groot probleem: deze modellen werken als een telefoonnummer die je één voor één moet intikken.

Om een getal als "123,45" te voorspellen, moet het model eerst "1" zeggen, dan "2", dan "3", dan de komma, en pas dan "4" en "5". Dit heet autoregressie. Het is traag, kost veel energie en als je twijfelt over de voorspelling (bijvoorbeeld: "Hoe zeker is hij?"), moet je het model 100 keer laten "praten" om 100 verschillende antwoorden te krijgen. Dat is als een kok die 100 keer dezelfde soep moet koken om te zien of hij goed is.

De auteurs van dit paper hebben een slimme truc bedacht: Waarom wachten tot het model praat? Laten we gewoon in zijn hoofd kijken.

De Kern van het Onderzoek: De "Gedachte-lezer"

De onderzoekers vragen zich af: Zit het antwoord al in het brein van het model, voordat het ook maar één woord (token) heeft uitgesproken?

Het antwoord is een volmondig JA.

Ze hebben een soort "Gedachte-lezer" (in de vaktaal een probing model) gebouwd. Dit is een klein, simpel hulpmiddel dat in het brein van het grote model kijkt en de antwoorden eruit haalt zonder dat het grote model hoeft te "praten".

Hoe werkt het? (De Analogie van de Koerier)

Stel je voor dat het grote model een koerier is die een pakketje (het getal) moet bezorgen.

1. De oude manier (Autoregressie): De koerier loopt naar de deur, pakt het pakket, schrijft het adres op, loopt naar de bus, rijdt naar de stad, stopt bij de straat, loopt naar het huis en belt aan. Dit duurt lang.
2. De nieuwe manier (Probing): De onderzoekers kijken naar de spieren en de houding van de koerier terwijl hij nog in het magazijn staat. Ze zien dat zijn spieren al gespannen zijn in een specifieke richting. Ze kunnen al zien: "Ah, hij gaat naar huis nummer 123,45!" Zonder dat hij ook maar één stap heeft gezet.

De Uitdaging: Getallen zijn lastig

Getallen zijn lastig voor AI. Een getal kan klein zijn (0,001) of gigantisch (1.000.000). Als je een AI vraagt om een getal te raden, is het alsof je vraagt: "Is het een muis of een olifant?" en "Hoe groot is hij precies?".

De onderzoekers hebben een slimme oplossing bedacht, een tweestaps-plan:

1. De Grootte-Gok: Eerst kijken ze of het een muis, een hond of een olifant is (de orde van grootte, bijvoorbeeld: "Is het in de duizenden?").
2. De Precisie-Gok: Daarna kijken ze hoe groot die olifant precies is (de exacte cijfers).

Door dit in twee stappen te doen, kunnen ze zelfs de onzekerheid van het model voorspellen. Ze kunnen zeggen: "Het model denkt dat het 123 is, maar hij is 80% zeker dat het ergens tussen 120 en 126 ligt."

Waarom is dit geweldig?

1. Snelheid: Het is als het verschil tussen een brief per post sturen (langzaam, stap voor stap) en een e-mail sturen (direct). Je hoeft niet 100 keer te wachten op het antwoord; je haalt het er direct uit.
2. Energie: Het kost veel minder rekenkracht. Je hoeft het grote model niet 100 keer te laten "denken", maar slechts één keer te "kijken".
3. Vertrouwen: Je krijgt direct te zien hoe zeker het model is. Dit is cruciaal voor dingen zoals medische diagnoses of beursvoorspellingen, waar je niet alleen een antwoord wilt, maar ook een waarschuwing als het antwoord twijfelachtig is.

Conclusie

Deze paper laat zien dat grote taalmodellen veel meer "weten" dan ze laten horen. Het antwoord zit al in hun interne gedachten, voordat ze ook maar een letter hebben uitgesproken. Met deze nieuwe "Gedachte-lezer" kunnen we sneller, goedkoper en slimmer gebruik maken van AI voor het voorspellen van getallen, zonder dat we hoeven te wachten tot het model langzaam zijn antwoord opbouwt.

Het is alsof we eindelijk de telepathische vaardigheden van de AI hebben ontdekt, in plaats van te wachten tot hij zijn mond opent.

Technische samenvatting

Probleemstelling

Grote Taalmodellen (LLMs) worden steeds vaker succesvol toegepast op regressietaken, zoals tijdreeksvoorspelling en tabulair dataverwachten, dankzij hun vermogen tot in-context learning. Een fundamentele beperking van deze toepassing is echter de autoregressieve decoderingsprocess. Omdat reële getallen vaak meerdere tokens beslaan, vereist het genereren van één numerieke voorspelling meerdere forward passes door het model.

Dit proces wordt extreem kostbaar en tijdrovend wanneer men voorspellingsverdelingen (predictive distributions) nodig heeft om onzekerheid te kwantificeren (bijvoorbeeld via Monte Carlo-sampling). Het herhaaldelijk genereren van samples om statistieken zoals het gemiddelde, de mediaan of kwantielen te schatten, leidt tot hoge computationele kosten en lange inferentietijden. De centrale vraag van dit paper is: Kan men de numerieke voorspellingsverdeling en de onderliggende onzekerheid van een LLM afleiden uit de interne representaties zonder de dure autoregressieve generatie uit te voeren?

Methodologie

De auteurs onderzoeken of de interne verborgen staten (hidden states) van een LLM voldoende informatie bevatten om statistische functionals van de numerieke uitkomstverdeling direct te voorspellen. Ze introduceren een probing-benadering waarbij onafhankelijke modellen worden getraind op de embeddings van de LLM.

1. Magnitude-gefactoriseerde regressie (Magnitude-Factorised Regression):
Een grote uitdaging bij het voorspellen van numerieke waarden is de enorme variatie in orde van grootte (magnitude). Standaard regressieverliezen (zoals MSE) falen vaak bij deze schaalvariatie. De auteurs lossen dit op met een tweeledig model:

• Magnitude Classificator ($f_{order}$): Een classifier die de orde van grootte ($m = \lfloor \log_{10}(|y|) \rfloor$) van het doelgetal voorspelt.
• Schaal-invariante Regressor ($f_{val}$): Een regressor die de geschaalde waarde voorspelt, geconditioneerd op de voorspelde orde van grootte.
  De uiteindelijke voorspelling wordt berekend als $\hat{y} = r_k \cdot 10^{m_k}$, waarbij $r_k$ de voorspelde schaalwaarde is voor de $k$-de magnitude-klasse.

2. Doelstellingen voor Point Estimates (Sectie 2):
Er worden probes getraind om drie statistieken van de LLM-verdeling te recupereren:

• De greedy output (de meest waarschijnlijke token-sequentie).
• Het gemiddelde (mean) van de verdeling.
• De mediaan (median) van de verdeling.
  De targets worden geschat via empirische sampling van de LLM (100 samples per input) tijdens het trainen van de probe.

3. Kwantielregressie voor Onzekerheid (Sectie 3):
Om de volledige verdelingsvorm en onzekerheid te vangen, gebruiken de auteurs kwantielregressie met de "pinball loss". Ze trainen probes om meerdere kwantielen (bijv. 2.5%, 25%, 50%, 75%, 97.5%) direct te voorspellen. Dit stelt hen in staat om betrouwbaarheidsintervallen en het interkwartielbereik (IQR) af te leiden zonder sampling.

4. Data en Setup:

• LLM: Voornamelijk Llama-2-7B (waarbij de tokenizer elke digit als apart token behandelt, wat de uitdaging vergroot).
• Data: Synthetische tijdreeksen (sinus, Gaussisch, ruis) met verschillende schalen ($[-1, 1]$ tot $[-10000, 10000]$) en real-world datasets (Monash, Darts).
• Input: De verborgen staten van de laatste 8 lagen van de LLM worden geconcateneerd als input voor de probe.

Belangrijkste Resultaten

1. Recuperatie van Numerieke Waarden:
De probes kunnen de interne representaties van de LLM gebruiken om nauwkeurige punt-schattingen te maken.

• Er is een sterke correlatie (Pearson $R \approx 0.98$) tussen de door de probe voorspelde waarden en de werkelijke statistieken (gemiddelde/median) van de LLM-verdeling.
• De probe presteert aanzienlijk beter dan simpele baselines (zoals het laatste token of het gemiddelde van de reeks).
• De resultaten suggereren dat de LLM de "redenering" voor het volgende getal al volledig heeft uitgevoerd in de verborgen staten, voordat er ook maar één token wordt gegenereerd.

2. Recuperatie van Onzekerheid:

• De kwantiel-probes kunnen het IQR (Interquartile Range) en de spreiding van de verdeling nauwkeurig voorspellen.
• De voorspelde betrouwbaarheidsintervallen zijn goed gekalibreerd: de empirische dekking (coverage) komt zeer dicht bij de theoretische verwachting (bijv. een voorspelde 95% interval bevat inderdaad ~95% van de LLM-samples).

3. Efficiëntie en Kosten:

• Inferentietijd: Het gebruik van de probe is aanzienlijk sneller dan autoregressieve sampling. Het genereren van één sample via autoregressie is ongeveer 47 keer trager dan het uitvoeren van de volledige inferentie-pijplijn met de probe (inclusief het halen van de hidden states).
• De probe kan de nauwkeurigheid van 20-25 LLM-samples benaderen met slechts één forward pass.

4. Generalisatie:

• Contextlengte: Probes getraind op een bepaald bereik van contextlengten generaliseren redelijk naar onbekende lengtes, hoewel training op een breder bereik de robuustheid verbetert.
• Real-world Data: Probes getraind op synthetische data tonen beperkte maar veelbelovende generalisatie naar real-world datasets. Probes getraind op real-world data presteren goed binnen hun domein, maar er is sprake van een daling in prestaties bij grote schaalverschillen (out-of-distribution).

Bijdragen en Significantie

1. Fundamenteel Inzicht: Het paper levert bewijs dat LLMs uitgebreide numerieke informatie en onzekerheidssignalen coderen in hun interne verborgen staten, vóór het autoregressieve decoderingsproces. Dit daagt het idee uit dat autoregressie noodzakelijk is om numerieke uitkomsten te extraheren.
2. Efficiënte Alternatieven: De auteurs introduceren een lichtgewicht, single-pass methode om zowel punt-schattingen als onzekerheidsverdelingen te verkrijgen. Dit opent de deur voor het toepassen van LLMs in regressietaken waar computerefficiëntie en onzekerheidsschatting cruciaal zijn (bijv. veiligheidskritische systemen, real-time controle).
3. Nieuwe Architectuur: De magnitude-factorised probe is een innovatieve aanpak om het probleem van schaalvariatie in numerieke regressie op te lossen, wat een verbetering biedt ten opzichte van bestaande lineaire probes.
4. Praktische Toepasbaarheid: De resultaten tonen aan dat het mogelijk is om de "planning" van een LLM voor numerieke taken te "lezen" zonder de kosten van herhaaldelijk genereren, wat een nieuwe richting opent voor mechanische interpretatie en efficiëntere inferentie.

Conclusie:
Dit werk toont aan dat de "wiskundige redenering" van LLMs al aanwezig is in hun embeddings. Door slimme probes te gebruiken, kunnen we deze informatie efficiënt blootleggen, waardoor we de dure autoregressieve cyclus kunnen omzeilen voor numerieke voorspellingen en onzekerheidskwantificering.