SatReg: Regression-based Neural Architecture Search for Lightweight Satellite Image Segmentation

Probleemstelling

Met de verschuiving van aardobservatie-taken naar onboard en edge-verwerking op kleine satellieten, ontstaan er strenge beperkingen op het gebied van energie, geheugen, latentie en downlink-capaciteit. Semantische segmentatie voor remote sensing is hierbij bijzonder veeleisend omdat het zowel fijne lokale details als brede contextuele redenering vereist.
Recente hybride modellen (zoals CM-UNet), die convoluties combineren met Mamba- en transformer-mechanismen, bieden hoge nauwkeurigheid, maar zijn vaak te zwaar voor de beperkte hardware van satellieten. Traditionele methoden voor Neural Architecture Search (NAS) zijn te duur en tijdrovend omdat ze vaak een exhaustive search vereisen waarbij veel kandidaat-modellen getraind en geprofilleerd moeten worden onder specifieke deployment-beperkingen. Er is dus behoefte aan een efficiëntere methode om deze hybride modellen aan te passen aan ruimte-edge systemen zonder de prestaties drastisch te verliezen.

Methodologie: SatReg

Het paper introduceert SatReg, een regressie-gebaseerd framework voor hardware-bewuste tuning van het CM-UNet-architectuur. In plaats van de volledige architectuurruimte te doorzoeken, beperkt SatReg de zoekruimte tot twee dominante variabelen en gebruikt het surrogate-modellen om optimale configuraties te vinden.

1. Gereduceerde Ontwerpruimte:
   De architectuur wordt gereduceerd tot twee variabelen:

   • $b$ (Encoder breedte): Schaleert de kanaal-dimensies van de vier encoder-stadia, de tussenliggende MSAA- en fusiemodules, en de buitenste decoder-kanaalbreedtes. Dit bepaalt voornamelijk de macro-kenmerken en de kosten van de backbone.
   • $h$ (Decoder bottleneck breedte): Definieert de latente bottleneck binnen de CSMamba-blokken. Dit ontkoppelt de dure SS2D-berekeningen van de bredere decoder-kanaalbreedtes en heeft een directe invloed op latentie en vermogen met minimale impact op de nauwkeurigheid (mIoU).

2. Surrogate Modeling (Regressie):

   • Een kleine set van student-modellen wordt getraind en geprofilleerd op een NVIDIA Jetson Orin Nano.
   • Er worden lage-orde regressiemodellen (surrogaten) gefit voor:
     • mIoU: Een rationele regressieformule.
     • Latentie en Vermogen: Kwadratische polynomen ($\hat{L}(b,h)$ en $\hat{P}(b,h)$).
   • Deze surrogaten voorspellen de prestaties over het continue $(b, h)$-gebied zonder dat elk punt fysiek getraind hoeft te worden.

3. Optimalisatie:
   Een genormaliseerde objectief-functie wordt geminimaliseerd die een afweging maakt tussen lage latentie, laag vermogen en hoge mIoU. De gevonden continue oplossing wordt gemapt naar de dichtstbijzijnde geldige discrete CM-UNet-configuratie.

4. Kennisdistillatie (Knowledge Distillation):
   Omdat het veranderen van de breedte ($b$) de backbone-kanaal-dimensies wijzigt, kunnen vooraf getrainde gewichten niet direct worden overgenomen. Daarom wordt Channel-Wise Knowledge Distillation (CWD) gebruikt. Een groot "teacher"-model (basale CM-UNet) distilleert kennis naar de kleinere "student"-modellen, waardoor een deel van de verloren segmentatieprestaties wordt hersteld.

Belangrijkste Bijdragen

• Compacte Ontwerpruimte: Formulering van een efficiënte, twee-variabele ontwerpruimte voor CM-UNet die het afwegen van nauwkeurigheid, latentie en energie voor remote sensing mogelijk maakt.
• Efficiënte Surrogate-Modellen: Aantonen dat een klein aantal geprofileerde configuraties voldoende is om nauwkeurige surrogate-modellen te bouwen die trends in mIoU, latentie en vermogen vastleggen.
• Surrogate-Gestuurde Selectie: Een praktische aanpak om een bijna-optimale, gedistilleerde CM-UNet-configuratie te selecteren voor embedded hardware, zonder brute-force zoektochten.

Resultaten

De methode werd getest op twee datasets: LoveDA en ISPRS Vaihingen, met de Jetson Orin Nano als target hardware.

• Prestatieverbetering:

  • Latentie: Vermindering van 32,9% tot 38,6% ten opzichte van de baseline CM-UNet.
  • Energie: Vermindering van 43,3% tot 51,0% per afbeelding.
  • Doorvoer (Throughput): Toename van 5,60 FPS naar 9,11 FPS (LoveDA) en 8,35 FPS (Vaihingen).
  • Nauwkeurigheid: De mIoU bleef behouden met slechts een beperkte daling (bijv. van 50,21% naar 47,54% op LoveDA), wat een uitstekende trade-off is gezien de enorme winst in efficiëntie.

• Analyse van Variabelen:
  De contourplots toonden aan dat $b$ (encoder breedte) de sterkste invloed heeft op de mIoU, terwijl $h$ (decoder bottleneck) de belangrijkste factor is voor latentie en vermogen. De geoptimaliseerde modellen behouden daarom een brede encoder (voor nauwkeurigheid) maar verkleinen de decoder bottleneck (voor efficiëntie).

• Vergelijking:
  De geselecteerde SatReg-modellen presteerden beter in termen van energie-efficiëntie en doorvoer dan andere hybride Mamba-modellen (zoals RS3Mamba en UNetMamba), die ofwel niet draaiden op de Jetson ofwel een lagere resolutie vereisten.

Betekenis en Conclusie

SatReg demonstreert dat regressie-gestuurde, gereduceerde ruimte-tuning een zeer praktische strategie is voor het aanpassen van moderne, zware remote-sensing modellen aan toekomstige ruimte-edge AI-systemen. Door de zoekruimte slim te beperken en surrogate-modellen te gebruiken, kunnen ontwikkelaars snel near-optimal deployment points vinden die voldoen aan de strikte energielimieten van satellieten, zonder in te leveren op de essentiële segmentatiekwaliteit. Dit opent de deur voor meer geavanceerde onboard perceptie in kleine satellieten.