Axon-SSM: 선택적 상태 공간 모델¶

Axon-SSM은 Bit-Axon의 Mamba 스타일 선택적 상태 공간 모델입니다. 표준 Transformer 자기 어텐션을 선형 순회로 대체하여, 토큰당 \(O(1)\) 메모리를 달성하고 KV cache를 완전히 제거합니다.

개요¶

Axon-SSM은 24개 층 중 16개에 등장합니다:

• 층 1–8: 순수 SSM (AxonSSMBlock) — SSM의 내부 확장이 FFN/MLP 역할을 완전히 대체
• 층 17–24: SSM + MoE (AxonSSMMoEBlock) — SSM이 순회를 담당하고 MoE가 희소 연산을 추가

알고리즘¶

Axon-SSM의 순전파는 다음 단계를 따릅니다:

입력 x: (batch, seq_len, hidden_dim=2560)
            │
            ▼
    ┌─── in_proj ───┐
    │  (D → 2E)     │
    └───────┬───────┘
            │ split
      ┌─────┴─────┐
      ▼           ▼
  x_branch    z_branch     (각 차원 E=7680)
      │           │
      ▼           │
 ┌──────────┐     │
 │ conv1d   │     │   인과적, kernel=4, groups=E
 │ (depthwise)│   │
 └────┬─────┘     │
      ▼           │
  ┌───────┐       │
  │ SiLU  │       │
  └───┬───┘       │
      │           │
      ▼           │
 ┌─── x_proj ───┐ │
 │ (E → 2·d_state+1) │
 └──────┬───────┘ │
   ┌────┼────┐    │
   ▼    ▼    ▼    │
   B    C   dt_raw   (B, C: 각각 d_state; dt_raw: 1)
   │    │    │      │
   │    │    ▼      │
   │    │ ┌────────┐│
   │    │ │dt_proj ││  (1 → E)
   │    │ │+softplus││
   │    │ │+clip   ││
   │    │ └───┬────┘│
   │    │     ▼     │
   │    │    dt     │  (채널별 스텝 크기, 차원 E)
   │    │     │     │
   ▼    ▼     ▼     │
 ┌──────────────┐   │
 │  SSM 스캔    │   │  seq_len에 대한 순차 순회
 │  (아래 참조) │   │
 └──────┬───────┘   │
        ▼           │
        y           │
        │           │
        ▼           ▼
    ┌─────────────────┐
    │  y * SiLU(z)    │   게이팅: SSM 출력에 활성화된 z 브랜치를 곱함
    └────────┬────────┘
             ▼
    ┌─── out_proj ───┐
    │  (E → D=2560)  │
    └────────┬───────┘
             ▼
      출력: (batch, seq_len, 2560)

핵심 투영¶

SSM 순회¶

핵심 스캔은 각 타임스텝 \(t\)에서 이산화된 선형 순회를 계산합니다:

여기서:

이산화¶

스텝 크기 \(\Delta t\)는 softplus 투영과 클램핑을 통해 계산됩니다:

이를 통해 \(\Delta t\)가 수치적으로 안정적인 범위를 유지하면서 입력 의존적(선택적)으로 유지됩니다.

상태 초기화¶

\(A\) 행렬은 반복된 대각행렬로 초기화됩니다:

런타임에: \(A = -\exp(A_{\log})\), \(-1\)에서 \(-d_\text{state}\)까지의 대각선을 생성합니다. 음의 지수는 시간에 따른 은닉 상태의 안정적인 감쇠를 보장합니다.

메모리 특성¶

토큰당 일정한 메모리¶

KV cache가 시퀀스 길이에 따라 \(O(n)\)으로 증가하는 표준 어텐션과 달리, SSM은 고정 크기 상태를 유지합니다:

16개 어텐션 층에 대해 64K 컨텍스트로 전체 KV cache를 사용할 경우 수 GB가 필요한 것과 비교해 보세요.

KV cache 없음¶

SSM 층은 캐시로 [conv_cache, ssm_state]를 반환합니다 — 작고 고정 크기의 텐서입니다. 모델의 _create_caches() 메서드는 모든 SSM 층에 대해 None을 반환하고, 8개 SWA 층(9–16)에 대해서만 KVCache 객체를 반환합니다.

JIT 컴파일된 커널¶

두 리프 함수는 퓨즈된 Metal 커널 생성을 위해 @mx.compile로 데코레이션됩니다 (Jamba 패턴을 따름):

_ssm_fma¶

@mx.compile
def _ssm_fma(a: mx.array, b: mx.array, c: mx.array) -> mx.array:
    return a * b + c    # dA * h + dB * x_t  (퓨즈된 곱셈-덧셈)

이는 상태 갱신 \(h_t = dA \cdot h_{t-1} + dB \cdot x_t\)를 단일 커널로 퓨즈하여 중간 텐서 할당을 피합니다.

_compute_dt¶

@mx.compile
def _compute_dt(dt: mx.array, dt_bias: mx.array, lo: float, hi: float) -> mx.array:
    return mx.clip(nn.softplus(dt + dt_bias), lo, hi)

바이어스 덧셈, softplus 활성화, 클램핑을 하나의 커널로 퓨즈합니다.

자기회귀 디코딩¶

점진적(토큰별) 생성 중 캐시 메커니즘은 다음과 같이 작동합니다:

1. 첫 호출 (prefill, cache=None): 전체 프롬프트를 처리하고, ssm_state를 0으로 초기화하며, 마지막 \(K-1\) 위치에서 conv_cache를 구축합니다.
2. 이후 호출 (decode, cache=[conv_cache, ssm_state]): conv_cache와 새로운 단일 토큰을 연결하고, 이전 ssm_state를 사용하여 스캔의 한 스텝을 실행한 뒤 갱신된 캐시를 반환합니다.

스캔 루프는 seq_len 위치만큼 반복됩니다 — prefill 중에는 전체 프롬프트 길이이고, decode 중에는 정확히 1입니다.

Parameters¶

SSM 층당 parameter 수 (\(D=2560\), \(E=7680\), \(d_\text{state}=16\) 기준):

SSM을 포함하는 16개 층(8 순수 + 8 MoE 포함)으로, SSM은 전체 3.2B 중 약 960M parameters를 차지합니다.

← 아키텍처로 돌아가기