Technische Samenvatting: Model Fusie via Retrofitting

Probleemstelling

Modelfusie beoogt onafhankelijk getrainde neurale netwerken te combineren tot één enkel model zonder hertraining, en biedt zo een oplossing voor scenario's waarin hertraining onbetaalbaar duur of onuitvoerbaar is vanwege gegevensprivacybeperkingen (bijvoorbeeld Federated Learning). Deze taak wordt echter bemoeilijkt door representatiedivergentie veroorzaakt door permutatie-invariantie, willekeurige initialisatie en heterogene verdelingen van trainingsdata.

Bestaande methoden, zoals OTFusion (Singh & Jaggi, 2020) en Git Re-Basin (Ainsworth et al., 2023), vertrouwen op de Linear Mode Connectivity (LMC) conjectuur, die stelt dat onafhankelijk getrainde netwerken in hetzelfde laag-verliesbekken liggen. Hoewel LMC geldt voor modellen die op identieke data zijn getraind, geven eerdere studies aan dat zelfs modellen die op dezelfde data zijn getraind, divergente interne kenmerken kunnen leren. Bijgevolg kampen bestaande fusiemethoden met aanzienlijke problemen in zero-shot settings (fusie zonder verdere training) en onder non-IID (non-Independent and Identically Distributed) data-regimes, waarbij ze vaak instorten tot bijna-willekeurige prestaties. Bovendien missen huidige benaderingen vaak generaliteit over architecturen heen (bijvoorbeeld moeite met Transformers) en vertrouwen ze op basismodellen die mogelijk niet volledig geoptimaliseerd zijn, waardoor het ware potentieel van fusie-algoritmen wordt verduisterd.

Methodologie

De auteurs stellen een neuron-gerichte familie van fusie-algoritmen voor die modelfusie herformuleert als een principieel representatie-matching probleem. In plaats van gewichten simpelweg te middelen of permutatiematrices te vinden, behandelt de methode fusie als een tweestapsproces: het groeperen van neuronen van oudermodellen in doelrepresentaties en vervolgens het trainen van de sub-netwerken van het gefuseerde model om deze doelen te benaderen.

Kernkader

1. Representatiekosten: Het doel wordt gedefinieerd door de gekwadrateerde $L_2$-afstand tussen de activaties van neuronen van het basismodel en de dichtstbijzijnde neuronen in het gefuseerde model, gewogen door neuron-impactscores (toeschrijvingsscores). Dit straft de verkeerde representatie van opvallende kenmerken af.
2. Decompositie (Stelling 1): De auteurs introduceren een hulpdoelvector $T$ om de representatiekosten te ontleden in twee interpreteerbare componenten:
   • Groeperingsfout: Meet hoe goed neuronen van het basismodel clusteren rondom impact-gewogen centroiden ($T_k$).
   • Benaderingsfout: Meet hoe goed de uitvoerneuronen van het gefuseerde model deze centroiden benaderen.
3. Tweestaps-algoritme:
   • Stap 1 (Groepering): Neuronen van de basismodellen worden geclusterd om de groeperingsfout te minimaliseren. Dit levert impact-gewogen centroiden op die als doelen dienen.
   • Stap 2 (Retrofitting): De gewichten van het gefuseerde model op het huidige niveau worden geoptimaliseerd (via een gesloten-formule oplossing voor lineaire lagen of SGD voor niet-lineaire lagen) om de benaderingsfout ten opzichte van deze doelen te minimaliseren. Dit proces wordt van onder naar boven door de netwerkniveaus uitgevoerd.

Algoritmevarianten

Het artikel introduceert twee primaire algoritmen gebaseerd op de groeperingsstrategie:

• Hongaarse Fusie (HF): Ontworpen voor modellen met gelijke grootte met een-op-een matching per niveau. Het maakt gebruik van het Hongaarse algoritme om de optimale toewijzing te vinden.
• K-means Fusie (KF): Ontworpen voor het algemene geval, inclusief modellen met willekeurige breedtes en dieptes. Het maakt gebruik van heuristische K-means clustering om neuronen te groeperen.

Beide algoritmen hebben Lineaire (gesloten-formule oplossing voor affiene lagen) en Op gradiënten gebaseerde (SGD/Adam voor niet-lineaire lagen) varianten.

Belangrijkste Innovaties

• Neuron-toeschrijving: In tegenstelling tot eerdere werken, integreert deze aanpak neuron-impactscores (bijvoorbeeld Conductance, DeepLIFT) om de uitlijning te beïnvloeden richting opvallende kenmerken, wat de prestaties in heterogene settings verbetert.
• Architecturale flexibiliteit: De methode is toepasbaar op elke architectuur die modulariseerbaar is als een Directed Acyclic Graph (DAG) van niveaus, en is empirisch gevalideerd op VGG's, ResNets en Vision Transformers (ViTs).
• Zero-shot capaciteit: Door expliciet de benadering van geclusterde representaties te optimaliseren, bereikt de methode sterke prestaties zonder de uitgebreide fine-tuning die door eerdere methoden wordt vereist.

Experimentele Resultaten

De auteurs evalueerden hun methoden over drie trainingsregimes: Sharded (disjuncte klassensets), Non-IID (scheve klassenverdelingen) en Volledige Dataset.

• Sharded/Non-IID Regimes: In extreme non-IID settings (bijvoorbeeld 6-weg klassensplitsen op CIFAR-100) instortten bestaande methoden zoals OTFusion en Git Re-Basin tot bijna-willekeurige nauwkeurigheid (vaak <5%). Daarentegen behielden KF-Gradient en HF-Gradient robuuste zero-shot nauwkeurigheid (bijvoorbeeld 34,7% op CIFAR-100 voor 6-weg splitsen versus <1,5% voor baselines).
• Volledige Dataset Regimes: Zelfs wanneer modellen op de volledige dataset zijn getraind, presteren de voorgestelde methoden beter dan baselines in zero-shot fusie. Bijvoorbeeld, op CIFAR-100 bereikte KF-Gradient 63,0% zero-shot nauwkeurigheid vergeleken met 4,3% voor Transformer OTFusion. Na fine-tuning benaderden de voorgestelde methoden consistent de ensemble-prestaties, terwijl ze de inferentiekosten van een enkel model behielden.
• Compressie en Heterogeniteit: De methode slaagde erin modellen van verschillende breedtes te fuseren en grotere modellen (ResNet-50) te comprimeren tot kleinere modellen (ResNet-18) met aanzienlijk hogere nauwkeurigheid dan standaard Knowledge Distillation (KD) wanneer data beperkt was.
• Robuustheid: De aanpak toonde robuustheid tegen variërende groottes van fusiedatasets, waarbij data-augmentatie effectief de behoefte aan grote fusiesets mitigeerde.

Betekenis en Claims

Het artikel claimt drie belangrijke gaten in het eerdere onderzoek aan te pakken:

1. Reproduceerbaarheid: Het bieden van een flexibele, open-source herimplementatie die systematische benchmarking over architecturen heen faciliteert.
2. Kwaliteit van basismodellen: Het aantonen dat fusiemethoden aanzienlijke verbeteringen kunnen opleveren, zelfs wanneer basismodellen goed getraind zijn, wat de notie uitdaagt dat slechte fusieresultaten uitsluitend te wijten zijn aan zwakke baselines.
3. Heterogene data: Het tonen aan dat bestaande methoden falen in zero-shot fusie onder non-IID condities, terwijl de voorgestelde neuron-gerichte aanpak modellen succesvol integreert die zijn getraind op disjuncte of scheve data-verdelingen.

De auteurs benadrukken dat hun aanpak de eerste is die neuron-impactscores succesvol integreert in modelfusie, wat leidt tot consistente verbeteringen ten opzichte van bestaande baselines, met name in zero-shot en non-IID scenario's waar andere methoden falen. Ze positioneren de methode als een praktische oplossing voor Federated Learning en privacy-bewuste modelaggregatie, waar hertraining geen optie is. Hoewel ze erkennen dat varianten op gradiënten gebaseerd computatiekosten hoger zijn dan lichtgewicht baselines, betogen de auteurs dat deze overhead gerechtvaardigd is door het vermogen om bruikbare gefuseerde modellen te produceren in regimes waar concurrenten onbruikbare resultaten opleveren.