기술 요약: 밀도 필드 상태 공간 모델 (Density Field State Space Models, DF-SSM)

1. 문제 정의

13억 개의 파라미터를 가진 FP16 정밀도의 대규모 언어 모델(LLM)은 약 2.7 GB의 메모리를 요구하며, 이는 대부분의 모바일, 임베디드 및 엣지 디바이스의 용량을 초과합니다. 양자화(Quantization)를 통해 모델 크기를 줄일 수 있지만, 극단적인 양자화(1~4비트)는 일반적으로 심각한 품질 저하를 초ку하거나 방대한 데이터셋(예: 100B+ 토큰)을 통한 처음부터의 학습을 필요로 합니다. 또한, Mamba-2와 같은 상태 공간 모델(SSM)은 고정된 크기의 은닉 상태를 통해 선형적으로 증가하는 KV 캐시 문제를 해결하여 온디바이스 배포에 구조적 이점을 제공하지만, 트랜스포머 양자화에 비해 가중치 압축 연구는 아직 미진한 상태입니다.

2. 방법론: DF-SSM 파이프라인

본 논문은 SSM을 1비트 스캐폴드(scaffold)와 int8 저계수(low-rank) 보정으로 압축하도록 설계된 프레임워크인 **밀도 필드 상태 공간 모델 (DF-SSM)**을 소개합니다. 이 접근 방식은 세 가지 주요 단계로 구성됩니다.

A. 밀도 필드 가중치 (Density Field Weight, DFW) 학습

저자들은 사전 학습된 FP16 교사 모델(Mamba-2 1.3B)로부터 양자화 인식 증류(quantization-aware distillation)를 채택했습니다.

• 17단계 양자화: 가중치는 -1과 +1 사이의 17개 이산 레벨로 매핑되는 연속적인 잠재 값으로 파라미터화됩니다.
• 양자화 어닐링 (Quantization Annealing): 어닐링 파라미터 $\alpha$는 연속 가중치에서 양자화된 가중치로의 전환을 보간하며, 학습 과정에서 0에서 1까지 선형적으로 증가합니다. 이는 모델이 이산적 제약에 적응하기 전에 연속 가중치로 언어 모델링 목적 함수를 먼저 학습할 수 있게 함으로써 "파괴적인 품질 절벽(catastrophic quality cliff)"을 방지합니다.
• 그래디언트 처리: 그래디언트 계산에는 Straight-through estimation이 사용됩니다. 저자들은 그래디언트 감쇠를 피하기 위해 시그모이드 매핑 대신 선형 클램핑(linear clamping)을 선택했습니다.

B. 동결된 스캐폴드와 LoRA 보정

학練 후, 양자화된 스캐폴드 가중치는 결정론적으로 반올림된 값으로 동결됩니다.

• 저계수 적응 (Low-Rank Adaptation): 잔차 오차를 수정하기 위해 작은 학습 가능한 int8 LoRA 모듈($W_A \in \mathbb{R}^{d_{out} \times r}, W_B \in \mathbb{R}^{r \times d_{in}}$, 랭크 $r=16$)이 추가됩니다.
• 아키텍처: LoRA는 48개 레이어 전체(in_proj 및 out_proj 모두)에 적용됩니다.
• 메커니즘: LoRA 보정은 단순히 개별 가중치의 반올림 오차를 수정하는 것이 아니라(반올림 잔차와의 상관관계는 ~0임), 입력 가중치가 적용된 출력 오차를 수정합니다.
• 양자화: LoRA 행렬과 임베딩/LM 헤드 행렬은 모두 per-channel/row 스케일링이 적용된 int8로 양자화되어, 사실상 손실 없는 퍼플렉시티(PPL 49.1 vs. 49.2 for FP16)를 달전합니다.

C. 최적화된 추론 파이프라인

저자들은 GPU와 CPU 모두를 위한 커스텀 추론 엔진을 개발했습니다.

• GPU (CUDA):
  • 스캐폴드 행렬 곱 (Scaffold Matmul): 1비트 패킹된 가중치는 온더플라이(on-the-fly)로 int8로 확장되어 cuBLAS INT8 텐서 코어(torch._int_mm 사용)를 통해 처리됩니다. 이를 통해 8배 감소된 메모리 전송과 더 높은 INT8 처리량을 활용합니다.
  • 상태 연산 (Stateful Ops): 커스텀 CUDA 커널이 상태 시프트 레지스터(convolution)와 SSM 재귀(warp-level reduction)를 처리합니다.
  • CUDA 그래프: 48개 레이어 전체의 순전파 과정을 CUDA 그래프로 캡처하여, 배치 크기 1 성능에서 치명적인 커널 런칭 오버헤드를 제거합니다.
• CPU (AVX-512):
  • 비트 패킹된 커널이 가중치를 int8로 확장하여 VNNI 명령어(_mm512_dpbusd_epi32)를 수행하며, 표준 int8 가중치 대비 8배 적은 데이터를 로드하면서도 높은 처리량을 달성합니다.

D. 밀도 필드 상태 표현 (Density Field State Representation)

SSM 은닉 상태를 압축하기 위해, 저자들은 각 상태 요소당 $K \times K$ 이진 필드($K^2 \approx 4 \times d_{state}$)를 유지합니다. 이 과정은 정규화, 시그마-델타 주입(오차 피드백), 이진화, 그리고 연속 값을 복구하기 위한 popcount 판독으로 구성되며, 청크 경계에서 적용됩니다.

3. 주요 결과

모델 효율성 및 크기

• 압축: 결과 모델은 278 MB로, 2.7 GB의 FP16 교사 모델 대비 9.7배 감소했습니다.
• 구성: 스캐폴드(1-bit) 56%, 임베딩/LM 헤드(int8) 37%, LoRA(int8) 4%, 기타(FP16/FP32) 3%.
• 학습 비용: 증류에는 단일 A100 GPU에서 32M 토큰과 6시간이 소요되었습니다. 이는 BitMamba-2를 처음부터 학습시키는 데 필요한 150B 토큰과 비교됩니다.

추론 성능

• GPU 속도 향상: 최적화된 mamba-ssm 레퍼런스 구현 대비 GPU(배치=1)에서 21.4배 빠른 추론을 달 achieves 합니다. 배치=512에서 처리량은 A100의 HBM 대역폭 제한(~5,081 tok/s)에 근접합니다.
• CPU 속도 향상: Intel Xeon, AVX-512 환경에서 567 MB RAM 사용량과 함께 2.9배의 속도 향상을 달성했습니다.

다운스트림 태스크 성능

DF-SSM은 BitMamba-2(150B 토큰으로 처음부터 학습된 1.58-bit 모델)와 2~4 퍼센트 포인트 내의 성능 차이를 유지하며, FP16 교사 모델 성능의 78~95%를 유지합니다.

• BoolQ: 60.8% (교사: 64.2%)
• PIQA: 67.1% (교사: 73.2%)
• WinoGrande: 54.7% (교사: 60.9%)
• ARC-easy: 50.2% (교사: 64.1%)
• HellaSwag: 41.4% (교사: 59.9%) – 다단계 추론을 위한 증류 토큰 부족으로 인한 예외 사례로 언급됨.

4. 해석 가능성 발견 (Interpretability Findings)

본 논문은 SSM의 사실적 지식 조직화에 대한 최초의 체계적인 분석을 제시하며, 세 가지 뚜렷한 처리 단계를 밝혀냈습니다.

1. 의도 분류 (Layers 0–3): 어휘 정렬이 없는 추상 공간에서 작동합니다. 분류는 의미론적(semantic)이기보다 **구문론적(syntactic)**입니다(예: "The capital of..."와 같은 템플릿 구조에 의해 구동됨). 동일한 템플릿을 가진 카테고리는 즉시 클러스터링되지만, 혼합된 템플릿의 카테고리는 그렇지 않습니다.
2. 지식 검색 (Layers 25–35): 사실적 연관성은 5개 레이어 윈도우(L32–L36에서 피크)에 국한됩니다. 인과적 개입 실험에 따르면, 이 윈도우의 은닉 상태를 교체하면 사실적 예측이 뒤집힙니다. 지식은 2048차원의 상태 중 저차원 부분 공간(10–16 차원)을 점유합니다.
3. 출력 포매팅 (Layers 36–47): 카테고리 구조가 해체되면서 모델은 토큰을 방출하기 위해 공유된 출력 표현으로 수렴합니다.

핵심 관찰: 모델은 약한 사실적 회상 능력(정답이 top-1 순위에 도달하는 경우가 드묾)에도 불구하고 잘 조직된 내부 지식 구조(깨끗한 클러스터링, 인과적 국소화)를 보여줍니다. 이는 표현 구조가 사실적 강도보다 앞설 수 있음을 시사합니다.

5. 의의 및 주장

본 논문은 다음과 같은 기여를 주장합니다.

• 효율성: 최소한의 증류 데이터(32M 토큰)를 사용하여 6시간 만에 1.3B SSM을 배포 가능한 바이너리 모델로 압축하는 완전한 파이프라인을 제공합니다.
• 성능: 메모리 대역폭 최적화와 INT8 텐서 코어 활용을 통해 FP16 교사 모델의 성능을 거의 유지하면서 GPU에서 21.4배, CPU에서 2.9배의 추론 속도 향상을 달성했습니다.
• 아키텍처: 균일한 극단적 양자화의 품질 저하를 피하는 혼합 정밀도 아키텍처(1-bit 스캐폴드 + int8 LoRA)를 제안합니다.
• 해석 가능성: SSM의 3단계 처리 모델을 발견하고, 초기 레이어의 분류가 구문론적이며 템플릿 기반으로 작동하며, 어휘와 정렬되지 않은 공간에서 작동함을 밝혀냈습니다.
• 시사점: 결과는 증류가 지식 표현의 조직 구조를 효율적으로 전달할 수 있음을 시사합니다. 이는 사실적 내용이 부분적으로만 유지되더라도 추가 학습 토큰이 기존 구조를 재구성하기보다는 채울 수 있음을 의미합니다.

인정된 한계점:

• PPL 저하(49.2 vs. 14.3)는 더 긴 학습이나 더 높은 LoRA 랭크를 통한 개선 여지가 있음을 나타냅니다.
• 해석 가능성 결과는 단일 증류 모델(1.3B Mamba-2)에 특화되어 있으며, 처음부터 학습된 다른 규모나 아키텍처로의 일반화는 검증되지 않았습니다.
• "1-bit"라는 명칭은 스캐폴드(가중치의 56%)에 적용되는 것이며, 임베딩과 LoRA는 여 still int8로 유지됩니다. 이는 다른 바이너리 SSM 방식들과 일치합니다.