HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces

Technische Samenvatting: HASTE – Hardware-bewuste Dynamische Sparse Training voor Grote Outputruimtes

Probleemstelling
Extreme Multi-Label Classificatie (XMC) omvat het trainen van modellen over outputruimtes die miljoens labels bevatten. Hoewel methoden op basis van sparsity (ijdelheid) zijn voorgesteld om de rekenkundige complexiteit van de outputlaag te verminderen, leveren deze in de praktijk vaak geen evenredige versnellingen op. De primaire knelpunten zijn onregelmatige geheugentoegangspatronen, slechte hardware-utilisatie op moderne GPU's (die geoptimaliseerd zijn voor dichte, aaneengesloten blokken), en de "memory wall" veroorzaakt door ongecoalesceerde reads. Bovendien leidt training in long-tailed regimes met sparse connectiviteit vaak tot instabiliteit door zwakke gradiëntsignalen, wat regelmatig aanvullende objectieven vereist die hun eigen computationele overhead en potentiële gradiëntconflicten introduceren.

Methodologie
De auteurs stellen HASTE voor, een framework dat een nieuw semi-gestructureerd sparsity-patroon combineert met een datagedreven optimalisatiestrategie.

1. Group-Shared Fixed Fan-In Sparsity:
   In tegenstelling tot standaard fixed fan-in sparsity (waarbij elk label een onafhankelijke set inputfeatures heeft) of rigide block sparsity, introduceert HASTE een "group-shared" aanpak.

   • Structuur: Labels worden gepartitioneerd in groepen op basis van semantische gelijkenis (gebruikmakend van label-embeddings afgeleid van trainingsinstanties). Alle labels binnen een groep delen dezelfde sparse inputfeature-indices (fan-in patroon), hoewel ze onafhankelijke gewichtsvectoren behouden.
   • Hardware-alignement: Dit ontwerp maakt "gather-once and dense MMA" executiepatronen mogelijk. Door labels te groeperen, kunnen de inputfeatures voor een groep één keer naar het shared memory worden opgehaald en worden hergebruikt door alle labels in die groep. Dit faciliteert het gebruik van Tensor Core primitives via aangepaste CUDA-kernels, waardoor rekenkundige reducties worden omgezet in werkelijke wall-clock versnellingen.
   • Geheugenefficiëntie: Deze aanpak vermindert de indexgeheugen-overhead drastisch. In plaats van indices voor elk label afzonderlijk op te slaan, worden indices één keer per groep opgeslagen, wat de overhead vermindert met een factor proportioneel aan de groepsgrootte.

2. Head–Tail Decompositie (Auxiliary-Free Optimalisatie):
   Om de trainingsinstabiliteit in long-tailed distributies aan te pakken zonder te vertrouwen op aanvullende objectieven, decomponeert HASTE de labelruimte:

   • Head: Een kleine subset van frequente labels wordt gemodelleerd met een standaard dichte classifier.
   • Tail: De overige, enorme meerderheid van zeldzame labels wordt afgehandeld door de group-shared sparse laag.
   • Mechanisme: Beide takken delen dezelfde encoder maar maken gebruik van aparte, lichtgewicht feature projecties. De dichte head zorgt voor een consistent, sterk gradiëntsignaal naar de encoder tijdens de training, wat het leerproces van de sparse tail stabiliseert zonder de noodzaak voor aanvullende supervisie-branches.

Belangrijkste Bijdragen

• Nieuw Sparsity-patroon: De introductie van group-shared fixed fan-in sparsity, die de expressiviteit van onafhankelijke featureselectie balanceert met de hardware-efficiëntie van gestructureerde, gedeelde toegang.
• Hardware-bewuste Kernels: De ontwikkeling van aangepaste CUDA-kernels die de groepsstructuur benutten om dense matrixvermenigvuldigingen (GEMM) uit te voeren op verzamelde feature-tiles, wat efficiënt gebruik van Tensor Cores mogelijk maakt.
• Auxiliary-Free Stabiliteit: Een decompositiestrategie die training in long-tailed regimes stabiliseert door de inherente frequentieverdeling van labels te benutten, waardoor de noodzaak voor aanvullende objectieven en de bijbehorende overhead wordt geëlimineerd.
• Empirische Validatie: Uitgebreide evaluatie over grootschalige benchmarks (tot 8,6 miljoen labels) die aantoont dat HASTE de prestatiekloof met dense modellen verkleint terwijl het aanzienlijk beter presteert dan eerdere sparse baselines op het gebied van snelheid en geheugenefficiëntie.

Experimentele Resultaten

• Versnelling: Microbenchmarks op NVIDIA A100 GPU's laten zien dat HASTE een tot 4,4× versnelling bereikt in de forward pass en tot 25× versnelling in backward passes vergeleken met standaard fixed fan-in sparsity (SPARTEX). Het opereert binnen enkele procenten van een FLOPs-gematchte dense bottleneck.
• Nauwkeurigheid: Op benchmarks inclusief Amazon-670K, Amazon-3M, en LF-Paper2Keywords-8.6M matcht of verbetert HASTE de Precision@k (P@k) en Propensity-Scored Precision@k (PSP@k) ten opzichte van state-of-the-art sparse baselines (bijv. SPARTEX, BLOCK SPARSE).
• Geheugenefficiëntie: De methode vermindert het piekgeheugen tijdens training aanzienlijk (bijv. 3,2 GiB vs. 5,0 GiB voor SPARTEX op AmazonTitles-670K) en verkort de tijd per epoch.
• Groepsgrootte Trade-off: Experimenten wijzen op een trade-off waarbij kleinere groepsgroottes (bijv. 16) iets hogere nauwkeurigheid opleveren vanwege betere representationele flexibiliteit, terwijl grotere groepen (bijv. 64) de kernel-efficiëntie verbeteren door een hogere feature-reuse.

Betekenis en Claims
Het artikel claimt dat HASTE er succesvol in slaagt de kloof tussen de theoretische rekenkundige voordelen van sparsity en de praktische realiteit van moderne GPU-hardware te overbruggen. Door structuur aan te brengen in het sparsity-patroon (group-sharing) en dit af te stemmen op de long-tailed aard van XMC-data (head-tail split), zet de aanpak FLOPs-reducties om in tastbare wall-clock winst. De auteurs stellen dat hun methode het trainen van modellen met miljoens labels toegankelijker maakt onder vaste compute-budgetten door zowel de trainingstijd als de geheugenvoetafdruk te verkleinen, terwijl een concurrerende nauwkeurigheid behouden blijft. Ze merken op dat hoewel de efficiëntie is verbeterd, het niet inherent problemen oplost die gerelateerd zijn aan databias of imperfecte labeling in long-tailed datasets, en dat gebruikers waakzaam moeten blijven met betrekking tot de prestaties van tail-labels.