Mamba-RS-Engine: High-Throughput Inference for Satellite Change Detection
This project is an ongoing effort to build a high-throughput inference engine for remote sensing change detection, using State Space Models (SSMs) as the core architecture and C++/TensorRT for deployment. The target application is rapid identification of environmental change events — wildfires and flooding — from high-resolution satellite imagery.
The project serves two purposes: demonstrating practical competency in
inference optimisation and systems-level ML engineering, and as a genuine
upskilling exercise in an area I find technically interesting. Progress and
implementation details are documented on GitHub.
1. Motivation: Why Not a Vision Transformer?
The dominant approach to image understanding tasks over the past several years has been the Vision Transformer (ViT). ViTs work by dividing an image into patches and computing self-attention across all patch pairs — a powerful mechanism, but one with a fundamental scaling problem.
Self-attention in a standard ViT scales as O(N²) in the number of image patches. For a high-resolution satellite image divided into small patches, N is large — making full attention computationally prohibitive at the resolutions required for meaningful environmental monitoring.
Remote sensing imagery compounds this problem: scenes are large, the features of interest (a spreading fire front, a flood boundary) may span long spatial distances, and inference needs to be fast enough to be operationally useful. A model architecture that sidesteps the O(N²) bottleneck is therefore not just theoretically appealing — it's practically necessary.
2. Architecture: RS-Mamba and State Space Models
State Space Models (SSMs) — and specifically the Mamba architecture — offer a compelling alternative. Rather than computing pairwise attention, SSMs process sequences through a learned recurrent state, achieving near-linear scaling with sequence length while retaining the ability to capture long-range dependencies.
RS-Mamba adapts this to the 2D structure of satellite imagery using an
omnidirectional scanning strategy. Rather than processing image patches
in a single raster-order sequence (which imposes an arbitrary spatial ordering
on inherently 2D data), the model scans patches in multiple directions — horizontal,
vertical, and diagonal — and aggregates the resulting representations. This
allows the model to capture spatial context from all directions without the
quadratic cost of global attention.
For change detection specifically, the model takes two co-registered
satellite images of the same scene (before and after a potential change event)
and produces a pixel-level binary mask indicating where change has occurred.
The datasets used are LEVIR-CD and/or OSCD,
both standard benchmarks for remote sensing change detection.
3. Optimisation: C++ Inference Server with TensorRT
The second half of the project moves from model architecture to deployment. Once a trained RS-Mamba model is performing well in PyTorch, the goal is to export it and build a standalone C++ inference server optimised via NVIDIA TensorRT.
TensorRT applies a series of graph-level optimisations — layer fusion,
precision calibration (FP16/INT8), and kernel auto-tuning — to maximise
throughput on a target GPU. Combined with a lean C++ serving layer that avoids
the overhead of the Python runtime and PyTorch's general-purpose execution
graph, the target is a >2× throughput improvement over the
PyTorch baseline under equivalent hardware conditions.
This component of the project targets edge-constrained environments —
scenarios where inference needs to run on limited hardware in the field rather
than a cloud data centre, which is a realistic constraint for disaster response
and environmental monitoring applications.
4. Project Status and Roadmap
The project is currently in its early stages. The table below summarises the planned phases and their current status.
| Phase | Description | Status |
|---|---|---|
| 1. Literature Review | Review of State Space Models and the Mamba architecture, Vision Transformer limitations in remote sensing contexts, and existing change detection approaches. Primary references include the RS-Mamba and Mamba papers. | Complete |
| 2. Use Case & Dataset | Identification of target environmental monitoring use cases (wildfire and flood detection) and selection of appropriate benchmarks. Dataset access, loading pipelines, and preprocessing for LEVIR-CD and/or OSCD. | Complete |
| 3. PyTorch Implementation | Implement the RS-Mamba architecture in PyTorch and train on LEVIR-CD or OSCD. This serves as both the primary model and the throughput baseline for subsequent optimisation. | In Progress |
| 4. TensorRT Export | Export the trained model and apply TensorRT optimisations (layer fusion, FP16/INT8 calibration). Validate accuracy retention post-quantisation. | Planned |
| 5. C++ Inference Server | Build a standalone C++ inference server wrapping the TensorRT engine. Benchmark end-to-end throughput against the PyTorch baseline, targeting >2× improvement. | Planned |