SMA-Attention

Symmetric Mamba Attention — a dual-path attention module for multi-scale feature fusion.

Unlike standard attention where query and key come from the same source, SMA takes two distinct inputs (e.g., different temporal scales, modalities, or branches) and produces a single fused output. Both Q and K are generated through a shared Mamba state-space model — the symmetric design ensures neither input path is privileged.

Key Features

• Symmetric Q/K via Mamba: A single shared Mamba layer processes both inputs. The selective state-space mechanism captures long-range dependencies that standard projections miss.
• Fused Value Projection: V1 and V2 are merged via concat(V1+V2, V1*V2) — capturing both additive and multiplicative cross-path interactions.
• Sigma-Modulated Attention: Element-wise gating sigma(Q) * sigma(K)^T modulates attention logits, enabling dimension-aware soft-thresholding.
• Multiple Residual Paths: Residual connections at Q/K, V fusion, and final output ensure stable gradient flow.

Installation

pip install sma-attention

Note: This package depends on mamba-ssm, which requires a CUDA-compatible GPU and CUDA toolkit. See mamba-ssm installation guide for details.

Quick Start

import torch
from sma_attention import SymmetricMambaAttention

# B=batch, D=feature_dim, N=sequence_length
x1 = torch.randn(2, 128, 60)   # e.g., high-frequency component
x2 = torch.randn(2, 128, 60)   # e.g., low-frequency component

sma = SymmetricMambaAttention(d_model=128, n_heads=4)
output = sma(x1, x2)            # shape: (2, 128, 60)

API Reference

SymmetricMambaAttention

SymmetricMambaAttention(
    d_model: int,        # feature dimension
    n_heads: int,        # number of attention heads
    d_state: int = 16,   # Mamba state dimension
    d_conv: int = 4,     # Mamba conv kernel size
    expand: int = 2,     # Mamba inner expansion factor
)

forward(x1, x2)

Architecture

The figure above illustrates the processing flow:

1. Inputs: Two paths x₁ and x₂ (e.g., adjacent scales in a circular topology), both shape (B, D, N)
2. Symmetric Q/K: Shared Mamba SSM with residual connections generates query from x₁ and key from x₂
3. Dual-Path Value: Independent linear projections → fusion via concat(sum, product) → linear projection
4. Multi-Head Reshape: Standard head splitting for attention computation
5. Sigma Modulation: σ(Q)·σ(K)ᵀ element-wise gating modulates attention logits before softmax
6. Output: Weighted V aggregation + average residual from both inputs

Requirements

• Python >= 3.10
• PyTorch >= 2.0
• einops >= 0.6
• mamba-ssm >= 2.0 (CUDA required)

License

MIT