Technische Samenvatting: SAFE-KD

Probleemstelling

Moderne vision-backbones bereiken een hoge nauwkeurigheid, maar brengen aanzienlijke inferentiekosten met zich mee, wat hun inzet op edge-apparaten en latency-gevoelige systemen beperkt. Early-exit netwerken pakken dit aan door tussenliggende classificators toe te voegen, waardoor "gemakkelijke" inputs de berekening voortijdig kunnen stoppen. De huidige implementatie van deze techniek leunt echter op heuristische stopregels (bijv. vertrouwen of entropiedrempels) die geen formele garanties bieden. Deze heuristieken zijn vaak onbetrouwbaar door miskalibratie van het model en verslechteren snel onder distributieverschuivingen, waardoor het moeilijk is om nauwkeurigheid op een voorspelbare, veilige manier af te wegen tegen latentie. Er is een kritieke behoefte aan een adaptief inferentiebeleid dat eindige-steekproefgaranties biedt op het risico van vroege exits.

Methodologie

SAFE-KD (Risk-Controlled Early-Exit Distillation) is een universele multi-exit wrapper die is ontworpen om gekalibreerde exits met statistische garanties mogelijk te maken. De methode integreert drie kerncomponenten:

1. Universele Multi-Exit Architectuur

SAFE-KD koppelt lichtgewicht exit-heads aan tussenliggende dieptes van moderne backbones (CNN's en Vision Transformers) met minimale architecturale aannames:

• Plaatsing: Voor CNN's (bijv. ResNet, ConvNeXt) worden exits geplaatst na geselecteerde residuele blokgroepen. Voor ViT's (bijv. ViT-S, Swin-T) worden exits gekoppeld na geselecteerde transformer-blokken, waarbij CLS- of mean token embeddings worden gepoold.
• Head Design: Een kleine predictie-head (Global Average Pooling of Token Pooling gevolgd door optionele MLP/LayerNorm en een lineaire classifier) zorgt voor minimale overhead terwijl er goed gevormde logits worden geproduerd.
• Signaal: Het primaire vertrouwenssignaal is de Maximum Softmax Probability (MSP).

2. Hiërarchische Distillatie

Om de zwakte en ruis van tussenliggende exits aan te pakken, gebruikt SAFE-KD een hiërarchische distillatiestrategie met een Exponential Moving Average (EMA) teacher:

• Decoupled Knowledge Distillation (DKD): Elke exit wordt door de teacher gesuperviseerd met behulp van DKD, wat de componenten van de doelklasse en de niet-doelklasse scheidt in de distillatieverliesfunctie om "dark knowledge" beter te benutten.
• Deep-to-Shallow Consistentie: Een consistentieverlies brengt vroege exits in lijn met de finale exit, wat de voorspellingsstabiliteit over de netwerkdiepte heen waarborgt.
• Doelstelling: De totale loss combineert standaard cross-entropy, DKD-loss tegen de teacher, en consistentieverlies tegen de finale exit, gewogen door hyperparameters $\alpha$ en $\beta$.

3. Risicogestuurde Exiting via Conformal Risk Control (CRC)

In plaats van heuristische drempels gebruikt SAFE-KD Conformal Risk Control om gekalibreerde drempels voor vroege exits te bepalen op een apart gehouden kalibratieset ($D_{cal}$):

• Selectieve Risicodefinitie: Vroege exit wordt geframed als selectieve classificatie. De onzekerheidsscore wordt gedefinieerd als $r_j(x) = 1 - \max_c p_j(c|x)$.
• Kalibratie: Voor een door de gebruiker gespecificeerd risiconiveau $\delta$, berekent CRC een drempelwaarde $\hat{\tau}_j(\delta)$ voor elke exit $j$. Deze drempel garandeert dat het selectieve misclassificatierisico (de kans op fouten bij samples die bij exit $j$ worden geaccepteerd) begrensd is door $\delta + O(1/|D_{cal}|)$ onder de aanname van uitwisselbaarheid (exchangeability).
• Inferentiebeleid: Tijdens testtijd stopt het model bij de vroegste head $j$ waar de onzekerheidsscore $r_j(x) \leq \hat{\tau}_j(\delta)$ is. Als geen enkele vroege exit aan de voorwaarde voldoet, gaat het model door naar de finale exit.

Belangrijkste Bijdragen

1. Universele Multi-Exit Wrapper: Een modulair ontwerp dat tussenliggende exits toevoegt aan CNN's en ViT's met minimale overhead, waarbij de kenmerken van de backbone behouden blijven.
2. Hiërarchische Distillatie: Een trainingsbenadering die gebruikmaakt van DKD en deep-to-shallow consistentie om de nauwkeurigheid en stabiliteit van tussenliggende voorspellingen te verbeteren.
3. Risicogestuurde Exiting: Een kalibratieprocedure die vertrouwensscores omzet in door de gebruiker instelbare, risicogebonden exit-drempels, waardoor broze heuristieken worden vervangen door een statistisch valide test.
4. Empirische Validatie: Getoond over meerdere datasets en architecturen, waarbij verbeterde nauwkeurigheid-rekenkracht-afwegingen, sterkere kalibratie en robuustheid onder corruptie worden aangetoond.

Experimentele Resultaten

De auteurs evalueerden SAFE-KD op datasets zoals CIFAR-10/100, STL-10, Oxford-IIIT Pets, Flowers102 en FGVC Aircraft, met backbones zoals ResNet-50, MobileNetV3-S, EfficientNet-B0, ConvNeXt-T, ViT-S en Swin-T.

• Nauwkeurigheid en Efficiëntie: SAFE-KD evenaart of overtreft de finale-exit nauwkeurigheid van standaard DKD en presteert aanzienlijk beter dan MultiExit-baselines. Op ViT-S met CIFAR-100 behaalde het een $+0,80\%$ verbetering in nauwkeurigheid ten opzichte van ERM.
• Pareto-Dominantie: SAFE-KD bereikt een strikt dominante Pareto-frontier vergeleken met vaste vertrouwen- en entropiegebaseerde gating. Bijvoorbeeld, op ResNet-50 met een doelrisico van $\delta=0,05$, verminderde SAFE-KD de computationele diepte met 41% terwijl een relatieve nauwkeurigheid van 94,1% werd behouden.
• Risicogeldigheid: Empirische selectieve risicocurven volgen consistent de $y=x$ diagonaal of blijven er iets onder, wat de eindige-steekproef validiteit bevestigt. In tegen staat hiermee vaak de heuristische baselines (bijv. vast vertrouwen van 0,95), die de risicoconstricties overtraden, wat resulteerde in geobserveerde risico's (bijv. 12% op Flowers102) die het doel ver veel overschreden.
• Robuustheid: Onder CIFAR-10-C corruptie degradeerde SAFE-KD minder sterk dan standaard modellen. De gekalibreerde drempels wezen gecorreleerde samples op natuurlijke wijze af bij vroege exits, waardoor het model gedwongen werd om diepere, robuustere lagen te gebruiken.
• Ablatie: Het verwijderen van DKD veroorzaakte een significante daling in nauwkeurigheid (-2,1%) en dwong tot conservatievere drempels (verhoging van diepte naar 0,85). Het verwijderen van consistentie verhoogde de risicovariantie. De volledige methode bood de beste balans tussen lage diepte en hoge nauwkeurigheid.

Betekenis en Claims

Het artikel stelt dat SAFE-KD adaptieve inferentie transformeert van een heuristisch optimalisatieprobleem naar een betrouwbaar, risicogestuurd systeem. Door het foutenpercentage van vroege exits wiskundig te begrenzen, biedt het een betrouwbare interface tussen modeloutputs en systeemvereisten. Dit stelt systeemontwerpers in staat om de risicoparameter $\delta$ te behandelen als een harde restrictie in plaats van een afstelbare hyperparameter, waardoor early-exit netwerken levensvatbaar worden voor de inzet in veiligheidskritische omgevingen waar voorspelbare latentie en foutmarges essentieel zijn.

De auteurs merken beperkingen op, specifelijk de afhankelijkheid van een kalibratieset (kleine sets leiden tot conservatievere drempels) en de afhankelijkheid van de aanname van uitwisselbaarheid, die kan breken onder zware distributieverschuivingen. Binnen deze grenzen biedt SAFE-KD echter een principieel kader voor efficiënte en veilige inferentie.