Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

Technische Samenvatting: Functionele-Niveau Kwantificatie van Onzekerheid voor Gecalibreerde Fine-Tuning op LLM's

Probleemstelling
Nauwkeurige kwantificatie van onzekerheid (UQ) is cruciaal voor de betrouwbare implementatie van Large Language Models (LLM's), met name in veiligheidskritieke domeinen zoals gezondheidszorg en wetenschappelijk redeneren. Gefinetuned LLM's vertonen echter vaak oververtrouwen wanneer ze worden aangepast aan taken met beperkte data, waardoor ze de werkelijke geloofwaardigheid van hun output niet weerspiegelen. Bestaande UQ-methoden voor Parameter-Efficiënte Fine-Tuning (PEFT), zoals die gebaseerd op Low-Rank Adaptation (LoRA), zijn grotendeels post hoc. Ze schatten onzekerheid nadat het model is getraind door de parameterruimte te analyseren (bijvoorbeeld via ensembles of Bayesiaanse benaderingen) of door prompts te verstoren. Deze benaderingen slagen er vaak niet in de oorzaak van oververtrouwen aan te pakken: het onvermogen van het model om tijdens het aanpassingsproces zelf onzekerheid veroorzaakt door fine-tuning en generalisatiekloven vast te leggen. Bovendien introduceren methoden die vertrouwen op meerdere forward passes of uitgebreide bemonstering aanzienlijke rekenkundige overhead, wat de schaalbaarheid beperkt.

Methodologie: UQ4CT
De auteurs stellen Functionele-Niveau Kwantificatie van Onzekerheid voor Gecalibreerde Fine-Tuning (UQ4CT) voor, een raamwerk dat de focus verschuift van parameterruimte naar functionele-ruimte onzekerheid. In plaats van geleerde parameters na training te analyseren, calibreren UQ4CT de onzekerheid over de functionele ruimte die wordt veroorzaakt door prompt-afhankelijke mengsels van LoRA-experts tijdens het fine-tuning-proces.

De methodologie bestaat uit drie kerncomponenten:

1. Hiërarchische Functionele Decompositie via MixLoRA:
   Om de combinatorische explosie van het modelleren van alle mogelijke expertcombinaties over lagen heen te vermijden, hanteert UQ4CT een hiërarchische decompositie. Het maakt gebruik van een Mixture-of-Experts (MoE)-architectuur waarbij, op elke laag, een spaarzame router dynamisch een subset van LoRA-experts selecteert op basis van de input verborgen toestand. Dit bouwt een flexibele functionele ruimte op waarin de output van het model een gewogen som is van diverse basisfuncties (LoRA-experts) die afhankelijk zijn van de input. De functionele ruimte wordt weergegeven als een recursieve compositie van laag-voor-laag mengsels, waardoor het model computationele paden adaptief kan samenstellen.

2. Functionele-Niveau Vertrouwen (FLC) Schatting:
   Het model kwantificeert onzekerheid door de concentratie van de expertselectie te analyseren. Specifiek wordt het Functionele-Niveau Vertrouwen (FLC) afgeleid uit de routeringsgewichten van de top-$k$ experts die door de router op elke laag worden geselecteerd. Een hoge concentratie in de routeringsgewichten (d.w.z. de router is zeker van het selecteren van specifieke experts) komt overeen met een hoog functioneel vertrouwen, terwijl verspreide gewichten wijzen op hogere onzekerheid. Dit biedt een principiële, interne signaal van de zekerheid van het model zonder dat meerdere bemonsteringspasses nodig zijn.

3. Calibratieverlies ($L_{cal}$):
   De kerninnovatie is een nieuw calibratieverlies dat het functionele-niveau vertrouwen tijdens training afstemt op voorspellende juistheid. Het verlies wordt gedefinieerd als:
   $$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$$
   waarbij $1\{\cdot\}$ een indicatorfunctie is voor correcte voorspellingen. Dit doel dwingt de FLC om de waarschijnlijkheid te benaderen dat de voorspelling van het model correct is. Door dit verlies te minimaliseren naast de standaard cross-entropy en load-balancing verliezen, leert het model om alleen hoog vertrouwen toe te wijzen wanneer het waarschijnlijk correct is, en hogere onzekerheid wanneer dat niet het geval is. Dit moedigt de router aan om alternatieve expertmengsels te verkennen wanneer de onzekerheid hoog is, waardoor de generalisatie wordt verbeterd.

Belangrijkste Bijdragen

• Functionele-Niveau Perspectief: Het paper introduceert een verschuiving van parameterruimte naar functionele-ruimte kwantificatie van onzekerheid, waarbij het vertrouwen van het model expliciet wordt gekalibreerd over de input-output-mappingen die het tijdens fine-tuning kan realiseren.
• Calibratiebewuste Fine-Tuning: Er wordt een nieuw calibratieverlies voorgesteld dat prompt-afhankelijke LoRA-mengselgewichten dynamisch afstemt op voorspellende juistheid, waardoor oververtrouwen wordt verminderd zonder de nauwkeurigheid te verslechteren.
• Efficiënte Architectuur: De methode maakt gebruik van een hiërarchische MixLoRA-architectuur die diverse functionele relaties vastlegt met minimale rekenkundige overhead, en vereist slechts één forward pass tijdens inferentie.
• Theoretische Garantie: De auteurs leveren een bewijs dat het optimaliseren van het calibratierisico over de data-distributie leidt tot een optimale FLC die overeenkomt met de ware waarschijnlijkheid van voorspellende juistheid.

Experimentele Resultaten
De auteurs hebben UQ4CT geëvalueerd op vier meerkeuze benchmarks (ARC-E, ARC-C, OBQA, ClimateQA) en twee open-ended generatieve QA-taken (TruthfulQA, TriviaQA) met behulp van Llama-3.1-8B en Mistral-7B.

• In-Distributie Prestaties: UQ4CT verlaagde de Expected Calibration Error (ECE) met meer dan 25% over de benchmarks in vergelijking met sterke baselines (waaronder LoRA, MC Dropout, Deep Ensembles en Bayesiaanse LoRA), terwijl het een concurrerende of superieure nauwkeurigheid behield. Opmerkelijk is dat UQ4CT deze resultaten behaalde met één forward pass, terwijl methoden zoals BLoB (N=10) tien passes vereisten maar toch een hogere ECE opleverden.
• Distributieverandering: Onder distributieveranderingen (getest op MMLU-subtaken in CS, Engineering, Recht en Gezondheid) behield UQ4CT superieure kalibratie en concurrerende nauwkeurigheid, wat een verbeterde robuustheid aantoont tegenover out-of-distribution input in vergelijking met baselines.
• Generatieve Taken: Bij open-ended generatie behaalde UQ4CT de laagste ECE op TruthfulQA en TriviaQA, wat suggereert dat het perspectief van functionele-ruimte kalibratie verder reikt dan classificatie-achtige settings.
• Efficiëntie: Geheugenanalyses geven aan dat het piek GPU-geheugengebruik van UQ4CT slechts ongeveer 7,5% hoger is dan standaard LoRA, waardoor het rekenkundig haalbaar is voor grootschalige implementatie.

Betekenis en Claims
Het paper stelt dat UQ4CT een fundamentele beperking in huidige fine-tuning-praktijken aanpakt: het gebrek aan bewustzijn van onzekerheid tijdens het aanpassingsproces. Door onzekerheidscalibratie direct te integreren in het fine-tuning-doel via functionele-niveau signalen, zorgt de methode ervoor dat het vertrouwen van het model beter zijn werkelijke voorspellende betrouwbaarheid weerspiegelt. De auteurs positioneren dit werk als een fundamentele stap richting betrouwbare gefinetunde LLM's, met name in data-schaarse omgevingen waar oververtrouwen aanzienlijke risico's met zich meebrengt. Ze merken bescheiden op dat hoewel de aanpak uitbreidt naar open-ended generatie met F1-gebaseerde proxies, schalen naar langere redeneringen of meerstaps-generatie een richting blijft voor toekomstig werk.