LLM REgression with a Latent Iterative State Head

Probleemstelling

Hoewel Large Language Models (LLMs) uitstekend presteren bij generatieve taken (zoals vertalen of samenvatten) binnen het "tekst-naar-tekst"-paradigma, zijn ze minder geschikt voor voorspellende taken zoals regressie, waarbij het doel is om een continue numerieke waarde (scalar) te voorspellen.

Bestaande methoden voor LLM-regressie hebben drie belangrijke beperkingen:

1. Autoregressieve decoding: Hierbij genereert het model de numerieke waarde als tekst (bijv. "3.5") en wordt deze vervolgens geparsed. Dit is suboptimaal omdat taalmodellen zijn getraind om token-kansverdelingen te minimaliseren, niet numerieke fouten (zoals MSE). Een verschil tussen 0.9 en 1.0 wordt als een volledig ander token behandeld, terwijl numeriek gezien de afstand klein is.
2. Regressie-bewuste inferentie (RAIL): Deze methoden proberen de numerieke fout te minimaliseren door de posterior-verdeling te benaderen via bemonstering. Dit leidt echter tot hoge rekentijd tijdens de inferentie (multi-sample) en lost het fundamentele mismatch-probleem slechts gedeeltelijk op, omdat de posterior nog steeds over getextualiseerde getallen wordt gedefinieerd.
3. Voorspellende koppen (Predictive Heads): Deze methoden omzeilen tekstgeneratie door direct een regressor op de interne representaties van het LLM te plaatsen. Echter, bestaande benaderingen gebruiken vaak simpele, statische pooling-methoden (zoals gemiddelde pooling over tokens). Dit kan leiden tot een "samenvattingsbottleneck" waarbij fijne, voor regressie relevante signalen verloren gaan.

Methodologie: RELISH

De auteurs stellen RELISH (REgression with a Latent Iterative State Head) voor, een nieuwe, lichtgewicht architectuur die de efficiëntie van voorspellende koppen combineert met een robuuster mechanisme voor het extraheren van informatie.

Kernarchitectuur:
In plaats van statische pooling of tekstgeneratie, gebruikt RELISH een iteratief verfijningsmechanisme over de token-level representaties van het LLM.

1. Frozen Backbone: Het onderliggende LLM blijft bevroren (geen fine-tuning van de backbone), wat zorgt voor parameter-efficiëntie.
2. Latente Toestand: Het model onderhoudt een geleerde latente toestand $r^{(0)}$.
3. Iteratieve Refinement: Deze toestand wordt $L$ keer bijgewerkt via cross-attention. De latente toestand fungeert als query, terwijl de token-representaties van het LLM als keys en values dienen. Dit stelt het model in staat om dynamisch en selectief informatie uit de sequentie te halen, in plaats van een statisch gemiddelde te nemen.
   • Formule: $r^{(i)} = \text{Refine}(r^{(i-1)}, H(x))$
4. Linear Regressor: De finale latente toestand $r^{(L)}$ wordt gemapt naar een scalair voorspelling via een lineaire laag: $\hat{y} = w^\top r^{(L)} + b$.
5. Training: De training gebruikt de Huber-loss op genormaliseerde targets, waarbij alleen de parameters van de RELISH-kop (projectie, latente toestand, attention-blokken, regressor) worden geoptimaliseerd.

Belangrijkste Bijdragen

1. Nieuwe Architectuur: Een innovatieve voorspellende kop die statische pooling vervangt door een iteratief, attention-gebaseerd mechanisme voor het samenvatten van token-representaties.
2. Uitgebreide Vergelijking: Een unificatie en empirische evaluatie van de drie grote families van LLM-regressie (autoregressief, regressie-bewust, en voorspellende koppen) over vijf datasets, vier verschillende LLM-backbones (8B tot 32B parameters) en twee trainingsregimes (bevroren vs. LoRA fine-tuning).
3. Parameter-efficiëntie: RELISH vereist slechts ~3,4–3,7 miljoen trainbare parameters, ongeacht de grootte van het LLM-backbone. Dit is slechts 0,01–0,04% extra overhead, wat aanzienlijk lager is dan LoRA-baselines (die 0,26–0,42% vereisen en meegroeien met het model).

Resultaten

RELISH presteert consequent beter dan state-of-the-art baselines op alle geteste configuraties:

• Prestatie: Over vijf datasets (inclusief semantische tekstuele gelijkenis en machine vertaalkwaliteit) en vier LLM-modellen (Llama 3.1, Qwen3, Gemma 3) behaalde RELISH de hoogste correlaties (Pearson en Spearman) en de laagste range-genormaliseerde RMSE (NRMSE).
• Verbetering: Ten opzichte van de tweede beste methode (RAFT), verbeterde RELISH de Pearson-correlatie met 6,7%, Spearman met 7,2% en NRMSE met 19,9%.
• Robuustheid: In tegenstelling tot andere methoden die soms onder de zero-shot prestaties zakken, levert RELISH in bijna alle gevallen een positieve verbetering op.
• Efficiëntie: De methode behoudt de "single-pass" inferentie-efficiëntie van voorspellende koppen, maar vermijdt de hoge rekentijd van regressie-bewuste inferentie (die meerdere samples vereist).

Betekenis en Impact

De paper toont aan dat de beperkingen van LLM-regressie niet noodzakelijk liggen in het gebrek aan numeriek inzicht van het model, maar in hoe die informatie wordt geëxtraheerd.

• Architecturale Superioriteit: De resultaten tonen aan dat iteratieve latent-state refinement superieur is aan statische pooling of tekstgeneratie voor numerieke taken. Het model kan "lezen" wat er in de hidden states zit zonder de ruis van tekstgeneratie.
• Schalbaarheid: Omdat de trainbare footprint van RELISH niet schaalt met de grootte van het LLM (maar wel met de hidden dimension), is het een ideale oplossing voor het toepassen van grote LLMs op regressietaken zonder de enorme kosten van full-parameter fine-tuning of LoRA.
• Toepassingen: De methode is breed toepasbaar op taken zoals beloningmodelling (reward modeling), "LLM-as-a-judge" (het beoordelen van outputkwaliteit met een score), en betrouwbaarheidsschattingen, waar continue scores nodig zijn.

Kortom, RELISH biedt een elegante, efficiënte en krachtige oplossing voor het integreren van regressie in het LLM-ecosysteem, waarbij de sterktes van de pre-trained representaties optimaal worden benut zonder de nadelen van tekstgeneratie of zware fine-tuning.