아키텍처 개요¶

Bit-Axon 3.2B: Mamba 스타일 상태 공간 모델, 슬라이딩 윈도우 어텐션, 공유 전문가 mixture-of-experts를 결합한 24층 하이브리드 언어 모델 — Apple Silicon 전용 설계.

설계 철학: "GPU도, 클라우드도 없이"¶

Bit-Axon은 팬 없는 MacBook Air M4(16 GB 통합 메모리)에서 학습-추론-배포 전체 사이클을 실행할 수 있도록 처음부터 설계되었습니다. 세 가지 아키텍처 기둥이 이를 가능하게 합니다:

기둥	기법	효과
선형	Axon-SSM (Mamba 스타일 SSM)	토큰당 \(O(1)\) 메모리, KV cache 없음
희소	Shared-Expert MoE	3.2B parameters 중 토큰당 약 ~1.4B만 활성화
양자화	4-bit 가중치 + TurboQuant KV cache	약 ~1.76 GB 가중치 사용량

24층 샌드위치 구조¶

24개 층은 세 가지 기능 영역으로 나뉩니다. SSM 층은 컨텍스트 길이와 무관하게 일정한 메모리를 유지하며, 중간 SWA 영역만 KV cache를 생성합니다.

                    Bit-Axon 3.2B — 24층 샌드위치
 ┌──────────────────────────────────────────────────────────────┐
 │  입력: embed_tokens (vocab 32K) → input_proj (2048→2560)    │
 ├──────────────────────────────────────────────────────────────┤
 │                                                              │
 │  영역 1: 기반                 층 1–8  (순수 Axon-SSM)        │
 │  ┌────────────────────────────────────────────────────────┐  │
 │  │  AxonSSMBlock: RMSNorm → AxonSSM (MLP 없음)            │  │
 │  │  • 선형 순회를 통한 컨텍스트 흡수                         │  │
 │  │  • KV cache 없음 — 토큰당 O(1) 메모리                    │  │
 │  └────────────────────────────────────────────────────────┘  │
 │                         × 8층                                │
 │                                                              │
 │  영역 2: 심층 추론            층 9–16 (SWA + MoE)            │
 │  ┌────────────────────────────────────────────────────────┐  │
 │  │  AxonSWAMoEBlock:                                       │  │
 │  │    RMSNorm → SWA (window=4096) + 잔차                   │  │
 │  │    RMSNorm → SharedExpertMoE (8 전문가, top-2) + 잔차    │  │
 │  │  • 국소 추론을 위한 슬라이딩 윈도우 어텐션                 │  │
 │  │  • KV cache 필요 (층 9–16만)                             │  │
 │  └────────────────────────────────────────────────────────┘  │
 │                         × 8층                                │
 │                                                              │
 │  영역 3: 출력 합성            층 17–24 (SSM + MoE)           │
 │  ┌────────────────────────────────────────────────────────┐  │
 │  │  AxonSSMMoEBlock:                                       │  │
 │  │    RMSNorm → AxonSSM + 잔차                             │  │
 │  │    RMSNorm → SharedExpertMoE (8 전문가, top-2) + 잔차    │  │
 │  │  • 선형 순회 + 희소 전문가                                │  │
 │  │  • KV cache 없음 — 열 효율적인 출력 생성                  │  │
 │  └────────────────────────────────────────────────────────┘  │
 │                         × 8층                                │
 │                                                              │
 ├──────────────────────────────────────────────────────────────┤
 │  출력: output_proj (2560→2048) → lm_head (2048→32000)       │
 └──────────────────────────────────────────────────────────────┘

왜 이 구조인가?¶

질문	답변
왜 층 1–8에 SSM을 사용하는가?	원시 컨텍스트에는 추론이 아닌 흡수가 필요합니다. SSM은 \(O(1)\) 메모리를 제공하고 모델의 1/3에 대해 KV cache를 제거합니다.
왜 층 9–16에만 SWA를 사용하는가?	심층 추론은 국소 어텐션의 이점을 얻습니다. SWA를 8개 층으로 제한하면 KV cache 메모리를 관리 가능한 범위로 제한할 수 있습니다.
왜 층 17–24에서 어텐션을 제거하는가?	출력 생성은 자기회귀적입니다. SSM + MoE는 최소 열 부하로 토큰을 생성합니다 — 팬 없는 기기에서 매우 중요합니다.

모델 구성¶

Parameter	값	비고
`vocab_size`	32,000	Qwen의 151K에서 축소
`hidden_dim`	2,560	모델 폭 (\(d_\text{model}\))
`num_layers`	24	샌드위치: 8 + 8 + 8
`num_heads`	32	SWA 어텐션 헤드
`head_dim`	80	\(2560 / 32\)
`d_source_model`	2,048	Qwen2.5-3B 브릿지 차원
`ssm_d_state`	16	SSM 상태 벡터 크기
`ssm_d_conv`	4	인과적 conv1d 커널
`ssm_expand`	3	SSM 중간 차원 = \(2560 \times 3 = 7680\)
`swa_window_size`	4,096	슬라이딩 윈도우 범위
`moe_num_experts`	8	총 라우팅 전문가
`moe_top_k`	2	토큰당 활성 전문가
`moe_intermediate_dim`	4,096	전문가 FFN 차원
`moe_shared_expert`	`true`	공유 전문가 항상 활성
`max_seq_len`	65,536	최대 컨텍스트 (64K)
`weight_tying`	`true`	`embed_tokens.weight = lm_head.weight`
`rms_norm_eps`	1e-6	RMSNorm 엡실론

핵심 설계 결정¶

차원 브릿지 (\(d_\text{source} = 2048\))¶

가중치는 hidden_size=2048을 사용하는 Qwen2.5-3B에서 가져왔습니다. Bit-Axon은 더 넓은 표현을 위해 hidden_dim=2560으로 동작합니다. 두 투영 층이 차원 불일치를 처리합니다:

embed_tokens(vocab=32000, dim=2048)
       ↓
  input_proj(2048 → 2560)     ← Bit-Axon의 더 넓은 공간으로 투영
       ↓
  [2560 차원의 24층 샌드위치]
       ↓
  output_proj(2560 → 2048)    ← 원본 차원으로 역투영
       ↓
  lm_head(2048 → 32000)

Qwen의 36개 층 중 첫 24개 층이 1:1로 매핑되며, 층 24–35는 제거됩니다. SSM parameters는 무작위 초기화됩니다. MoE 전문가 가중치: 공유 전문가는 Qwen의 MLP에서 가져오고, 전문가 0은 복사본, 전문가 1–7은 섭동된 복사본입니다.

가중치 공유¶

embed_tokens.weight와 lm_head.weight는 동일한 parameter 텐서를 공유합니다. 이를 통해 약 ~64 MB의 중복 저장 공간을 제거합니다 (\(2048 \times 32000 \times 2\) 바이트, FP16 기준).

MLX 통합¶

전체 모델은 Apple의 MLX 프레임워크로 구축되었으며, PyTorch가 아닙니다:

JIT 컴파일: 리프 함수(_ssm_fma, _compute_dt, swiglu)는 @mx.compile로 데코레이션되어 퓨즈된 Metal 커널을 생성합니다
통합 메모리 제로 카피: 양자화된 가중치는 한 번 로드되어 복사 없이 CPU와 GPU 모두에서 접근됩니다
Metal 최적화 quantization: nn.QuantizedLinear은 4-bit matmul을 위해 퓨즈된 Metal 커널을 사용합니다

MLX 컴파일 제약사항

모델 수준의 mx.compile이 사용됩니다 (모듈 참조 문제로 인해 층 수준은 작동하지 않음)
shapeless=True는 MLX ≤ 0.31에서 matmul에 대해 작동하지 않음 — shape 종속 컴파일 사용
MoE 라우팅의 NumPy interop은 mx.compile 추적을 중단시킴 — 순수 MLX 디스패치 필요

하위 페이지¶

페이지	내용
Axon-SSM	Mamba 스타일 선택적 상태 공간 모델: 알고리즘, 수학, 구현
SWA + MoE	슬라이딩 윈도우 어텐션과 공유 전문가 mixture-of-experts
메모리 예산	메모리 분석, quantization 전략, 열 관리