Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration

Technische Samenvatting: Diffusie-geïnspireerde Reconfiguratie van Transformers voor Uncertainty Calibration

Probleemstelling
Vooraf getrainde transformers zijn het fundament geworden van grote schaal foundation modellen binnen de domeinen visie, taal en spraak. Deze modellen missen echter doorgaans een principiële mechanisme voor de propagatie van onzekerheid door hun stacks van feature-transformaties. Bijgevolg wijzen ze vaak hoge zekerheid toe aan onjuiste voorspellingen, wat veiligheidsrisico's oplevert in toepassingen met hoge stakes. Hoewel er technieken voor uncertainty calibration (UC) bestaan voor kleinere netwerken, vertrouwen deze vaak op het benaderen van model-posteriors via variational inference of posterior sampling, wat voor modellen op grote schaal computationeel onhaalbaar en onnauwkeurig wordt.

Bestaande benaderingen die proberen attention-blokken te herparametriseren als Gaussische processen (GPs) stuiten op een kritieke beperking: ze behandelen blokken onafhankelijk van elkaar. Dit houdt geen rekening met de correlaties tussen feature-transformaties die tijdens het vooraf trainen tot stand zijn gekomen. Als gevolg daarvan dwingen deze methoden vaak een trade-off af waarbij een verbeterde uncertainty calibration ten koste gaat van een aanzienlijk verslechterde voorspellende prestatie. Bovendien verstoort het afwisselen van probabilistische attention-blokken met punt-geschatte componenten (zoals MLP's en Layer Normalization) de stroom van onzekerheidspropagatie, waardoor het tot stand brengen van een goed gedefinieerd stochastisch pad wordt verhinderd.

Methodologie
De auteurs stellen DIRECTOR (Diffusion-Inspired REConfiguration of TransfORmers) voor, een raamwerk dat vooraf getrainde transformers reconfigureert tot een unificerend diffusieproces om een principiële onzekerheidspropagatie mogelijk te maken.

1. Structurele Reconfiguratie als Kanspad:
   De methode begint met het herindelen van de transformer-architectuur zodat elke transformatie-blok eindigt met een Multi-Head Self-Attention (MHSA) module. Geput uit de gevestigde connectie tussen MHSA en Gaussische Processen, wordt de output van elke MHSA-blok geïnterpreteerd als een Gaussische verdeling over intermediate features.
   Cruciaal worden de deterministische segmenten (MLP en Layer Normalization) die doorgaans de onzekerheidsstroom onderbreken, niet geherinterpreteerd als waargenomen functies van onbekende priors, maar als extra parametrisaties van de geneuraleerde Gaussische overgang die wordt veroorzaakt door de door GP herparametriseerde MHSA. Deze herstructurering zorgt ervoor dat de volledige reeks feature-transformaties een coherent stochastisch proces $\{X_t\}_{t=0}^T$ vormt met Gaussische overgangen, waardoor de architectuur van de transformer effectief wordt afgestemd op een diffusieproces in omgekeerde tijd.

2. Distillatie tot een Unificerend Diffusiemodel:
   In plaats van de overgang van elke blok onafhankelijk te behandelen, distilleert DIRECTOR de waargenomen reeks feature-transformaties tot een unificerend spatiotemporale overgangskern. Dit wordt bereikt door een op diffusie gebaseerd model (geïmplementeerd als een DiT met één blok) te trainen om het proces in omgekeerde tijd te leren dat de finale feature-verdeling terugmapt naar de invoerdata-verdeling.
   Het trainingsdoel combineert twee componenten:

   • Onzekerheidsbewuste Loss ($L_1$): Een Kullback-Leibler (KL) divergentie-term die de geleerde overgangskern afstemt op de waargenomen overgangen van de bevroren vooraf getrainde backbone. Dit zorgt ervoor dat het model de inherente correlaties tussen attention-blokken vastlegt.
   • Prestatiebewuste Loss ($L_2$): Een standaard voorspellingsloss (bijvoorbeeld cross-entropy) toegepast op steekproeven getrokken uit het geleerde diffusieproces. Dit zorgt ervoor dat het gereconfigureerde model de voorspellende nauwkeurigheid behoudt of verbetert.

Belangrijkste Bijdragen

• Probabilistische Reconfiguratie: Het artikel introduceert een nieuwe structurele reconfiguratie die een vooraf getrainde transformer transformeert tot een reeks geneuraleerde Gaussische overgangen, waarbij het model wordt gegeneraliseerd tot een diffusieproces geparametriseerd door een unificerende overgangskern.
• Unificerende Onzekerheidspropagatie: Door het feature-pad van de transformer te distilleren tot een enkel diffusiemodel, legt de methode de tijdens het vooraf trainen tot stand gekomen correlaties tussen blokken vast, waardoor een principiële onzekerheidspropagatie mogelijk wordt zonder de onuitvoerbaarheid van het berekenen van volledige modelposteriors.
• Efficiënte Post-hoc Calibratie: De aanpak fungeert als een post-hoc embedding van probabilistisch redeneren. Het unificerende overgangsmodel is aanzienlijk kleiner dan de oorspronkelijke backbone (bijvoorbeeld 2,7 miljoen parameters versus 6,24 miljoen voor ViT), wat een verbeterde geheugenefficiëntie biedt terwijl de voorspellende capaciteiten van het oorspronkelijke model behouden blijven.

Experimentele Resultaten
De auteurs hebben DIRECTOR geëvalueerd op visie-benchmarks (CIFAR-10, CIFAR-100, CIFAR-10-C) en taal-benchmarks (IMDB, CoLA, CoLA-OOD).

• Calibratie en Nauwkeurigheid: DIRECTOR presteert consequent beter dan bestaande onzekerheidsbewuste transformers (inclusief KEP-SVGP varianten) en standaard baselines (Temperature Scaling, Monte Carlo Dropout, Deep Ensembles). Op CIFAR-10 bereikte DIRECTOR een hogere nauwkeurigheid (84,74% versus 82,68% voor KEP-7/7) en een lagere Expected Calibration Error (5,18% versus 6,95%). Vergelijkbare verbeteringen werden waargenomen op alle datasets, wat aantoont dat de methode de typische trade-off tussen calibratie en prestatie oplost.
• Robuustheid: Onder distributieveranderingen (CIFAR-10-C en CoLA-OOD) behield DIRECTOR concurrerende voorspellende nauwkeurigheid terwijl het de calibratiemetrics aanzienlijk verbeterde in vergelijking met vooraf getrainde modellen.
• OOD-detectie: De methode toonde superieure prestaties in taken voor Out-of-Distribution detectie, met hogere AUROC- en AUPR-scores.
• Schaalbaarheid: Experimenten op grotere modellen (ViT-B-16 en Qwen-2.5) bevestigden dat DIRECTOR effectief schaalt en betrouwbare calibratie levert met minder parameters dan de oorspronkelijke modellen.

Betekenis en Claims
Het artikel claimt dat DIRECTOR een schaalbaar en principiële paradigma biedt voor uncertainty calibration in grote vooraf getrainde modellen. Door de stack van feature-transformaties van de transformer te herinterpreteren als een diffusieproces, embedt de methode expressieve onzekerheidsbewuste structuren direct in de interne architectuur van het model. De auteurs stellen dat deze aanpak een betrouwbare implementatie mogelijk maakt in risicosensitieve toepassingen door ervoor te zorgen dat zekerheidsscores trouw de empirische nauwkeurigheid weerspiegelen, terwijl tegelijkertijd de voorspellende prestaties worden behouden of verbeterd en de geheugenefficiëntie wordt verhoogd. Het werk opent een nieuwe richting voor het verbeteren van de betrouwbaarheid van foundation modellen zonder hun kerncapaciteiten op te offeren.