Split View: Mamba와 상태 공간 모델(SSM) 완전 정복: Transformer를 넘어서

Attention matrix size: n × n
= 1,000 tokens   → 1,000,000 elements
= 10,000 tokens  → 100,000,000 elements
= 100,000 tokens → 10,000,000,000 elements (infeasible!)

x'(t) = A·x(t) + B·u(t)
y(t)  = C·x(t) + D·u(t)

A_bar = exp(Delta · A)
B_bar = (Delta · A)^(-1) · (exp(Delta · A) - I) · Delta · B

A_bar = (I - Delta/2 · A)^(-1) · (I + Delta/2 · A)
B_bar = (I - Delta/2 · A)^(-1) · Delta · B

x[t] = A_bar · x[t-1] + B_bar · u[t]
y[t] = C · x[t]

y[0] = C · B_bar · u[0]
y[1] = C · A_bar · B_bar · u[0] + C · B_bar · u[1]
y[2] = C · A_bar^2 · B_bar · u[0] + C · A_bar · B_bar · u[1] + C · B_bar · u[2]
...

K = (C·B_bar, C·A_bar·B_bar, C·A_bar^2·B_bar, ...)
y = K * u  (convolution)

A[n,k] = -sqrt((2n+1)(2k+1))  if n > k
A[n,k] = -(n+1)               if n == k
A[n,k] = 0                    if n < k

A = Λ - P·Q^T

import torch
import torch.nn as nn
import numpy as np

class S4Layer(nn.Module):
    """Simplified implementation of S4 layer"""
    def __init__(self, d_model, d_state=64, dropout=0.0):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state

        # HiPPO-LegS initialization
        A = self._make_hippo(d_state)
        # DPLR decomposition
        self.A = nn.Parameter(torch.tensor(A, dtype=torch.float32))
        self.B = nn.Parameter(torch.randn(d_state, 1) * 0.01)
        self.C = nn.Parameter(torch.randn(1, d_state) * 0.01)
        self.log_delta = nn.Parameter(torch.zeros(d_model))
        self.D = nn.Parameter(torch.ones(d_model))

    def _make_hippo(self, N):
        """Generate HiPPO-LegS matrix"""
        P = np.sqrt(1 + 2 * np.arange(N))
        A = P[:, np.newaxis] * P[np.newaxis, :]
        A = np.tril(A) - np.diag(np.arange(N))
        return -A

    def discretize(self, A, B, C, delta):
        """ZOH discretization"""
        dA = torch.matrix_exp(delta.unsqueeze(-1) * A)
        dB = torch.linalg.solve(A, (dA - torch.eye(A.shape[0])) @ B)
        return dA, dB, C

    def forward(self, u):
        # u: (batch, seq_len, d_model)
        delta = torch.exp(self.log_delta)
        # Discretization and convolution computation
        # Actual implementation is more complex, simplified here for concept
        return u

class H3Layer(nn.Module):
    """H3 Layer"""
    def __init__(self, d_model, d_state=64):
        super().__init__()
        # Two SSMs (shift SSM + diagonal SSM)
        self.shift_ssm = S4Layer(d_model, d_state)
        self.diag_ssm = S4Layer(d_model, d_state)

        # Projections
        self.Q_proj = nn.Linear(d_model, d_model)
        self.K_proj = nn.Linear(d_model, d_model)
        self.V_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        # Q, K, V projections
        q = self.Q_proj(x)
        k = self.K_proj(x)
        v = self.V_proj(x)

        # SSM 1: Apply Shift SSM to K
        k_ssm = self.shift_ssm(k)

        # Multiplicative gating: Q element-wise K_ssm
        gated = q * k_ssm

        # SSM 2: Apply Diagonal SSM to V, then multiply with gated
        v_ssm = self.diag_ssm(v)
        output = gated * v_ssm

        return self.out_proj(output)

B(t) = Linear_B(u(t))   # B that varies with input
C(t) = Linear_C(u(t))   # C that varies with input
Delta(t) = softplus(Linear_Delta(u(t)))  # Delta that varies with input

# Parallel Scan concept (actual implementation in CUDA)
def parallel_scan(gates, tokens):
    """
    Compute linear recurrence in parallel
    x[t] = gates[t] * x[t-1] + tokens[t]

    Achieves O(log n) parallel depth with binary tree structure
    """
    n = len(tokens)
    # Up-sweep phase
    log_n = int(np.log2(n))
    for d in range(log_n):
        step = 2 ** (d + 1)
        for i in range(step - 1, n, step):
            gates[i] = gates[i] * gates[i - 2**d]
            tokens[i] = gates[i - 2**d] * tokens[i - 2**d] + tokens[i]
    # Down-sweep phase (omitted)
    return tokens

Input x (B, L, D)
    │
    ├──────────────────────┐
    │                      │
  Linear(D→ED)          Linear(D→ED)
  + SiLU activation        │
    │                    SSM (S6)
    │                      │
    └─────── ⊙ ────────────┘
    (element-wise multiply)
    │
  Linear(ED→D)
    │
  Output y (B, L, D)

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat

class MambaBlock(nn.Module):
    """
    Mamba Block implementation
    Paper: "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
    """
    def __init__(
        self,
        d_model,       # Model dimension D
        d_state=16,    # SSM state dimension N
        d_conv=4,      # Convolution kernel size
        expand=2,      # Expansion ratio E
        dt_rank="auto",
        dt_min=0.001,
        dt_max=0.1,
        dt_init="random",
        dt_scale=1.0,
        dt_init_floor=1e-4,
        bias=False,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)

        if dt_rank == "auto":
            self.dt_rank = max(1, int(d_model / 16))
        else:
            self.dt_rank = dt_rank

        # Input projection (D → 2*ED, compute both branches at once)
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=bias)

        # Local convolution (depthwise conv)
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            bias=True,
            kernel_size=d_conv,
            groups=self.d_inner,
            padding=d_conv - 1,  # causal padding
        )

        # SSM parameter projections
        self.x_proj = nn.Linear(
            self.d_inner, self.dt_rank + self.d_state * 2, bias=False
        )
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)

        # A initialization (HiPPO-based)
        A = repeat(
            torch.arange(1, self.d_state + 1),
            "n -> d n",
            d=self.d_inner
        )
        self.A_log = nn.Parameter(torch.log(A))
        self.A_log._no_weight_decay = True

        # D (skip connection)
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.D._no_weight_decay = True

        # Output projection
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias)

    def forward(self, hidden_states):
        """
        hidden_states: (B, L, D)
        Returns: (B, L, D)
        """
        batch, seqlen, dim = hidden_states.shape

        # Input projection
        xz = self.in_proj(hidden_states)
        x, z = xz.chunk(2, dim=-1)  # each (B, L, ED)

        # Convolution (causal 1D conv)
        x = rearrange(x, "b l d -> b d l")
        x = self.conv1d(x)[:, :, :seqlen]  # causal trimming
        x = rearrange(x, "b d l -> b l d")
        x = F.silu(x)

        # SSM
        y = self.ssm(x)

        # Gating (element-wise multiply with z branch)
        y = y * F.silu(z)

        # Output projection
        output = self.out_proj(y)
        return output

    def ssm(self, x):
        """Selective SSM (S6) computation"""
        d_in, n = self.d_inner, self.d_state

        # A matrix (-exp(-A_log) ensures always negative = stability)
        A = -torch.exp(self.A_log.float())  # (ED, N)

        # x_proj: compute Delta, B, C
        x_dbl = self.x_proj(x)  # (B, L, dt_rank + 2N)
        delta, B, C = x_dbl.split(
            [self.dt_rank, n, n], dim=-1
        )

        # Delta: softplus ensures positive values
        delta = F.softplus(self.dt_proj(delta))  # (B, L, ED)

        # ZOH discretization and selective scan
        y = self.selective_scan(x, delta, A, B, C, self.D)
        return y

    def selective_scan(self, u, delta, A, B, C, D):
        """
        Selective scan algorithm
        In practice, uses CUDA kernels from mamba_ssm
        Here shown in pure PyTorch for conceptual understanding
        """
        b, l, d_in = u.shape
        n = A.shape[1]

        # Discretization
        # delta: (B, L, ED), A: (ED, N) → dA: (B, L, ED, N)
        deltaA = torch.exp(
            torch.einsum("bld,dn->bldn", delta, A)
        )
        # delta: (B, L, ED), B: (B, L, N), u: (B, L, ED)
        # → deltaB_u: (B, L, ED, N)
        deltaB_u = torch.einsum("bld,bln,bld->bldn", delta, B, u)

        # Recurrent scan
        x = torch.zeros(b, d_in, n, device=u.device, dtype=u.dtype)
        ys = []
        for i in range(l):
            x = deltaA[:, i] * x + deltaB_u[:, i]
            y = torch.einsum("bdn,bn->bd", x, C[:, i, :])
            ys.append(y)
        y = torch.stack(ys, dim=1)  # (B, L, ED)

        # D (skip connection)
        y = y + u * D

        return y


class MambaModel(nn.Module):
    """Sequence model composed of stacked Mamba blocks"""
    def __init__(self, d_model, n_layers, d_state=16, expand=2):
        super().__init__()
        self.layers = nn.ModuleList([
            MambaBlock(d_model, d_state=d_state, expand=expand)
            for _ in range(n_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x):
        """x: (B, L, D)"""
        for layer in self.layers:
            x = x + layer(self.norm(x))  # Pre-norm residual
        return x

# SSD operation concept
# SSM recurrence: x[t] = A[t] * x[t-1] + B[t] * u[t]
#                 y[t] = C[t]^T * x[t]
#
# Process in chunks:
# - Within chunk: parallel computation via matrix multiplication
# - Between chunks: recurrence for state propagation

class Mamba2Block(nn.Module):
    """Mamba 2 Block"""
    def __init__(self, d_model, d_state=64, n_heads=8, chunk_size=64):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.n_heads = n_heads
        self.chunk_size = chunk_size
        self.d_head = d_model // n_heads

        # Multi-head structure
        self.norm = nn.LayerNorm(d_model)
        self.in_proj = nn.Linear(d_model, d_model * 2 + d_state * 2 + n_heads)
        self.out_proj = nn.Linear(d_model, d_model)

        # A parameters (per head)
        self.A_log = nn.Parameter(torch.randn(n_heads))

    def forward(self, x):
        """x: (B, L, D)"""
        # Actual implementation uses CUDA kernels from mamba_ssm package
        return x

Layer 1: Mamba Block   (capture local patterns)
Layer 2: Attention     (capture global dependencies)
Layer 3: Mamba Block
Layer 4: Attention
...

Jamba = Mamba + Transformer + Mixture of Experts

Architecture:
- 52B parameters (active: 12B)
- Layer ratio: Attention 1 : Mamba 7
- MoE applied to some layers
- Supports 256K context window

# Jamba-style hybrid block (concept)
class JambaLayer(nn.Module):
    def __init__(self, d_model, layer_idx, attn_every_n=8):
        super().__init__()
        self.use_attention = (layer_idx % attn_every_n == 0)

        if self.use_attention:
            self.mixer = nn.MultiheadAttention(d_model, num_heads=8)
        else:
            self.mixer = MambaBlock(d_model)

        # MoE (only for some layers)
        self.use_moe = (layer_idx % 2 == 0)
        if self.use_moe:
            self.ffn = MixtureOfExperts(d_model, n_experts=16)
        else:
            self.ffn = nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.SiLU(),
                nn.Linear(d_model * 4, d_model)
            )

    def forward(self, x):
        x = x + self.mixer(x)
        x = x + self.ffn(x)
        return x

- Training: Parallelized like Transformer (matrix form)
- Inference: Recurrent like RNN (O(1) state)
- Token mixing without attention

Three computational paradigms:
1. Parallel: O(n^2) parallel computation during training (lower constant than Transformer)
2. Recurrent: O(1) memory during inference
3. Chunkwise Recurrent: balanced intermediate approach

Normalizing performance at sequence length 2K to 1x:

Sequence Length   Transformer   Mamba
1K                0.5x          0.5x
2K                1.0x          1.0x
4K                3.5x          2.0x    ← 1.75x faster than Transformer
8K                13x           4.0x    ← 3.25x faster than Transformer
16K               50x           8.0x    ← 6.25x faster than Transformer
100K              ~1800x        ~50x    ← 36x faster than Transformer

Model           Inference Memory (per 1K tokens)
Transformer     O(n) KV Cache
Mamba           O(1) SSM state (fixed size!)

Example: 130M parameter model, 1M token sequence
- Transformer: ~16GB KV Cache
- Mamba: ~1MB state (constant!)

Model        ListOps  Text  Retrieval  Image  Path-X  Avg
Transformer  36.4     65.0  57.5       42.4   0.0     40.3
LSTM         35.9     63.7  65.0       43.3   0.0     41.6
S4           59.6     86.8  90.9       88.7   86.1    82.4
Mamba        ~S4-level performance, faster processing

# Example of Mamba for biological sequence modeling
from mamba_ssm import MambaLMHeadModel
import torch

# DNA sequence modeling (A, T, G, C, N tokens)
DNA_VOCAB = {'A': 0, 'T': 1, 'G': 2, 'C': 3, 'N': 4}
VOCAB_SIZE = len(DNA_VOCAB)

# Process very long sequences efficiently
model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba-130m",
    device="cuda",
    dtype=torch.float16
)

# VMamba core: 2D selective scan
# Scan image in 4 directions to capture 2D structure

# Directions:
# 1. Left→Right, Top→Bottom (standard raster scan)
# 2. Right→Left, Bottom→Top (reverse)
# 3. Top→Bottom, Left→Right (column-first)
# 4. Bottom→Top, Right→Left (reverse)

class VMambaBlock(nn.Module):
    """VMamba: Visual Mamba Block"""
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        # 4-directional SSMs
        self.ssms = nn.ModuleList([
            MambaBlock(d_model, d_state) for _ in range(4)
        ])
        self.out_proj = nn.Linear(d_model * 4, d_model)

    def forward(self, x):
        """x: (B, H, W, D) - image patch embeddings"""
        b, h, w, d = x.shape
        x_flat = x.view(b, h*w, d)

        outputs = []
        # 4-directional scan
        for i, ssm in enumerate(self.ssms):
            if i == 0:   # forward
                seq = x_flat
            elif i == 1: # reverse
                seq = x_flat.flip(1)
            elif i == 2: # column-first
                seq = x.permute(0, 2, 1, 3).reshape(b, h*w, d)
            else:        # column-first reverse
                seq = x.permute(0, 2, 1, 3).reshape(b, h*w, d).flip(1)

            out = ssm(seq)
            if i % 2 == 1:
                out = out.flip(1)
            outputs.append(out)

        # Combine 4 directional results
        combined = torch.cat(outputs, dim=-1)
        return self.out_proj(combined).view(b, h, w, d)

# Required packages
pip install torch torchvision torchaudio
pip install causal-conv1d>=1.2.0
pip install mamba-ssm

# Or install from source (latest features)
git clone https://github.com/state-spaces/mamba
cd mamba
pip install -e ".[dev]"

import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer

# Load pre-trained model
model_name = "state-spaces/mamba-2.8b-slimpj"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")

model = MambaLMHeadModel.from_pretrained(
    model_name,
    device="cuda",
    dtype=torch.bfloat16
)
model.eval()

# Text generation
def generate_text(prompt, max_new_tokens=200, temperature=0.7):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.9,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated = output[0][input_ids.shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True)

# Example
prompt = "State Space Models are powerful because"
result = generate_text(prompt)
print(result)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer, get_linear_schedule_with_warmup

class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=512):
        self.encodings = tokenizer(
            texts,
            truncation=True,
            max_length=max_length,
            padding="max_length",
            return_tensors="pt"
        )

    def __len__(self):
        return len(self.encodings['input_ids'])

    def __getitem__(self, idx):
        return {
            'input_ids': self.encodings['input_ids'][idx],
            'attention_mask': self.encodings['attention_mask'][idx],
        }

def finetune_mamba(
    model_name="state-spaces/mamba-130m",
    texts=None,
    num_epochs=3,
    learning_rate=1e-4,
    batch_size=8,
):
    """Fine-tune Mamba model"""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
    tokenizer.pad_token = tokenizer.eos_token

    # Load model
    model = MambaLMHeadModel.from_pretrained(
        model_name,
        device=device,
        dtype=torch.bfloat16
    )

    # Dataset
    dataset = TextDataset(texts, tokenizer)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # Optimizer (AdamW recommended for Mamba)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=learning_rate,
        weight_decay=0.1
    )

    # Scheduler
    total_steps = len(dataloader) * num_epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=total_steps // 10,
        num_training_steps=total_steps
    )

    # Training loop
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch_idx, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(device)

            # Language modeling: next token prediction
            outputs = model(input_ids)
            logits = outputs.logits

            # Shift for next token prediction
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = input_ids[..., 1:].contiguous()

            loss = nn.CrossEntropyLoss()(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1)
            )

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            total_loss += loss.item()

            if batch_idx % 100 == 0:
                print(f"Epoch {epoch+1}, Step {batch_idx}: Loss = {loss.item():.4f}")

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1} complete: Avg Loss = {avg_loss:.4f}")

    return model

from mamba_ssm import Mamba
from mamba_ssm.models.config_mamba import MambaConfig

# Custom configuration
config = MambaConfig(
    d_model=1024,
    n_layer=48,
    vocab_size=50280,
    d_state=16,
    d_conv=4,
    expand=2,
    dt_rank="auto",
    dt_min=0.001,
    dt_max=0.1,
    dt_init="random",
    dt_scale=1.0,
    dt_init_floor=1e-4,
    rms_norm=True,
    residual_in_fp32=True,
    fused_add_norm=True,
    pad_vocab_size_multiple=8,
)

# Create model (~1.4B parameters)
model = MambaLMHeadModel(config, device="cuda", dtype=torch.bfloat16)
print(f"Parameter count: {sum(p.numel() for p in model.parameters()):,}")