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는 지원하지 않습니다.

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

기본 사용법

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 투영도 방향별로 분리할 수 있습니다.

# 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 영역에서 시작해 실제 토큰까지 영향을 미칩니다.

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 고유의 수치를 확인할 수 있습니다.

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가지 기준을 검사합니다.

python benchmarks/verify_bidirectional.py

4. SSD token-mixing 행렬 (mixing_matrix)

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

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 고정)

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

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
• VMamba — Liu et al., VMamba: Visual State Space Model, 2024. arXiv:2401.10166
• SwiGLU FFN — Shazeer, GLU Variants Improve Transformer, 2020. arXiv:2002.05202
• Mamba-2 / SSD — Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, 2024. arXiv:2405.21060