StateSMix: Mamba 상태 공간 모델과 희소 N-gram 컨텍스트 믹싱을 통한 온라인 무손실 압축

수억 개의 파라미터를 가진 거대 모델이나 GPU 없이도, 압축 과정 중에 실시간으로 학습하는 가벼운 Mamba 모델이 기존의 강력한 압축 도구인 xz보다 더 높은 압축률을 기록했습니다. 이는 특정 데이터에 즉각적으로 적응하는 온라인 학습 모델이 범용 압축 분야에서 실질적인 경쟁력을 가질 수 있음을 보여줍니다.

기존의 신경망 기반 압축은 미리 대규모 데이터로 학습된 거대한 모델을 사용해야 하므로 모델 파일 자체가 너무 커서 배포에 불리하거나, 압축 시점에 학습을 하더라도 연산량이 너무 많다는 한계가 있었다. StateSMix는 이 문제를 해결하기 위해 연산 효율이 뛰어난 Mamba 아키텍처를 선택했다. Mamba는 시퀀스 길이에 따라 연산량이 선형적으로 증가하므로, 파일의 앞부분을 읽으며 배운 지식을 바탕으로 뒷부분을 예측하는 '온라인 학습'에 최적이다.

핵심 아이디어는 Mamba가 문장의 문법이나 의미 같은 '전역적인 패턴'을 학습하는 동안, n-gram 모델이 자주 반복되는 '구체적인 문구'를 기억하도록 역할을 분담시킨 것이다. 예를 들어 'Hugging Face'라는 단어가 파일에 자주 등장한다면, Mamba는 문맥상 고유 명사가 나올 타이밍임을 예측하고, n-gram은 구체적으로 'Face'가 올 차례임을 확신하게 된다.

이 과정에서 '엔트로피 적응형 믹싱'이 사용된다. 모델이 다음 단어를 예측할 때 확률 분포의 균일함(엔트로피)을 측정하여, 모델이 스스로 '잘 모르겠다'고 판단하면 n-gram의 기억력을 더 강하게 반영하고, 모델이 확신에 차 있으면 n-gram의 간섭을 줄인다. 결과적으로 사전 학습 없이도 파일 하나를 읽는 과정만으로 기존 압축 알고리즘인 xz를 능가하는 효율을 보여준다.

전체 시스템은 BPE 토큰화, 어휘 재매핑, 온라인 예측-부호화-업데이트 루프, 직렬화의 4단계로 구성된다. 입력 데이터는 GPT-NeoX 어휘집을 사용하는 BPE 토큰화 과정을 거쳐 정수 ID 시퀀스로 변환된다. 이때 해당 파일에 실제로 등장하는 토큰들로만 어휘집을 재구성하는 'Compact Vocabulary Remapping'을 수행하여 연산 대상 범위를 49,152개에서 약 18,000~44,000개 수준으로 줄인다.

핵심 예측 모델은 2개의 Mamba 레이어로 구성된 SSM이다. 각 레이어는 입력값 x를 받아 Layer Normalization을 거친 후, 두 갈래(SSM 브랜치와 게이트 브랜치)로 나누어 처리한다. [입력 벡터 x → 선형 투영 → 1D Depthwise Convolution → SiLU 활성화 함수 → 선택적 재귀 연산(Selective Recurrence) → 게이트 결합] 순으로 연산이 진행된다. 재귀 연산 단계에서는 입력에 따라 변화하는 ∆(step size)를 사용하여 현재 상태 h를 업데이트하며, 이는 모델이 중요한 정보는 오래 기억하고 불필요한 정보는 빠르게 잊도록 제어한다.

학습은 32개 토큰으로 구성된 청크(chunk) 단위로 수행된다. 각 청크를 처리한 후 Adam 옵티마이저를 사용하여 교차 엔트로피 손실(Cross-Entropy Loss)을 최소화하는 방향으로 가중치를 갱신한다. [예측 확률 p와 실제 토큰 t → -log(p) 계산 → 손실값 도출 → 역전파를 통한 그라디언트 계산 → 가중치 업데이트] 과정을 거치며 모델은 실시간으로 파일의 통계적 특성에 최적화된다.

enwik8(100MB) 벤치마크 실험 결과, StateSMix는 10MB 이하의 파일 크기에서 범용 압축 도구인 xz -9e를 일관되게 능가했다. 1MB 파일에서는 xz 대비 8.7% 개선된 2.123 bpb(bits per byte)를 기록했으며, 10MB에서는 0.7% 우위를 점했다. 다만 파일 크기가 30MB를 넘어서면 xz의 사전 기반 복사 방식이 더 효율적으로 작용하여 역전되는 양상을 보였다.

Ablation Study에 따르면 SSM 모델이 압축률 향상의 가장 큰 기여자로 나타났다. 빈도수 기반 기준 모델 대비 SSM만 사용했을 때 46.6%의 크기 감소 효과가 있었으며, n-gram 모델만 단독으로 사용했을 때는 16.1% 감소에 그쳤다. 이는 SSM이 단순 반복을 넘어 언어의 구조적 패턴을 효과적으로 학습하고 있음을 증명한다.

성능 측면에서는 4코어 CPU(AVX2 SIMD 활용) 환경에서 초당 약 2,000 토큰(약 700 KB/s)의 속도를 기록했다. 이는 GPU나 대규모 사전 학습 가중치 없이 순수하게 실시간 학습만으로 달성한 수치로, 임베디드 시스템이나 보안이 중요한 환경에서 외부 리소스 없이 고성능 압축이 가능함을 시사한다.

StateSMix의 아키텍처는 DM=32, NL=2의 매우 경량화된 Mamba 설정을 사용한다. 총 파라미터 수는 파일당 약 120K개에 불과하며, 이는 온라인 학습의 연산 부담을 최소화하기 위한 선택이다. 모델은 매 32개 토큰마다 Adam 옵티마이저를 통해 업데이트되며, 초기 10개 청크에는 8회의 반복 학습을 수행하는 웜업 스케줄을 적용하여 초기 수렴 속도를 높였다.

n-gram 모델은 2-gram부터 32-gram까지 총 9개의 차수를 동시에 사용한다. 각 차수는 2^24개의 슬롯을 가진 해시 테이블로 관리되며, 충돌 해결을 위해 선형 탐사(Linear Probing) 방식을 채택했다. 특히 16-gram과 32-gram 같은 장거리 컨텍스트 테이블은 SSM의 재귀적 메모리 한계를 넘어선 반복 문구(예: 위키피디아의 인용 양식)를 포착하는 역할을 한다.

수학적으로 n-gram의 결합은 베이지안 업데이트로 해석될 수 있다. SSM의 출력을 사전 확률(Prior)로 보고, n-gram의 카운트 정보를 우도(Likelihood)로 사용하여 사후 확률(Posterior)을 계산하는 방식이다. 이때 사용되는 로짓 바이어스 공식 δ_j = λ log(1 + c_j/α)는 소프트맥스의 평행 이동 불변성을 활용하여, 카운트가 있는 토큰만 효율적으로 업데이트할 수 있게 한다.

구현 측면에서는 C 언어와 AVX2 SIMD 인트린직을 사용하여 병렬화를 극대화했다. 특히 연산 병목인 헤드 프로젝션(Head Projection) 단계에서 _mm256_fmadd_ps 명령어를 사용하여 32개의 부동 소수점 연산을 동시에 처리함으로써 CPU 환경에서도 실시간 학습이 가능한 수준의 성능을 확보했다.

가장 큰 한계는 압축 속도(약 700 KB/s)와 메모리 사용량(100MB 입력 시 약 6.1 GB)입니다. 또한 파일 크기가 30MB를 초과하면 xz와 같은 기존 사전 기반 압축기가 블록 단위 복사 기능을 통해 더 높은 효율을 내기 때문에, 대용량 파일에서는 경쟁력이 떨어집니다.