Technische Samenvatting: Voorbij Post-hoc Temperatuur-schaling: Bilevel Optimalisatie voor LLM Kalibratie

Probleemstelling

Grote taalmodellen (LLM's), met name die zijn uitgelijnd via Reinforcement Learning from Human Feedback (RLHF) of Direct Preference Optimization (DPO), vertonen frequent een slechte kalibratie, wat zich manifesteert als overmatige zelfverzekerdheid (overconfidence). Hoewel zelfverzekerdheid vaak wordt gebruikt als een proxy voor betrouwbaarheid, kunnen slecht gekalibreerde modellen in scenario's met hoge inzet (bijv. de gezondheidszorg, de juridische sector) leiden tot onterecht vertrouwen en slechte besluitvorming.

Bestaande oplossingen vertrouiken primair op post-hoc Temperatuur-schaling (TS), waarbij een enkele scalaire temperatuur wordt aangepast op een validatieset om de logits na de training te herschalen. Echter, de paper identificeert een kritieke beperking: TS is inherent dataset-afhankelijk. Een temperatuur die geoptimaliseerd is voor één domein, faalt vaak bij andere domeinen, wat het moeilijk maakt om een "universele" temperatuur te vinden die goed kalibreert over diverse taken heen. Bov�elijk laten post-hoc methoden de modelparameters vaststaan, waardoor de kans gemist wordt om een model te leren dat intrinsiek beter geschikt is voor kalibratie.

Methodologie: CALM (Calibration for Large Models via Bilevel Optimization)

De auteurs stellen CALM voor, een raamwerk tijdens de training dat kalibratie herformuleert als een bilevel optimalisatieprobleem. In plaats van een post-hoc parameter aan te passen, optimaliseert CALM gezamenlijk de modelparameters en "temperatuur-achtige" verlies-hyperparameters om de voorspellende entropie te maximalen terwijl de taakprestaties behouden blijven.

1. Bilevel Formulering

Het probleem is als volgt gestructureerd:

• Onderliggend Niveau (Taakoptimalisatie): Traint de modelparameters $\theta$ om een standaard parametrische loss (cross-entropy) op de trainingsset $S_T$ te minimaliseren, gegeven een set hyperparameters $\alpha$.
  $$\theta(\alpha) \in \arg\min_{\theta'} \mathcal{L}_{tr}(\theta'; \alpha)$$
• Bovenliggend Niveau (Kalibratieoptimalisatie): Optimaliseert de hyperparameters $\alpha$ om de entropie van de voorspellende distributie op een gehouden validatieset $S_V$ te maximaliseren, onder de voorwaarde dat $\theta$ optimaal blijft voor de onderliggende taak.
  $$\min_{\alpha} \mathcal{L}_{ent}(\alpha, \theta(\alpha))$$

De hyperparameters $\alpha$ generaliseren het concept van temperatuur-schaling. In plaats van een enkele scalaire $\tau$, gebruikt CALM een hogere-dimensie parameter $\alpha = (l, \tau)$, waarbij $l \in \mathbb{R}^K$ additieve logit-shifts introduceert en $\tau \in \mathbb{R}^K_{>0}$ klasse-afhankelijke temperaturen biedt. Dit maakt aanpassingen per vocabulaire mogelijk in plaats van een globale herschaling.

2. Objectieve Functies

• Bovenliggend Objectief: Maximaliseert de Shannon-entropie $H(f_{\alpha, \theta}(\cdot | x))$ om te voorkomen dat voorspellingen te geconcentreerd (overconfident) zijn. Dit biedt een controleerbare richting om het model van overconfidence naar onderconfidence te bewegen, wat potentieel een punt vindt waar zelfverzekerdheid overeenkomt met nauwkeurigheid.
• Onderliggend Beperking: Zorgt ervoor dat het model zijn onderscheidend vermogen behoudt door cross-entropy op de trainingsdata te minimaliseren.

3. Efficiënte First-Order Approximatie

Exacte bilevel optimalisatie vereist het berekenen van hypergradiënten die betrokken zijn bij tweede-orde afgeleiden (Hessian inverses), wat computationeel onhaalbaar is voor LLM's met miljarden parameters. Om dit aan te pakken, adopteert CALM een first-order approximatie geïnspireerd door BOME (Bilevel Optimization Made Easy):

• Het herformuleert het probleem als een geconstrueerde optimalisatie over de gezamenlijke variabele $u = (\alpha, \theta)$.
• Het benadert de onderliggende waarde-functie door $T$ innerlijke gradiëntstappen uit te voeren vanaf de huidige $\theta$ om een surrogaat $\tilde{\theta}$ te genereren.
• Het construeert een "plug-in residual" beperking en berekent een gezamenlijke update-richting met alleen first-order afgeleiden, waardoor expliciete Hessian-inversie wordt vermeden.

Belangrijkste Bijdragen

1. Perspectief tijdens de Training: De paper verschuift het kalibratieparadigma van post-hoc dataset-specifieke aanpassingen naar optimalisatie tijdens de training, met als doel modellen te leren die algemene kalibratie-aanpassingen over datasets heen toelaten.
2. Bilevel Loss Design: Het formuleert kalibratie-georiënteerd loss-ontwerp als een bilevel probleem waarbij het bovenliggende niveau temperatuur-achtige hyperparameters optimaliseert met een entropie-gebaseerd objectief, en het onderliggende niveau de training van het model verzorgt.
3. Entropie-maximalisatie: De auteurs stellen een entropie-maximalisatie bovenliggend objectief voor om overmatige zelfverzekerdheid in uitgelijnde LLM's te mitigeren zonder dat er extra harde labels of hulp-heads nodig zijn.
4. Schaalbaar Raamwerk: Ze introduceren een efficiënt first-order bilevel optimalisatieschema (CALM) dat dure tweede-orde berekeningen vermijdt, waardoor de aanpak praktisch bruikbaar is voor grootschalige LLM fine-tuning.

Experimentele Resultaten

De methode werd geëvalueerd op vier uitgelijnde LLM's (Llama-3.1-Tulu-8B, Vicuna-7B, OLMo-2-7B, Mistral-7B) over multiple-choice QA (MCQA) en open-ended generatieve QA.

Out-of-Domain (OOD) Generalisatie

• Superieure Kalibratie: CALM behaalde de beste confidence Expected Calibration Error (conf-ECE) op drie van de vier modellen in de OOD-setting, wat de fout aanzienlijk verminderde vergeleken met de baselines. Op Llama-3.1 verminderde het bijvoorbeeld de conf-ECE van 0.1784 naar 0.1050.
• Behoud van Nauwkeurigheid: In tegenstelling tot sommige ablaties (bijv. "Iterate" of "Regularization" zonder de bilevel-beperking) die leden onder catastrofale nauwkeurigheidsinstorting, behield CALM een hoge taaknauwkeurigheid terwijl de kalibratie verbeterde.
• Robuustheid: In open-ended generatieve QA (cross-domain transfer tussen PopQA en TriviaQA) behaalde CALM de laagste gemiddelde Semantic-ECE (0.0671), waarmee het zowel de ongekalibreerde baseline als sterke training-gebaseerde baselines zoals Calibration-Aware Fine-Tuning (CFT) versloeg.

In-Domain (ID) Prestaties

• In de ID-setting, waar de trainings- en evaluatiedistributies overeenkomen, presteerde Temperature Scaling het best op ECE (zoals verwacht voor post-hoc methoden). Echter, CALM was de enige training-gebaseerde methode die consistent zowel de nauwkeurigheid als de kalibratie verbeterde over alle vier de modellen vergeleken met de baseline.

Ablatie-studies

• Bilevel Structuur: Het verwijderen van de bilevel-beperking (Regularization) of de koppelingsterm (Iterate) resulteerde in ofwel slechtere kalibratie of significante nauwkeurigheidsverlies, wat aantoont dat de koppeling tussen taakprestatie en kalibratie essentieel is.
• Parametrisatie: Het gebruik van per-vocabulaire vectorparameters ($d_y, l_y$) was veel superieur aan een samengevallen scalaire parameter (geleerde TS), wat het belang benadrukt van fijnmazige aanpassingen om asymmetrische miskalibratie over antwoordopties heen te behandelen.

Betekenis en Claims

De paper claimt dat CALM een robuuste, tijdens de training gerichte oplossing biedt voor het probleem van overmatige zelfverzekerdheid in uitgelijnde LLM's. De primaire betekenis ligt in:

• Generalisatie: In tegenstelling tot post-hoc TS, leert CALM een kalibratiemechanisme dat generaliseert over domeinen, waarmee het de "dataset-afhankelijkheid" van huidige methoden aanpakt.
• Trade-off Management: Het bereikt een superieur evenwicht tussen nauwkeurigheid en kalibratie, waarbij de nut-instorting (utility collapse) wordt vermeden die vaak optreedt bij het proberen te verbeteren van kalibratie via eenvoudige regularisatie.
• Praktische Bruikbaarheid: Door gebruik te maken van een first-order approximatie, maakt het bilevel optimalisatie haalbaar voor moderne LLM's, wat een pad biedt naar "bij ontwerp gekalibreerde" modellen in plaats van "gekalibreerd door post-processing".

De auteurs concluderen dat hun methode bijzonder voordelig is in out-of-domain settings, waar alternatieve methoden vaak instabiel worden of aanzienlijk verlies aan nut lijden, wat het een levensvatbare aanpak maakt voor het inzetten van betrouwbare LLM's in diverse, real-world scenario's.