Unmixing-Guided Spatial-Spectral Mamba with Clustering Tokens for Hyperspectral Image Classification

1. Problem Statement

Hyperspectral Image (HSI) classification faces three primary challenges:

• Spectral-Mixture Effect: Pixels often contain mixed materials, making it difficult to identify pure spectral signatures.
• Spatial-Spectral Heterogeneity: HSI data exhibits complex variations across both spatial and spectral domains.
• Boundary and Detail Preservation: Existing models struggle to preserve fine-scale details and class boundaries.

While deep learning models like CNNs and Transformers have been applied, they face limitations:

• CNNs: Limited in modeling long-range spatial-spectral dependencies.
• Transformers: High computational cost and risk of overfitting due to quadratic complexity in attention mechanisms.
• Mamba Models (State-Space Models): While efficient with linear complexity, standard Mamba models suffer from correlation decay when processing an entire HSI image as a single long token sequence. This approach is inefficient for capturing subtle patterns in highly heterogeneous images.

2. Methodology

The authors propose a novel framework that integrates spectral unmixing with a Mamba-based architecture to guide token selection and sequencing.

A. Model Overview

The architecture consists of two parallel branches (Unmixing and Classification) connected via Abundance Maps.

1. Patch Embedding: The input HSI cube is processed into feature maps.
2. Dual Branches:
   • Unmixing Branch: Learns endmembers and abundance maps.
   • Classification Branch: Learns discriminative features for class labels.
3. Token Selection: Abundance maps are used to define clusters, from which Top-K tokens are adaptively selected for the Mamba blocks.
4. Fusion: Features from both branches are fused to produce the final classification map, abundance maps, and endmember variability maps.

B. Key Components

1. Semi-Blind Unmixing Network with Endmember Variability

• Abundance Extraction: Uses a 1×1 convolutional network to estimate abundance maps ($A$) while preserving spatial structure.
• Endmember Variability Modeling: Instead of assuming fixed endmembers, the model initializes endmembers using Vertex Component Analysis (VCA) and repeats them into $R$ global variants. It then learns pixel-wise variability weights to dynamically reconstruct endmembers as a convex combination of these variants. This addresses the issue where the same material appears differently under varying conditions.
• Reconstruction: The observed spectrum is modeled as a weighted sum of abundance and variable endmembers.

2. Adaptive Top-K Token Selection

• Motivation: To avoid processing the entire image as one long sequence (which causes correlation decay), the model selects only the most informative tokens based on abundance clusters.
• Temporal Abundance: Uses an Exponential Moving Average (EMA) of abundance maps across training epochs to stabilize token selection.
• Selection Strategy: The number of tokens ($K_p$) for each abundance cluster $p$ is calculated based on the peak and average abundance values. This ensures the model focuses on high-abundance regions while maintaining spatial diversity.
• Complexity Reduction: Reduces complexity from $O(L \cdot d \cdot s)$ (where $L$ is total pixels) to $O(\sum K_p \cdot d \cdot s)$, where $\sum K_p \ll L$.

3. Unmixing-Guided Spatial-Spectral Mamba Module

• Parallel Mamba Blocks: The model employs $P+1$ parallel Mamba blocks (where $P$ is the number of endmembers).
  • $P$ blocks process Abundance-Specific Spatial Tokens (selected from the classification feature map based on the $p$-th abundance map).
  • $P$ blocks process Abundance-Specific Spectral Tokens (reshaped to learn grouped spectral features).
  • 1 additional block handles unselected (sparse) pixels.
• Sequencing: Tokens are sorted by abundance value (high to low) before being fed into the Mamba layers, creating a non-local, sparse, and adaptive sequence.

4. Multi-Task Learning

• The model is supervised by a combined loss function:
  $$L_{total} = L_{CE} + 0.01 L_{unmix} + 0.001 ||A||_{1/2}$$
  • $L_{CE}$: Cross-entropy loss for classification.
  • $L_{unmix}$: Spectral Angle Distance (SAD) for reconstruction accuracy.
  • $||A||_{1/2}$: Non-convex sparse penalty on abundance maps.

3. Key Contributions

1. Novel Spectral Unmixing Network: Automatically learns endmembers and abundance maps while explicitly modeling endmember variability, disentangling the spectral mixture effect.
2. Top-K Token Selection Strategy: Introduces an efficient, adaptive strategy to sequence tokens based on abundance clusters, solving the correlation decay issue in standard Mamba models and reducing computational cost.
3. Unmixing-Guided Spatial-Spectral Mamba: A new module that learns relationships among selected sparse tokens in both spatial and spectral domains, significantly improving feature learning and detail preservation.
4. Unified Multi-Task Framework: Simultaneously optimizes for classification, unmixing, and variability modeling, outputting accurate classification maps alongside comprehensive spectral libraries and abundance maps.

4. Experimental Results

The model was evaluated on four datasets: QUH-Tangdaowan, Pavia University, Houston, and Fanglu Tea Farm.

• Performance: The proposed method ("Ours") consistently outperformed state-of-the-art methods (including RF, SSRN, SS-ConvNeXt, Transformers, and other Mamba variants like MambaHSI, SFMamba, SDMamba) in Overall Accuracy (OA), Average Accuracy (AA), and Kappa coefficients.
  • Example: On the QUH-Tangdaowan dataset, it achieved 98.25% OA, surpassing the second-best (MambaHSI at 97.76%).
• Qualitative Analysis:
  • Boundary Preservation: The model produced sharper class boundaries and preserved fine details (e.g., single trees, building edges, roads) better than patch-based or downsampling-heavy models.
  • Abundance Estimation: The unmixing branch provided distinct and physically meaningful abundance maps, capturing subtle variations (e.g., shadows under trees) that other models missed.
• Ablation Study: Removing components like positional coding, Top-K selection, or endmember variability resulted in significant performance drops, confirming the necessity of each module.
• Endmember Number Sensitivity: Experiments showed that setting the number of endmembers ($P$) equal to the number of classes did not yield the best results; a higher $P$ (accounting for variability) improved performance.

5. Significance

This paper bridges the gap between physical spectral unmixing and modern deep learning architectures (Mamba).

• Efficiency: By using unmixing to guide token selection, it overcomes the computational bottlenecks and correlation decay issues of applying Mamba to full HSI cubes.
• Interpretability: The framework provides not just a classification label but also physical insights (abundance maps and spectral libraries), which are crucial for environmental monitoring and resource exploration.
• Robustness: The explicit modeling of endmember variability makes the model more robust to real-world spectral variations, leading to superior generalization on complex, large-scale datasets.

The code is publicly available, facilitating further research in unmixing-guided deep learning for remote sensing.