Technische Samenvatting: Regularisatie voor Gewichtsconcentratie ter Verbetering van de Robuustheid van Pruning onder Hoge Sparsiteit

Probleemstelling

Diepe neurale netwerken (DNN's) hebben state-of-the-art prestaties behaald in visie- en taakopgaven, maar vereisen vaak aanzienlijke rekenkracht, wat hun inzet beperkt in omgevingen met beperkte middelen. Modelpruning beoogt de modelgrootte te verkleinen door minder belangrijke parameters te verwijderen. Hoewel iteratieve pruning met hertraining effectief is, is dit rekenkundig duur. One-shot pruning, waarbij redundante parameters in één stap worden verwijderd zonder hertraining, biedt een efficiënter alternatief. Echter, modellen die met standaarddoelstellingen zijn getraind, lijden vaak onder significante nauwkeurigheidsdalingen bij agressieve sparsiteit.

Vorig werk probeert dit te mitigeren via twee hoofdrichtingen:

1. Op regularisatie gebaseerde benaderingen (bijv. $\ell_1$, DeepHoyer) die de verdeling van gewichten tijdens het trainen vormgeven. Echter, $\ell_1$ verkleint alle gewichten uniform, terwijl DeepHoyer schaal-invariante sparsiteit induceert. Geen van beide concentreert de gewichtsenergie expliciet op een kleine subset van informatieve parameters.
2. Op pruning robuuste optimalisatoren (bijv. SAM, CrAM, S2SAM) die het verlieslandschap afvlakken om generalisatie te verbeteren.

De kernbeperking die is geïdentificeerd, is dat standaard sparsiteit-inducerende methoden de afweging tussen sparsiteit en informatiebehoud niet expliciet aanpakken. Wanneer de belangrijkheid van gewichten diffuus is verdeeld, kan het verwijderen van gewichten met een kleine grootte de prestaties toch verslechteren, omdat de "belangrijke" energie niet voldoende is geconcentreerd op een paar dominante parameters.

Methodologie: Regularisator voor Gewichtsconcentratie (WCR)

De auteurs stellen Weight Concentration Regularizer (WCR) voor, een regularisator voor het trainingsmoment die is ontworpen om de verdeling van de grootte van gewichten te herschikken. In tegenstelling tot conventionele methoden die gewichten uniform verkleinen, beoogt WCR de gewichtsenergie te concentreren op een kleine subset van parameters, terwijl de rest naar nul wordt gedreven. Dit creëert een "heavy-tailed" verdeling waarbij een paar parameters het merendeel van de gewichtsenergie bevatten, waardoor magnitude-pruning de verwaarloosbare parameters kan verwijderen met minimale prestatieverlies.

Wiskundige Formulering

WCR werkt door de reciproke van de variantie van absolute gewichten te bestraffen. Voor een geselecteerde laag met gewichten $w^{(\ell)}$ passen de auteurs eerst een gladde absolute-waarde-transformatie toe om differentieerbaarheid bij de oorsprong te waarborgen:
$$\tilde{w}^{(\ell)}_i = \sqrt{(w^{(\ell)}_i)^2 + r} \approx |w^{(\ell)}_i|, \quad r = 10^{-8}$$
De variantie wordt per laag berekend om schaalverschillen te verwerken:
$$\text{Var}(\tilde{w}^{(\ell)}) = \frac{1}{n_\ell} \sum_{i=1}^{n_\ell} \left( \tilde{w}^{(\ell)}_i - \frac{1}{n_\ell} \sum_{j=1}^{n_\ell} \tilde{w}^{(\ell)}_j \right)^2$$
De regularisatiestraf $\psi(w)$ is de som van de reciproke varianties over alle geselecteerde lagen:
$$\psi(w) = \sum_{w^{(\ell)} \in \mathcal{W}} \frac{1}{\text{Var}(\tilde{w}^{(\ell)}) + \epsilon}$$
waarbij $\epsilon = 10^{-8}$ zorgt voor numerieke stabiliteit. Het totale trainingsdoel wordt:
$$\mathcal{L}_{total}(w) = \mathcal{L}(w) + \lambda \psi(w)$$
waarbij $\lambda > 0$ een balancerende coëfficiënt is.

Theoretische Analyse

Het artikel biedt een convergentieanalyse die aantoont dat WCR de standaard convergentiesnelheid van Stochastic Gradient Descent (SGD) behoudt. Onder $\beta_1$-gladheid van de verliesfunctie en $\beta_2$-gladheid van de regularisator voldoet het gecombineerde doel aan standaard garanties voor convergentie naar stationariteit. Aangezien de regularisator een deterministische functie van de gewichten is, introduceert deze geen extra stochastische gradientvariantie. Bijgevolg blijft de asymptotische convergentiesnelheid $O(1/\sqrt{T})$, wat overeenkomt met standaard SGD.

Belangrijkste Bijdragen

1. Introductie van WCR: Een nieuwe regularisator voor het trainingsmoment die de gewichtsenergie concentreert op een kleine subset van parameters, en expliciet een heavy-tailed gewichtsverdeling aanmoedigt in plaats van uniforme verkleining.
2. Theoretische Garanties: Bewijs dat WCR de standaard convergentiesnelheid van SGD behoudt onder milde gladheidsaannames.
3. Empirische Validatie: Uitgebreide evaluatie in diverse settings, waaronder:
   • Beeldclassificatie: ResNet-18/50, WideResNet-28-10 op CIFAR-10/100, SVHN, en ViT-B/32 op Tiny-ImageNet.
   • Medische Segmentatie: ResNet-50–UNet op het LGG Brain MRI-dataset.
   • LLM Fine-tuning: Qwen-2.5-1.5B en LLaMA-3.2-1B met LoRA op Commonsense-170K.
4. Compatibiliteit: Demonstratie dat WCR compatibel is met en bestaande pruning-robuste optimalisatoren (SAM, CrAM, S2SAM) verbetert, en met hen kan worden gecombineerd voor verdere winst.

Experimentele Resultaten

De auteurs hebben WCR geëvalueerd onder one-shot magnitude-pruning zonder hertraining.

• Visietaken (CNN's & ViT's): WCR verbeterde consistent de nauwkeurigheid na pruning, met name bij hoge sparsiteitsniveaus (bijv. 96% pruning).
  • Op CIFAR-10 met ResNet-18 verbeterde het toevoegen van WCR aan SGD de nauwkeurigheid bij 96% pruning van 19,69% naar 25,70%.
  • In combinatie met SAM leverde WCR dramatische verbeteringen op: bij 96% pruning steeg de nauwkeurigheid van 48,36% (alleen SAM) naar 92,08% (SAM + WCR).
  • Vergelijkbare trends werden waargenomen met S2SAM (72,51% $\to$ 93,44%) en op SVHN (9,74% $\to$ 95,99%).
  • Voor ViT-B/32 verbeterde WCR de nauwkeurigheid bij 80% pruning op CIFAR-100 van 64,11% (SAM) naar 84,95%.
• Medische Segmentatie: Op het LGG Brain MRI-dataset verbeterde WCR de F1-scores aanzienlijk en verlaagde de Hausdorff-afstand onder 85% en 88% pruning. Opmerkelijk is dat bij SGD met 88% pruning de baseline geen geldige segmentaties produceerde, terwijl SGD + WCR een F1-score van 0,5280 bereikte. Kwalitatieve resultaten toonden aan dat WCR laesiegrenzen beter behield dan baselines.
• LLM Fine-tuning: Bij LoRA-gebaseerde fine-tuning presteerde WCR beter dan $\ell_1$ en DeepHoyer regularisatoren. Op Qwen-2.5-1.5B bereikte WCR de hoogste gemiddelde nauwkeurigheid bij 30% pruning (69,54%) onder de vergeleken regularisatoren, wat robuustheid aantoont in scenario's voor parameter-efficiënte fine-tuning.
• Generalisatie: Analyse van de principale Hessian-eigenwaarde gaf aan dat WCR de kromming van het verliesoppervlak verlaagt, wat suggereert dat het optimalisatie naar vlakkere minima leidt, wat correleert met verbeterde generalisatie.

Betekenis en Claims

Het artikel beweert dat het vormgeven van de verdeling van de grootte van gewichten tijdens het trainen een effectieve aanpak is om de robuustheid onder hoge sparsiteit te verbeteren. De betekenis van WCR ligt in het vermogen om:

• Een duidelijkere scheiding te creëren tussen grote-magnitude (informatieve) en kleine-magnitude (verwaarloosbare) parameters.
• Magnitude-gebaseerde pruning in staat te stellen een groot deel van de parameters te verwijderen met minimaal functioneel effect.
• Te werken als een standalone regularisator of als een complementair onderdeel voor methoden voor het afvlakken van het verlieslandschap zoals SAM.

De auteurs concluderen dat WCR een praktische oplossing biedt voor het trainingsmoment voor de nauwkeurigheidsdaling die gepaard gaat met agressieve one-shot pruning, toepasbaar op CNN's, Transformers en Large Language Models, zonder dat architecturale wijzigingen of hertraining na pruning nodig zijn.