Metadata-Version: 2.4
Name: sma-attention
Version: 0.1.1
Summary: Symmetric Mamba Attention: Dual-path attention with Mamba-enhanced Q/K generation and sigma-modulated attention weighting
Author: Wenqi Liao
License: MIT
Project-URL: Homepage, https://github.com/liaowenqi123/sma-attention
Project-URL: Repository, https://github.com/liaowenqi123/sma-attention
Keywords: deep-learning,attention,mamba,state-space-model,time-series,pytorch
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: einops>=0.6
Requires-Dist: mamba-ssm>=2.0
Dynamic: license-file

# 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

```bash
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](https://github.com/state-spaces/mamba) for details.

## Quick Start

```python
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`

```python
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

![SMA Architecture](https://raw.githubusercontent.com/liaowenqi123/sma-attention/master/Symmetric1.png)

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
