FairQuant: Fairness-Aware Mixed-Precision Quantization for Medical Image Classification

Dit artikel introduceert FairQuant, een framework voor eerlijkheidsbewuste mixed-precision kwantisatie dat de prestaties van medische beeldclassificatiemodellen onder strikte bitbudgetten optimaliseert door de nauwkeurigheid te behouden en de prestaties voor de slechtst presterende groepen te verbeteren.

De kern

Stel je voor dat je een heel slimme, maar ook heel grote en zware robot bouwt om huidafwijkingen op foto's te herkennen. Deze robot is zo goed dat hij bijna net zo goed werkt als een menselijke arts. Maar er is een probleem: deze robot is te groot en te traag om op een gewone telefoon of in een kleine kliniek te gebruiken. Hij heeft te veel "geheugen" en "rekenkracht" nodig.

Om dit op te lossen, proberen onderzoekers de robot te verkleinen. Ze doen dit door de getallen in de robot's hersenen (de "gewichten") minder precies te maken. In plaats van getallen met tien decimalen (zoals 3.1415926535) te gebruiken, gebruiken ze afgeronde getallen (zoals 3.14). Dit noemen ze kwantisatie.

Het probleem met de oude manier:
Stel je voor dat je de robot verkleint door alle onderdelen even hard aan te halen. Je maakt alles "4-bit" (zeer ruw).

• Het resultaat: De robot wordt wel klein en snel, maar hij wordt ook onbetrouwbaar.
• De ongelijkheid: Het ergste is dat hij niet voor iedereen even slecht wordt. Voor mensen met lichte huid werkt hij misschien nog redelijk, maar voor mensen met een donkere huid faalt hij volledig. De gemiddelde score is misschien nog oké, maar voor de kwetsbare groepen is de robot gevaarlijk.

De oplossing: FairQuant (De eerlijke robot)
De auteurs van dit paper, Thomas en Raghavendra, hebben een nieuwe methode bedacht die ze FairQuant noemen. Ze gebruiken een slimme strategie om de robot te verkleinen, zodat hij niet alleen klein wordt, maar ook eerlijk blijft.

Hier is hoe het werkt, vertaald in alledaagse termen:

1. De "Belangrijkheidskaart" (Wie heeft wat nodig?)

Stel je voor dat je een grote groep mensen hebt die een zware tas moeten dragen. Sommige mensen hebben sterke armen, anderen niet.

• De oude manier: Je geeft iedereen precies dezelfde zware tas. De zwakken zakken door hun knieën.
• FairQuant: Ze kijken eerst naar de robot en vragen: "Welke onderdelen van de robot zijn het belangrijkst om goed te zien voor mensen met een donkere huid, en welke voor mensen met een lichte huid?"
• Ze maken een belangrijkheidskaart. Ze zien dat bepaalde onderdelen van de robot cruciaal zijn om donkere huidtinten goed te zien, terwijl andere onderdelen minder belangrijk zijn.

2. Slimme Verdeling (Mixed-Precision)

Nu gaan ze de robot verkleinen, maar ze doen het niet voor iedereen gelijk.

• De belangrijke onderdelen: De onderdelen die nodig zijn om de "minderheidsgroepen" (zoals mensen met donkere huid) goed te behandelen, krijgen veel precisie (bijvoorbeeld 8 bits). Ze blijven bijna net zo goed als het origineel.
• De minder belangrijke onderdelen: De onderdelen die niet zo kritiek zijn, krijgen minder precisie (bijvoorbeeld 2 of 4 bits). Deze worden flink verkleind om ruimte te besparen.

Dit noemen ze Mixed-Precision: een mix van hoge en lage precisie, net als een orkest waar de viool (belangrijk) veel aandacht krijgt, maar de fluit (minder belangrijk) wat minder.

3. De "Leerzame" Robot (BAQ)

Het allercoolest aan FairQuant is dat de robot niet alleen een vaste kaart krijgt, maar leert hoe hij zichzelf moet verdelen terwijl hij traint.

• Stel je voor dat de robot tijdens het oefenen zelf mag beslissen: "Hey, ik merk dat ik bij deze foto's de donkere huid beter kan zien als ik hier iets meer precisie gebruik. Ik ga mijn 'precisie-batterij' hierheen verplaatsen."
• De robot probeert continu de balans te vinden tussen:
  1. Snelheid/Grootte: Niet te veel bits gebruiken.
  2. Eerlijkheid: Zorg dat hij voor alle huidtinten even goed werkt.

Wat zeggen de resultaten?

De onderzoekers hebben dit getest op echte medische foto's van mensen met verschillende huidtinten.

• De oude methode (Uniform 4-bit): De robot werd klein, maar hij was een ramp voor mensen met donkere huid. De diagnose was vaak fout.
• FairQuant: De robot werd net zo klein (gemiddeld 4 tot 6 bits), maar hij hield zijn belofte in. Hij deed het bijna net zo goed als de grote, zware robot, en hij maakte veel minder fouten bij de groepen die eerder werden genegeerd.

Conclusie

Kortom: FairQuant is als het herverdelen van middelen in een team. In plaats van iedereen even hard te straffen om ruimte te besparen, kijken ze wie wat nodig heeft. Ze geven de kwetsbaren extra steun (precisie) en nemen het minder belangrijke weg. Het resultaat is een snelle, kleine robot die voor iedereen eerlijk en betrouwbaar werkt, niet alleen voor de "gemiddelde" patiënt.

Dit is een enorme stap voor medische technologie, omdat het betekent dat slimme AI-toepassingen in de toekomst ook in arme of kleine klinieken kunnen werken, zonder dat ze de zorg voor minderheidsgroepen opofferen.

Technische samenvatting

Probleemstelling

Diepe neurale netwerken worden steeds vaker ingezet voor medische beeldanalyse, maar er zijn twee kritieke uitdagingen:

1. Efficiëntie: Klinische toepassingen vereisen vaak inferentie met beperkte rekenkracht, geheugen en energie. Quantisatie (het converteren van floating-point parameters naar lage precisie-integers) is een veelgebruikte techniek om modellen te comprimeren.
2. Fairness (Billijkheid): Medische datasets vertonen vaak ondervertegenwoordiging van bepaalde subgroepen (bijv. op basis van huidskleur of geslacht). Bestaande quantisatiemethoden (zoals Post-Training Quantization of Quantization-Aware Training) richten zich uitsluitend op het behoud van de gemiddelde nauwkeurigheid. Hierdoor kan het zijn dat de prestaties van de minst bediende subgroepen (worst-group performance) drastisch verslechteren, zelfs als de gemiddelde nauwkeurigheid hoog blijft. Dit is in de medische context onacceptabel.

Het paper onderzoekt hoe men fairness-bewuste mixed-precision quantisatie kan toepassen onder strikte bit-begrotingen, zodat de compressie niet ten koste gaat van de betrouwbaarheid voor kwetsbare groepen.

Methodologie: FairQuant Framework

Het voorgestelde framework, FairQuant, combineert groep-georiënte analyse met een leerbare bit-toewijzing. Het proces verloopt in drie hoofdstappen:

1. Groep-gevoelige Importantie Analyse (Group-Sensitive Importance Analysis):

   • Tijdens een kalibratiefase wordt het model bevroren.
   • Er wordt een "importance map" berekend voor elke kwetsbare groep (bijv. lichte vs. donkere huidskleur). Dit gebeurt door de gradiënten van de loss-functie per groep te analyseren (gebaseerd op een Taylor-uitbreiding).
   • Deze groep-specifieke signalen worden samengevoegd tot één importance map die zowel de algehele gevoeligheid van de netwerklagen als de dispariteit tussen groepen weerspiegelt.

2. Begrensde Mixed-Precision Toewijzing (Budgeted Mixed-Precision Allocation):

   • Op basis van de importance map wordt een statische toewijzing van bit-breedtes gemaakt onder een globale bit-begroting.
   • Lagen die belangrijk zijn voor specifieke groepen of algemeen gevoelig zijn, krijgen meer bits, terwijl minder belangrijke lagen meer gecomprimeerd worden.
   • Dit dient als een "warm start" voor de volgende stap.

3. Bit-Aware Quantization (BAQ):

   • Dit is de kerninnovatie. In plaats van statische bit-breedtes, worden de bit-breedtes per netwerkeenheid (scope) behandeld als leerbare parameters.
   • Het model leert continu proxies voor bit-breedtes ($b_{cont}$) die worden afgerond naar discrete waarden.
   • De trainingsloss bestaat uit drie componenten:
     • Taak-Loss: De standaard classificatiefout.
     • Fairness-Loss: Een regularisatieterm die de verspreiding van de fouten tussen groepen minimaliseert (maximale groepsfout minus minimale groepsfout).
     • Bitrate-Loss: Een $L_2$-strafterm die de leerbare bit-logits naar nul duwt, waardoor de bit-breedtes worden beperkt tot de gewenste begroting.
   • Aan het einde van de training worden de geleerde bit-breedtes vastgezet voor inferentie.

Belangrijkste Bijdragen

• FairQuant Framework: Een nieuw systeem dat groep-gevoelige analyse koppelt aan een begrensde mixed-precision allocatie.
• Bit-Aware Quantization (BAQ): Een methode waarbij bit-breedtes worden geoptimaliseerd als trainbare variabelen, gekoppeld aan zowel een bitrate- als een fairness-regularisator.
• Robuustheid: De methode is getest op verschillende architecturen (CNN's zoals ResNet en Transformers zoals DeiT/TinyViT) en blijkt stabiel over verschillende random seeds en hyperparameters.

Resultaten

Het onderzoek is uitgevoerd op twee dermatologische datasets: Fitzpatrick17k (met huidskleur als gevoelige attribuut) en ISIC 2019 (met geslacht als attribuut).

• Prestaties bij lage precisie: Uniforme 4-bit quantisatie leidt vaak tot een instorting van de prestaties, vooral voor de "worst-group" (bijv. bij TinyViT op Fitzpatrick17k daalt de nauwkeurigheid van ~50% naar ~3% bij uniforme 4-bit).
• FairQuant presteert beter:
  • Met een gemiddelde precisie van 4-6 bits herstelt FairQuant een groot deel van de nauwkeurigheid van het volledige 8-bit model.
  • Het verbetert de worst-group nauwkeurigheid aanzienlijk ten opzichte van uniforme 4-bit baselines.
  • Op ISIC 2019 behaalt FairQuant (BAQ) bij ~4.1 bits een gemiddelde nauwkeurigheid van 82.73% en een worst-group nauwkeurigheid van 81.13%, wat dicht in de buurt komt van het 8-bit model (84.6% / 82.5%), terwijl uniforme 4-bit veel slechter scoort.
• Stabiliteit: Ablatiestudies tonen aan dat de methode robuust is ten opzichte van de gewichten van de regularisatoren (bitrate en fairness), zolang deze binnen een redelijk bereik blijven.

Betekenis en Conclusie

FairQuant demonstreert dat compressie van neurale netwerken niet louter een technisch optimalisatieprobleem is, maar dat groepsinformatie essentieel is voor betrouwbare medische AI.

• Praktische impact: Het biedt een pad naar efficiënte, op apparaten draaiende medische modellen die niet alleen snel zijn, maar ook eerlijk en betrouwbaar voor alle patiëntgroepen, zelfs onder strikte hardware-beperkingen.
• Inzicht: Het paper laat zien dat "one-size-fits-all" quantisatie (zoals uniforme 4-bit) gevaarlijk kan zijn voor subgroepen, terwijl een slimme, groep-bewuste toewijzing van bits deze dispariteiten kan oplossen zonder de totale efficiëntie te verliezen.

Kortom, FairQuant creëert een betere trade-off tussen compressie, nauwkeurigheid en eerlijkheid, wat cruciaal is voor de veilige implementatie van AI in de gezondheidszorg.