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

1. Problem Statement

As Earth-observation missions shift toward onboard and edge processing (Space-Edge computing), remote-sensing models face severe constraints regarding power, memory, latency, and downlink bandwidth.

• The Challenge: State-of-the-art hybrid models, such as CM-UNet (combining CNNs with Mamba/Transformer mechanisms), offer high accuracy for semantic segmentation but are too computationally expensive for resource-constrained edge hardware (e.g., small satellites).
• The Gap: Traditional Neural Architecture Search (NAS) methods require training and profiling thousands of candidate models, which is prohibitively expensive in terms of time and energy. Furthermore, standard NAS often fails to account for the specific hardware characteristics of edge platforms like the NVIDIA Jetson Orin Nano.
• Goal: Develop a method to efficiently find a "near-optimal" lightweight architecture that balances segmentation accuracy (mIoU) with hardware costs (latency and power) without exhaustive search.

2. Methodology: SatReg Framework

The authors propose SatReg, a regression-based, hardware-aware tuning framework that reduces the search space to a manageable size and uses surrogate models to predict performance.

A. Reduced Design Space

Instead of searching over complex architectural choices (e.g., layer counts, topology changes), SatReg defines a compact two-variable design space:

1. Encoder Base Width ($b$): Scales the channel dimensions of the ResNet-style encoder, intermediate fusion modules, and decoder outer channels. This primarily controls the model's feature capacity and backbone cost.
2. Decoder Bottleneck Width ($h$): Defines the latent bottleneck width within the CSMamba (Mamba-based) decoder blocks. This decouples the expensive Selective Scan (SS2D) computation from the outer channel widths, acting as a direct control knob for latency and power with minimal impact on accuracy.

B. Regression-Based Surrogate Modeling

The framework avoids training all possible configurations by profiling a small subset and fitting mathematical models (surrogates) to predict performance across the continuous $(b, h)$ space:

• Sampling: A small set of student models is trained and profiled on an NVIDIA Jetson Orin Nano.
• Surrogate Functions:
  • Latency ($\hat{L}$) & Power ($\hat{P}$): Modeled using quadratic polynomials (2nd degree) because higher-order polynomials overfit the data.
  • Accuracy ($\hat{m}$ - mIoU): Modeled using a rational-form regression (similar to RegTune), which captures the non-linear relationship between width and accuracy.
• Optimization: The system minimizes a normalized objective function that balances low latency, low power, and high mIoU to find the optimal continuous $(b, h)$ point, which is then mapped to the nearest valid discrete configuration.

C. Knowledge Distillation

Since changing the width ($b$) alters the backbone dimensions, pre-trained weights cannot be directly transferred.

• Teacher-Student Setup: The baseline CM-UNet (trained from scratch) acts as the Teacher.
• Training Strategy: Student models are trained using Channel-Wise Knowledge Distillation (CWD) on backbone features combined with the segmentation loss. This allows smaller models to recover performance lost due to width reduction without needing to pretrain from scratch.

3. Key Contributions

1. Compact Design Space Formulation: Introduced a two-variable parameterization ($b$ and $h$) for hybrid CNN-Mamba models, enabling efficient exploration of accuracy/latency/energy trade-offs.
2. Surrogate-Guided Optimization: Demonstrated that a small number of sampled configurations (16 per dataset) is sufficient to fit accurate surrogate models (with $R^2 \approx 97\%$) that capture hardware trends, eliminating the need for brute-force search.
3. Practical Edge Deployment: Successfully identified near-optimal configurations for the Jetson Orin Nano, proving that regression-guided tuning is a viable strategy for adapting modern remote-sensing models to future space-edge systems.

4. Experimental Results

The framework was evaluated on two benchmarks: LoveDA and ISPRS Vaihingen, using an NVIDIA Jetson Orin Nano.

• Performance Gains:
  • Latency: Reduced by 32.9% – 38.6% compared to the baseline CM-UNet.
  • Energy Consumption: Reduced by 43.3% – 51.0% per image.
  • Throughput: Increased from 5.6 FPS to **9.1 FPS** (LoveDA) and ~8.4 FPS (Vaihingen).
  • Accuracy: Retained the majority of the baseline mIoU (e.g., on Vaihingen, baseline was 84.64% vs. SatReg's 80.64%).
• Hardware Efficiency:
  • Power: Dropped from ~7.2W to ~5.7W.
  • Model Size: Reduced significantly (e.g., from 51.56 MB to 12.36 MB on LoveDA).
• Insights from Surrogates:
  • mIoU is most sensitive to the Encoder Width ($b$).
  • Latency and Power are most sensitive to the Decoder Bottleneck ($h$).
  • Strategy: The optimal solution retains high encoder capacity (for accuracy) while shrinking the decoder bottleneck (for speed/energy).

5. Significance

• Efficiency: SatReg offers a "fast-track" alternative to traditional NAS, reducing the computational cost of finding optimal architectures by orders of magnitude.
• Space-Edge Applicability: It directly addresses the constraints of satellite hardware, providing a blueprint for deploying complex hybrid models (CNN-Mamba) on devices with strict power and thermal limits.
• Generalizability: The approach of using reduced-space regression and knowledge distillation can be applied to other hybrid architectures and edge platforms beyond satellite imagery.

In conclusion, SatReg demonstrates that by intelligently narrowing the search space and leveraging regression models, it is possible to deploy high-performance, lightweight segmentation models on edge devices without sacrificing critical accuracy.