- The paper introduces RS4D, which leverages knowledge distillation to transfer transformer segmentation priors into a compact linear-time state space model for remote sensing.
- It employs a two-stage training paradigm with domain-adaptive noise and masking, resulting in significant parameter and FLOPs reduction while achieving competitive AP scores.
- The approach demonstrates real-time efficiency and robust spatial aggregation, making it ideal for resource-constrained applications like edge devices and in-orbit processing.
Linear-Time Instance Segmentation for Remote Sensing via State Space Distilled Visual Foundation Models
Introduction
The RS4D framework addresses the computational constraints facing vision transformers in dense prediction tasks, notably remote sensing instance segmentation. Transformer-based visual foundation models such as SAM and its derivatives have dominated segmentation accuracy benchmarks but their quadratic self-attention scaling is prohibitive for high-resolution remote sensing imagery, which typically involves long token sequences and demands real-time, memory-efficient inference. RS4D leverages knowledge distillation to compress the expressive self-attention priors from base foundation models into a compact, linear-time state space model (SSM), thus enabling efficient instance segmentation on high-resolution remote sensing datasets.
Methodology
State Space Model Distillation
RS4D adopts a two-stage training paradigm. The backbone is distilled from a pretrained ViT-based SAM encoder using ~0.1–1% of SA-1B, with domain-adaptive noise and masking augmentations to bridge the visual domain gap between natural and remote sensing imagery. Three SSM architectures—VanillaMamba, TransMamba, and ScanningMamba—encode prior knowledge into linear state space, each designed for efficiency versus representation completeness. ScanningMamba employs bidirectional scanning with stacked SSM blocks to enhance spatial context aggregation and achieves superior segmentation performance and convergence stability.
Knowledge transfer is formulated as a masked/noised consistency distillation, where the student SSM encoder matches teacher outputs with an adaptive noise amplitude, ensuring robustness to high-frequency perturbations and masking. The supervision signal is L1 distance between teacher and student output feature maps. Empirical results show that 20% noise/masking yields peak segmentation robustness.
Model Architecture
The system consists of the distilled backbone, a lightweight neck with multi-level feature fusion (skip connections and SimpleFPN on shallow/deep SSM outputs), and flexible instance segmentation heads. Box-based (Mask R-CNN style) and prompt-based (SAM-injected) heads are explored. The neck avoids traditional FPN pyramids, instead fusing selected hidden layers to balance representational depth with computational cost; ablation reveals that increasing fusion layers improves segmentation accuracy.
Fine-tuning is performed on SSDD, WHU, and NWPU datasets with standard augmentation, AdamW optimization, and bfloat16 precision. Distilled SSM weights provide superior initialization and faster convergence compared to random or standard ImageNet pretraining.
Experimental Results
Quantitative Analysis
RS4D achieves a ×8 parameter reduction and ×9 FLOPs reduction versus ViT-Base/SAM backbones, with competitive or superior mask AP and box AP on remote sensing benchmarks. On SSDD and WHU, RS4D matches or surpasses ViT-based foundation models and high-quality CNN baselines (Mask R-CNN, Cascade Mask RCNN, HQ-ISNet). Specifically, RS4D(box) obtains 70.7 AP_bbox and 70.0 AP_mask, outperforming SAM-seg and RSPrompter with only a quarter of their parameters. Efficiency metrics show RS4D matches MobileNetV2/ResNet18 in compute cost while exhibiting state-of-the-art dense prediction performance.
On NWPU, RS4D remains competitive but slightly trails the most powerful foundation models, largely due to SSM’s tendency to compress small or adjacent instances during HiPPO-driven memory projection under limited downstream training. Increasing distillation volume from SA-1B directly correlates with improvements in instance segmentation accuracy, particularly for data-scarce, category-diverse tasks.
Frozen backbone tests validate the transferability of distilled SSM representations, as strong mask AP is achieved even when only neck/head layers are adjusted.
Ablation and Analysis
- Knowledge Distillation: Distilled SSM weights yield faster and more stable convergence versus random initialization.
- Neck Fusion: Optimal segmentation performance is attained by fusing multiple intermediate SSM hidden layers; shallow-only fusion drastically degrades accuracy.
- Noise/Masking: 20% perturbation optimizes robustness; excessive masking reduces performance.
- Backbone Depth/Width: Moderate increases improve AP but risk overfitting beyond 24 layers.
- Distillation Volume: Larger SA-1B subsets boost generalization for challenging downstream datasets.
- Frozen Backbone: Retaining pretrained SSM weights still delivers strong transfer, showing efficient representation learning.
Discussion
RS4D demonstrates that linear-time SSM backbones, when distilled from transformer-based foundation encoders, achieve high-quality instance segmentation in remote sensing with orders-of-magnitude gains in compute efficiency, parameter count, and GPU memory usage. This is particularly salient for edge devices and real-time contexts (e.g., in-orbit image processing), which are otherwise infeasible for transformer-heavy models. The bidirectional SSM scanning mechanism provides superior spatial aggregation compared to unidirectional or shallow SSM designs.
Limitations include the tendency of memory compression mechanisms (HiPPO) to lose information about small, densely packed objects, evident on challenging datasets like NWPU. Potential mitigation strategies involve expanding embedding dimensionality, refining local feature integration, and optimizing memory mechanisms to improve small-instance recall.
Implications and Future Directions
RS4D’s architecture paves the way for scalable, efficient deployment of foundation models on large remote sensing imagery. The distillation framework is readily extensible to other vision tasks, foundation encoder architectures, and domain adaptation contexts. Ongoing advances in state space modeling, memory gating, and domain-aware distillation may further accelerate dense prediction and unlock new capabilities for resource-constrained AI in remote sensing and beyond. The framework’s compact design could facilitate the transfer of priors from much larger, promptable foundation models with billions of parameters, enabling efficient storage and rapid inference for real-world applications.
Conclusion
RS4D presents a principled and efficient solution for remote sensing instance segmentation, distilling dense-prior transformer-based foundation model knowledge into lightweight, linear-time state space architectures. The approach achieves competitive dense prediction performance with dramatically reduced computational overhead, making it highly suitable for scalable deployment in resource-limited settings. Its generalizability and efficiency warrant further investigation for other vision tasks and foundational model transfer paradigms.