Dynamic Memory Transformer for Hyperspectral Image Classification

1. Problem Statement

Hyperspectral Image Classification (HSIC) faces three primary challenges:

• High Dimensionality & Limited Data: The "Hughes phenomenon" occurs where the number of spectral bands far exceeds the number of labeled samples, leading to overfitting and poor generalization.
• Complex Spatial-Spectral Dependencies: Effective classification requires modeling both fine-grained local textures and long-range contextual relationships simultaneously.
• Limitations of Existing Models:
  • CNNs: Struggle to capture long-range dependencies due to limited receptive fields.
  • Standard Transformers: While capable of modeling long-range dependencies, they often introduce redundant attention patterns, suffer from quadratic computational complexity, and require massive pre-training or large datasets, making them inefficient for HSIC with limited labels.
  • Hybrid Models: Often increase architectural complexity and training time without guaranteeing consistent fine-grained feature capture.

2. Proposed Methodology: MemFormer

The authors propose MemFormer, a lightweight transformer architecture designed to balance representational power with computational efficiency. The core components are:

A. Input Preprocessing

• The input HSI ($H \times W \times S$) is split into non-overlapping patches ($W_s \times W_s \times S$).
• A shallow convolutional layer projects these patches into feature embeddings ($Z \in \mathbb{R}^{N \times K}$).

B. Spatial-Spectral Positional Embedding (SSPE)

Unlike standard transformers that rely on learnable or purely sinusoidal positional encodings, MemFormer introduces a domain-aware SSPE to jointly encode:

• Spatial Continuity: Using sinusoidal functions based on 2D patch coordinates $(x_i, y_i)$.
• Spectral Ordering: Using sinusoidal functions based on spectral band indices ($s_j$).
• Mechanism: These encodings are projected to the embedding dimension and concatenated, ensuring the model understands the structural integrity of the hyperspectral cube without heavy convolutional pre-processing.

C. Dynamic Memory-Enhanced Attention

This is the core innovation. Instead of standard Multi-Head Self-Attention (MHSA) which computes attention solely between current tokens, MemFormer integrates a Global Dynamic Memory Module:

• Memory Structure: A compact, non-trainable memory bank ($M \in \mathbb{R}^{M \times K}$) is maintained.
• FIFO Policy: The memory is updated dynamically per batch using a First-In-First-Out (FIFO) policy. It discards the oldest entries and appends new aggregated features, ensuring the model retains recent global context without accumulating noise.
• Attention Mechanism:
  • Queries ($Q$): Derived from the current input tokens.
  • Keys ($K_m$) & Values ($V_m$): Derived exclusively from the dynamic memory.
  • Computation: The attention score is calculated as $Softmax(\frac{Q K_m^T}{\sqrt{K}}) V_m$. This allows current tokens to query the aggregated global context stored in memory.
• Update Rule: After each batch, the memory is updated by averaging the attention responses of the current batch, refining the global representation progressively across layers.

D. Classification

A classification token ($z_{CLS}$) is prepended to the sequence. The final representation is processed through a Feed-Forward Network (FFN) and Layer Normalization, with the final class prediction derived from the $z_{CLS}$ token.

3. Key Contributions

1. Dynamic Memory-Enhanced Attention: A novel mechanism that reduces attention redundancy and computational cost by decoupling queries (current input) from keys/values (global memory), enabling efficient long-range dependency modeling.
2. Spatial-Spectral Positional Embedding (SSPE): A specialized encoding scheme that explicitly preserves spatial continuity and spectral ordering, eliminating the need for convolution-heavy positional encodings.
3. Lightweight Architecture: MemFormer achieves state-of-the-art performance with significantly fewer parameters compared to heavy hybrid or hierarchical transformer models.

4. Experimental Results

The model was evaluated on three benchmark datasets: Indian Pines (IP), WHU-Hi-HanChuan (HC), and WHU-Hi-HongHu (HH).

• Performance Metrics:
  • Indian Pines: Achieved 99.55% Overall Accuracy (OA), 99.38% Average Accuracy (AA), and 99.49% Kappa ($\kappa$). This outperformed 2D/3D CNNs, Hybrid CNNs, and other Transformers (e.g., PyFormer, WaveFormer).
  • WHU-Hi-HanChuan: Achieved 99.27% OA and 98.67% AA.
  • WHU-Hi-HongHu: Achieved 98.51% OA and 97.26% AA.
• Efficiency:
  • MemFormer is highly parameter-efficient (e.g., ~0.82M parameters on HH dataset), significantly lower than baselines like Pyramid Transformer (PF) or 3D CNNs.
  • Training and inference times are competitive, often faster than complex hybrid models.
• Ablation Studies:
  • Memory vs. Standard Attention: Adding the memory module improved $\kappa$ by >4% on IP and ~2% on HC/HH, proving its value in retaining context.
  • Positional Encoding: SSPE consistently outperformed "None," "Learnable," and standard "Sinusoidal" encodings.
  • Memory Size: An optimal memory length of ~10 tokens was found; too small (1 token) lacked context, while too large (25-30) introduced redundancy.

5. Significance

• Solving the Data Scarcity Issue: MemFormer demonstrates that effective HSIC can be achieved without massive pre-training or huge labeled datasets by leveraging dynamic memory to aggregate context efficiently.
• Efficiency-Accuracy Trade-off: It breaks the traditional trade-off where high accuracy requires heavy computational cost. MemFormer achieves superior accuracy with a lightweight footprint, making it suitable for real-world deployment on edge devices or resource-constrained environments.
• Structural Integrity: The SSPE approach provides a robust way to handle the unique spatial-spectral nature of hyperspectral data, offering a template for future domain-specific transformer designs.
• Generalizability: The model shows strong robustness across diverse datasets with varying spatial resolutions and spectral characteristics, addressing the challenge of cross-domain generalization.