Metadata-Version: 2.4
Name: bidirectional-mamba
Version: 0.2.0
Summary: Bidirectional Mamba-2 SSD block with SwiGLU FFN (Vim/VMamba-style merge).
Author: redEddie
License-Expression: MIT
Project-URL: Repository, https://github.com/redEddie/bidirectional-mamba
Keywords: mamba,ssm,bidirectional,vision-mamba,swiglu
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: einops>=0.4
Requires-Dist: mamba-ssm>=2.0
Requires-Dist: causal-conv1d>=1.2
Requires-Dist: matplotlib>=3.0
Dynamic: license-file

# bidirectional-mamba

양방향 Mamba-2 SSD 블록에 SwiGLU FFN을 결합한 Vision-Mamba/VMamba 스타일 구현체입니다.  
forward scan과 backward scan을 독립적으로 수행하고, 출력을 더한 뒤 gate와 out projection을 공유합니다.

```
zxbcdt    = in_proj(u)                  # [z | x | B | C | dt | dt_b], 하나의 projection
y_fw      = SSD(conv  (x),     A,   D,   dt,       dt_bias  )
y_bw      = SSD(conv_b(x_rev), A_b, D_b, dt_b_rev, dt_bias_b)_rev   # reversed
y         = y_fw + y_bw                 # sum merge
out       = out_proj(RMSNormGated(y, z))
```

블록 구조는 `pre-norm` 기반이며 `(mixer residual) + (SwiGLU FFN residual)`입니다.  
스택 마지막에는 final RMSNorm이 붙습니다.

## 설치

CUDA + Triton 환경이 필요합니다 (`mamba-ssm`, `causal-conv1d` 필요). CPU-only는 지원하지 않습니다.

```bash
pip install torch einops matplotlib
pip install mamba-ssm causal-conv1d --no-build-isolation
pip install -e .
```

## 기본 사용법

```python
import torch
from bidirectional_mamba import BiMambaStack

stack = BiMambaStack(d_model=768, n_layers=4, d_state=64, headdim=64,
                     chunk_size=49).cuda()

x = torch.randn(2, 196, 768, device="cuda")   # [B, L, D]
y = stack(x)                                  # [B, L, D]
```

## 주요 기능

### 1. 방향별 분리 (`share_dt`, `share_bc`)

기본적으로 forward/backward는 각자의 `A`, `D`, `dt_bias`, `conv1d`를 가집니다.  
추가로 `dt` 투영과 `B/C` 투영도 방향별로 분리할 수 있습니다.

```python
# dt 분리: backward scan이 자신만의 dt를 가짐
BiMambaStack(d_model=768, n_layers=4, share_dt=False)

# B/C 분리: backward scan이 자신만의 B, C를 가짐
BiMambaStack(d_model=768, n_layers=4, share_dt=False, share_bc=False)
```

### 2. Padding mask (`valid_mask`)

가변 길이 배치를 다룰 때 반드시 `valid_mask`를 전달하세요.  
padding이 없으면 backward scan이 padding 영역에서 시작해 실제 토큰까지 영향을 미칩니다.

```python
mask = torch.zeros(B, L, dtype=torch.bool, device="cuda")
mask[i, :lengths[i]] = True

y = stack(x, valid_mask=mask)
```

mask를 쓰면 padding 위치는 0을 출력하고 실제 토큰에는 영향을 주지 않습니다.

### 3. 학습 건강 지표 (`return_stats`)

`BiMambaStack`, `BiMambaBlock`, `BiMambaMixer` 모두 `return_stats=True`를 지원합니다.  
loss 외에도 dt, A decay, B/C norm, forward/backward 기여도 등 Mamba 고유의 수치를 확인할 수 있습니다.

```python
out, stats = stack(x, return_stats=True)

# stats 예시
# {
#   "layer0/mixer/dt/effective_mean": ...,
#   "layer0/mixer/A/decay_mean": ...,
#   "layer0/mixer/y_fw_bw_ratio": ...,
#   ...
# }
```

자세한 예제는 `benchmarks/log_mamba_stats.py`를 참고하세요.

### 5. 구현 검증

`benchmarks/verify_bidirectional.py`는 양방향 Mamba 구현이 제대로 되었는지 확인하는 독립 스크립트입니다.  
flip equivariance, backward branch 사용 여부, explicit reference match, gradient flow, valid mask 등 8가지 기준을 검사합니다.

```bash
python benchmarks/verify_bidirectional.py
```

### 4. SSD token-mixing 행렬 (`mixing_matrix`)

Mamba의 암묵적인 token-mixing 행렬 L을 attention map처럼 시각화할 수 있습니다.  
forward scan은 하삼각, backward scan은 상삼각을 채웁니다.

```python
mixer = BiMambaMixer(d_model=192, d_state=32, headdim=32).cuda()
L = mixer.mixing_matrix(x)  # [B, L, L]
```

시각화 예제는 `benchmarks/visualize_mixing_matrix.py`를 참고하세요.

## CIFAR-10 벤치마크 결과

**공통 설정**: patch_size=4, share_dt=False, share_bc=True, use_mlp=True, epochs=30, seed=0

### Width scaling (n_layers=8 고정)

| d_model | n_layers | headdim | params(추정) | test acc |
|--------:|---------:|--------:|------------:|---------:|
| 768     | 8        | 64      | ~67M        | 80.70%   |
| 1024    | 8        | 128     | ~119M       | 81.76%   |
| 4096    | 4        | 128     | ~944M       | 80.13%   |

### Depth scaling (d_model=384 고정, LR per-depth scaling)

| d_model | n_layers | headdim | params(추정) | test acc |
|--------:|---------:|--------:|------------:|---------:|
| 384     | 4        | 32      | ~8.5M       | 84.39%   |
| 384     | 8        | 32      | ~17M        | 85.00%   |
| 384     | 12       | 32      | ~25M        | 85.04%   |
| 384     | 16       | 32      | ~34M        | 85.23%   |
| 384     | 20       | 32      | ~42M        | 84.65%   |

> d_model=384이 768/1024 width scaling보다 훨씬 적은 파라미터로 더 높은 성능을 냈습니다.  
> 이는 현재 width scaling 설정(특히 4096)이 CIFAR-10에 비해 과도하게 크다는 것을 의미합니다.

## 공개 API

- `BiMambaMixer` — 양방향 SSD mixer (shared in/out projection, 방향별 conv + A + D + dt + dt_bias).
- `BiMambaBlock` — pre-norm block: `mixer + SwiGLU FFN` residual.
- `BiMambaStack` — `n`개 block + final RMSNorm.
- `SwiGLUMLP` — 독립형 SwiGLU FFN (expand=8/3 기본).

세 클래스 모두 `share_dt`, `share_bc`를 받고, `forward`는 `valid_mask`를 받습니다.  
`BiMambaMixer/Block/Stack`은 추가로 `return_stats=True`를 받을 수 있습니다.

## 참고

- `A`/`A_b`, `D`/`D_b`, `dt_bias`/`dt_bias_b`, 두 `conv1d`, 그리고 (`share_dt=False`일 때) 두 `dt` 투영은 모두 묶여 있지 않습니다.
- `in_proj`, `RMSNormGated` (z로 gate), `out_proj`는 공유됩니다.  
  출력을 더한 뒤 gate를 적용하는 것은 `norm_before_gate=False`이므로 방향별로 따로 gate하는 것과 동일합니다.
- `share_dt=True`는 이전 동작으로 돌아가 forward `dt`를 reversed해서 backward가 재사용합니다.  
  softplus가 단조증가이고 per-head bias만 분리되므로 두 방향이 "어떤 토큰이 중요한지"에 대해 강제로 동의해야 합니다.
- autoregressive 디코딩용 `step()`은 없습니다. 양방향 scan은 전체 시퀀스가 필요합니다.
- `seq_idx` 기반 세그먼트 경계는 아직 구현되지 않았습니다.

## 참고 문헌

- **Vision Mamba (Vim)** — Zhu et al., *Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model*, 2024. [arXiv:2401.09417](https://arxiv.org/abs/2401.09417)
- **VMamba** — Liu et al., *VMamba: Visual State Space Model*, 2024. [arXiv:2401.10166](https://arxiv.org/abs/2401.10166)
- **SwiGLU FFN** — Shazeer, *GLU Variants Improve Transformer*, 2020. [arXiv:2002.05202](https://arxiv.org/abs/2002.05202)
- **Mamba-2 / SSD** — Dao & Gu, *Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality*, 2024. [arXiv:2405.21060](https://arxiv.org/abs/2405.21060)
