Cosine-Normalized Attention for Hyperspectral Image Classification

1. Problem Statement

Hyperspectral Image Classification (HSIC) faces two primary challenges:

• Geometric Mismatch in Attention: Standard Transformer models rely on scaled dot-product attention, which computes similarity based on the inner product of query ($Q$) and key ($K$) vectors. This metric conflates vector magnitude and orientation. In hyperspectral data, spectral signatures often vary in magnitude due to illumination changes, atmospheric effects, or sensor noise, even when the underlying material identity remains the same. Consequently, dot-product attention may prioritize magnitude over the more discriminative angular direction of the spectral signature.
• Label Scarcity: Pixel-level annotation for hyperspectral images is expensive and labor-intensive. Models often suffer from overfitting when trained on extremely limited labeled data (e.g., 1% of pixels), especially when using complex architectures with large parameter counts.

2. Methodology

The authors propose a Cosine-Normalized Attention mechanism integrated into a lightweight Spatial-Spectral Transformer framework.

A. Core Innovation: Cosine-Normalized Attention

Instead of standard dot-product similarity, the method reformulates attention scoring from a geometric perspective:

1. $\ell_2$ Normalization: Query ($Q$) and Key ($K$) embeddings are projected onto a unit hypersphere by normalizing them:
   $$\tilde{q} = \frac{q}{\|q\|_2}, \quad \tilde{k} = \frac{k}{\|k\|_2}$$
   This ensures that the similarity score depends only on the angle between vectors, ignoring magnitude variations.
2. Squared Cosine Similarity: The similarity score is computed as the element-wise square of the cosine similarity:
   $$S = (\tilde{q}^\top \tilde{k})^2 = \cos^2(\theta)$$
   This "sharpening" operation further emphasizes strong angular alignment while suppressing weak alignments, creating a more concentrated attention distribution.

B. Architecture

• Patch Tokenization: The HSI cube is divided into overlapping 3D patches ($P \times P \times C$). Spectral vectors within patches are projected into an embedding space and treated as tokens.
• Transformer Backbone: A lightweight Spatial-Spectral Transformer is used with:
  • Embedding dimension: 64
  • Depth: 4 layers
  • Attention heads: 4
  • MLP dimension: 128
• Training Regime: The model is evaluated under a 1% training / 1% validation / 98% testing split to simulate extreme label scarcity.

C. Controlled Analysis

The study evaluates 11 different attention score variants within the same backbone to isolate the effect of the similarity function. Variants include:

• Cosine-squared (Proposed)
• Cosine similarity
• Scaled dot-product (Standard)
• Additive attention
• Cross-attention variants (where $Q$ and $K$ come from different streams)

3. Key Contributions

1. Geometric Re-alignment: The paper identifies and addresses the mismatch between Euclidean dot-product attention and the angular nature of hyperspectral signatures. It proposes that angular consistency is a more robust inductive bias for HSIC than magnitude-inclusive metrics.
2. Lightweight Performance: The proposed method achieves state-of-the-art results using a significantly smaller and simpler architecture compared to recent complex Transformer and Mamba-based models.
3. Comprehensive Ablation: Through a controlled study, the authors demonstrate that the choice of similarity function is a more critical factor for performance than increasing model depth, width, or complexity.
4. Robustness: The method shows superior stability against spectral noise and magnitude variations, which are common in remote sensing data.

4. Experimental Results

The method was tested on three benchmark datasets: Salinas (SA), WHU Hi HongHu (HH), and QUH-Tangdaowan (TD).

• Performance Metrics: The proposed Cosine-Squared (CS2) variant consistently achieved top-tier performance in Overall Accuracy (OA), Kappa ($\kappa$), and Average Accuracy (AA).
  • Example (Salinas): CS2 achieved 99.18% OA and 99.33% Kappa, outperforming recent Mamba-based models (e.g., SSAM, DBMLLA) and other Transformers (MSST, S2CIFT).
• Comparison: The approach outperformed or matched complex baselines (including WaveFormer, DiffFormer, and various Mamba variants) despite using a much lighter backbone.
• Visual Quality: Classification maps generated by the proposed method showed more homogeneous regions, fewer isolated errors, and cleaner class boundaries, particularly in spectrally confusing areas.
• Ablation Insights:
  • Normalizing both Query and Key yielded the best results.
  • Squared cosine outperformed standard cosine and dot-product variants.
  • Architectural changes (e.g., increasing depth or heads) had a marginal impact compared to the switch to cosine-based scoring.
  • The model demonstrated graceful degradation under increasing spectral noise (10dB to 30dB), confirming its robustness.

5. Significance and Conclusion

This work fundamentally shifts the focus in HSIC from designing increasingly complex architectures to optimizing the geometric properties of the attention mechanism.

• Inductive Bias: It establishes that cosine-based scoring provides a crucial inductive bias for hyperspectral representation learning, aligning the model's internal geometry with the physical properties of spectral data.
• Efficiency: It proves that for label-scarce scenarios, improving the similarity function is more effective than simply scaling up model parameters.
• Generalizability: The findings suggest that attention mechanisms in remote sensing should prioritize angular relationships over Euclidean distances to handle illumination and sensor variability effectively.

In summary, the paper offers a simple, interpretable, and highly effective alternative to standard dot-product attention, demonstrating that geometry-aware scoring is a key driver of performance in hyperspectral image classification.