Technische Samenvatting: KLAS - Het Gebruik van Similariteit om Neuronale Netwerken te Naaien voor Verbeterde Afwegingen tussen Nauwkeurigheid en Efficiëntie

1. Probleemstelling

De implementatie van grote vooraf getrainde vision transformers (ViTs) en convolutionele neurale netwerken (CNN's) staat voor een fundamentele uitdaging: diverse omgevingen variëren van krachtige cloudservers tot middelenbeperkte randapparaten, wat flexibele modelselectie binnen specifieke rekenbudgetten vereist. Hoewel Neural Architecture Search (NAS) probeert nauwkeurigheid en efficiëntie in evenwicht te brengen, is traditionele één-op-één NAS computationeel onbetaalbaar, en zijn één-op-veel of zero-shot NAS-benaderingen vaak beperkt tot één modelontwerpruimte, waardoor ze geen gebruik maken van de rijke voorraad bestaande vooraf getrainde modellen.

Modelstitching biedt een veelbelovend alternatief door blokken van verschillende vooraf getrainde "anker"-modellen te combineren om geïnterpoleerde netwerken te creëren. Echter, bestaande veel-op-veel stitching-benaderingen, zoals SN-Net, vertrouwen op naïeve heuristieken (bijvoorbeeld "naaste stitching" van aangrenzende ankers of vaste blokkoppeling) om stitch-configuraties te selecteren. Deze heuristieken leveren vaak suboptimale afwegingen tussen nauwkeurigheid en efficiëntie op en missen generaliseerbaarheid over diverse modelfamilies. Bovendien is de zoekruimte voor stitch-configuraties enorm ($O(k^2n^2)$ voor $k$ modellen met diepte $n$), waardoor een exhaustieve zoektocht onuitvoerbaar is. Het kernprobleem is het ontbreken van een principiële, geautomatiseerde methode om compatibele stitch-configuraties te identificeren die de prestaties maximaliseren zonder hoge trainingskosten.

2. Methodologie: Het KLAS-kader

De auteurs stellen KLAS (KL-divergentie gebaseerde Ankerstitching) voor, een kader dat stitch-selectie automatiseert door gebruik te maken van de similariteit tussen intermediaire representaties van vooraf getrainde modellen.

Kernconcept: KL-divergentie als Similariteitsmaatstaf

KLAS stelt dat succesvolle stitching afhangt van zowel representatieve similariteit (geometrische uitlijning van activaties) als functionele similariteit (behoud van taakrelevante informatie en beslissingsgrenzen). De auteurs betogen dat Kullback-Leibler (KL)-divergentie uniek geschikt is om beide vast te leggen:

• Representatieve Uitlijning: Het kwantificeert de statistische afstand tussen activatiedistributies, wat aangeeft of een lichtgewicht transformatie bronactivaties kan afbeelden op doelactivaties.
• Functionele Uitlijning: Het beoordeelt hoe goed de output van een bronblok de beslissingsgrenzen van het doelmodel behoudt, zodat het gestikte netwerk de prestaties van downstream-taken behoudt met minimale finetuning.

In tegenstelling tot maatstaven zoals Centered Kernel Alignment (CKA), Mean Squared Error (MSE) of Cross-Entropy (CE), die vaak falen bij het voorspellen van stitching-succes, meet KL-divergentie direct de compatibiliteit van kansdistributies gegenereerd door lineaire probes.

Het KLAS-algoritme

Het kader werkt in twee hoofdfasen:

1. Ankerselectie:

   • KLAS berekent de KL-divergentie tussen de softmax-kansdistributies van de laatste blokken van kandidaat vooraf getrainde modellen.
   • Modellen met een lage KL-divergentie van het laatste blok worden beschouwd als compatibele ankers, omdat ze vergelijkbare beslissingsgrenzen en klassen-vertrouwensdistributies delen. Dit vervangt de heuristische "naaste stitching" uit eerder werk.

2. Blokselectie:

   • Zodra ankers zijn geselecteerd, evalueert KLAS blokkoppelingen met behulp van een stitch-score $\Gamma(i, j)$ gedefinieerd als:
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     Waarbij:
     • $\Theta(P^f_i, P^g_j)$ de cross-anker activatie-afstand is (KL-divergentie tussen bronblok $i$ en doelblok $j$).
     • $\Theta(P^g_j, P^g_{j+1})$ de intra-anker blokcapaciteit is (KL-divergentie tussen opeenvolgende blokken in het doelanker).
   • Deze score balanceert de transformatiebelasting op de hypothetische stitch-laag tegen het vermogen van het doelblok om distributieverschillen te absorberen.
   • ProbeNet: Om deze scores efficiënt te berekenen zonder volledige stitch-lagen te trainen, maakt KLAS gebruik van een unificeerde architectuur genaamd ProbeNet. Het koppelt lineaire classifier-probes aan elk blok van de ankermodellen en traint deze gezamenlijk in één forward-backward pass. Dit veroorzaakt een verwaarloosbare eenmalige kosten (bijvoorbeeld 0,25 GPU-dagen voor Swin-B) in vergelijking met het trainen van individuele probes.

3. Configuratie-selectie:

   • Kandidaat stitch-configuraties worden gefilterd op basis van een instelbare drempelwaarde $\tau$ ten opzichte van de minimale score binnen FLOPs-buckets. Dit zorgt voor een dichte Pareto-front van implementeerbare modellen over het spectrum van nauwkeurigheid-efficiëntie.

3. Belangrijkste Bijdragen

1. Identificatie van Heuristische Beperkingen: Het artikel toont aan dat naïeve heuristieken (zoals die in SN-Net) niet universeel optimaal zijn en vaak superieure "verre stitching"-configuraties missen (bijvoorbeeld het direct verbinden van Swin-Ti met Swin-B) die door heuristieken worden over het hoofd gezien.
2. Principiële Similariteitsmaatstaf: De introductie van KL-divergentie als maatstaf voor stitch-selectie die zowel representatieve als functionele uitlijning vastlegt zonder de instantiatie of training van stitch-lagen te vereisen.
3. KLAS-kader: Een generaliseerbaar, laagkosten kader dat anker- en blokselectie automatiseert over diverse modelfamilies (ViTs, CNN's en LLM's).
4. Empirische Validatie: Uitgebreide experimenten die consistente verbeteringen tonen ten opzichte van baselines over meerdere datasets (ImageNet-1K, CIFAR-100, ADE20K, TruthfulQA) en architecturen.

4. Experimentele Resultaten

De auteurs evalueerden KLAS tegen SN-Net en andere baselines onder identieke finetuning-omstandigheden (50 epochs voor visie, 10 epochs voor LLM's).

• Afwegingen tussen Nauwkeurigheid en Efficiëntie:

  • ImageNet-1K: KLAS bereikt tot 1,21% hogere Top-1-nauwkeurigheid bij dezelfde rekenkosten of een 1,33× reductie in FLOPs bij equivalente nauwkeurigheid in vergelijking met SN-Net.
  • AUC-scores: Op de Swin-familie behaalde KLAS een Area Under the Curve (AUC) van 0,8950, wat significant beter presteert dan SN-Net (0,8345) en andere similariteitsmaatstaven zoals CKA (0,8124).
  • Cross-Architectuur: KLAS slaagde erin ResNet-18 te stikken met Swin-B, en presteerde beter dan cascaded inference-baselines met 2-3% Top-1-nauwkeurigheid bij vergelijkbare FLOPs.

• Dichte Voorspelling (Semantische Segmentatie):

  • Op ADE20K met Mask2Former bereikten door KLAS geselecteerde stiksels tot +0,9% mIoU boven SN-Net-baselines bij vergelijkbare FLOPs-budgetten.

• Grote Taalmodellen (LLM's):

  • Met behulp van Llama 1B- en 3B-ankers op de TruthfulQA-benchmark presteerden door KLAS gestikte modellen beter dan ESTA (een heuristische LLM-stitchmethode) met tot 0,017 ROUGE-1 en 0,033 ROUGE-2, ondanks dat ze in sommige configuraties minder parameters hadden.

• Effectiviteit van Maatstaven:

  • KL-divergentie slaagde erin 88,9% van de levensvatbare Swin Ti-B stitch-configuraties te herstellen, terwijl CKA slechts 5,5% herstelde en SN-Net-heuristieken 61,1%.

5. Betekenis en Beweringen

Het artikel beweert dat KLAS de kritieke kloof in modelstitching aanpakt door ad-hoc heuristieken te vervangen door een principiële, similariteit-gedreven aanpak. De betekenis hiervan ligt in:

• Generaliseerbaarheid: Het werkt over diverse architectuurfamilies (Transformers, CNN's, Hybriden) en modaliteiten (Visie, Taal), waardoor het de beperkingen van methoden die zijn beperkt tot enkele families, overwint.
• Kosten-efficiëntie: Het genereert een continue Pareto-front van implementeerbare modellen met verwaarloosbare extra rekenkosten (alleen probe-training en finetuning van stitch-lagen), waardoor de onbetaalbare kosten van traditionele NAS worden vermeden.
• Systematische Optimalisatie: Door distributieverschillen expliciet vast te leggen via KL-divergentie, stelt KLAS de ontdekking van niet-intuïtieve, hoogpresterende stitch-configuraties (bijvoorbeeld "verre stitching") mogelijk die heuristische methoden missen.

De auteurs concluderen dat KLAS een robuuste, geautomatiseerde oplossing biedt voor het optimaliseren van de afweging tussen nauwkeurigheid en efficiëntie in implementatiescenario's, waarbij effectief gebruik wordt gemaakt van het bestaande ecosysteem van vooraf getrainde modellen zonder training vanaf nul.