Mamba: Linear-Time Sequence Modeling with Selective State Spaces 論文分析

# 従来SSM：固定パラメータ
B = nn.Parameter(torch.randn(N))  # 学習されるが入力に依存しない

# Mamba：入力依存パラメータ
B = nn.Linear(D, N)(x)      # xに応じてBが変化
C = nn.Linear(D, N)(x)      # xに応じてCが変化
delta = nn.Linear(D, D)(x)  # xに応じてstep sizeも変化

┌────────────────────┐
│   SRAM (20MB)      │  ← 非常に高速、非常に小さい
├────────────────────┤
│   HBM (80GB)       │  ← 低速、大容量
└────────────────────┘

# 簡略化したSelective Scan（概念コード）
def selective_scan(x, A, B, C, delta):
    """
    x: (batch, length, d_model)
    Returns: (batch, length, d_model)
    """
    batch, length, d = x.shape
    n = A.shape[1]  # state dimension

    # Discretize A, B using delta
    deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, D, N)
    deltaB_x = delta.unsqueeze(-1) * B.unsqueeze(2) * x.unsqueeze(-1)

    # Sequential scan (実際にはCUDAカーネルで並列化)
    h = torch.zeros(batch, d, n, device=x.device)
    ys = []
    for i in range(length):
        h = deltaA[:, i] * h + deltaB_x[:, i]
        y = (h * C[:, i].unsqueeze(1)).sum(-1)
        ys.append(y)

    return torch.stack(ys, dim=1)

Input (x)
    │
    ├──── Linear projection (expand) ────┐
    │                                     │
    ▼                                     ▼
  Conv1D                              SiLU (gate)
    │                                     │
    ▼                                     │
  SiLU                                    │
    │                                     │
    ▼                                     │
  SSM (selective scan)                    │
    │                                     │
    ▼                                     │
    ×─────────────────────────────────────┘
    │
    ▼
  Linear projection (reduce)
    │
    ▼
  Output

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = d_model * expand

        # Input projection
        self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False)

        # Conv1D
        self.conv1d = nn.Conv1d(
            d_inner, d_inner, d_conv,
            padding=d_conv - 1, groups=d_inner
        )

        # SSM parameters
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1, bias=False)  # B, C, delta
        self.dt_proj = nn.Linear(1, d_inner, bias=True)

        # A parameter (structured)
        A = torch.arange(1, d_state + 1).float().repeat(d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))

        self.D = nn.Parameter(torch.ones(d_inner))
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)

    def forward(self, x):
        batch, length, _ = x.shape

        # Dual path
        xz = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)

        # Conv + activation
        x = self.conv1d(x.transpose(1, 2))[:, :, :length].transpose(1, 2)
        x = F.silu(x)

        # SSM
        y = self.ssm(x)

        # Gate and project
        y = y * F.silu(z)
        return self.out_proj(y)