sma-attention 0.1.1

Symmetric Mamba Attention: Dual-path attention with Mamba-enhanced Q/K generation and sigma-modulated attention weighting

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
)

Parameter	Description
d_model	Feature dimension (must be divisible by n_heads)
n_heads	Number of attention heads
d_state	Mamba SSM state expansion factor
d_conv	Mamba 1D convolution kernel width
expand	Inner dimension multiplier for Mamba blocks

forward(x1, x2)

Argument	Shape	Description
x1	(B, D, N)	First input (e.g., one temporal scale)
x2	(B, D, N)	Second input (e.g., another temporal scale)
Returns	(B, D, N)	Fused output

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