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

왜 중요한가

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

핵심 기여

완전 온라인 Mamba 압축 아키텍처

사전 학습된 가중치나 외부 데이터 없이, 압축 대상 파일에서 직접 토큰 단위로 학습하는 2레이어 Mamba SSM 기반의 무손실 압축 시스템을 구축했다. 순수 C 언어와 AVX2 명령어를 사용하여 GPU 없이 CPU만으로 초당 약 2,000 토큰의 학습 및 압축 속도를 달성했다.

소프트맥스 불변 로짓 바이어스 메커니즘

n-gram 모델의 증거를 SSM의 로짓에 결합할 때, 카운트가 0이 아닌 토큰만 업데이트하는 희소 업데이트 공식을 유도했다. 이를 통해 고차 n-gram(최대 32-gram)을 메모리와 연산 효율적으로 처리하며 SSM의 전역 예측 능력을 보완했다.

엔트로피 적응형 컨텍스트 믹싱

SSM의 예측 엔트로피(불확실성)에 따라 n-gram 모델의 반영 비중을 실시간으로 조절한다. SSM이 확신할 때는 n-gram의 개입을 줄이고, SSM이 불확실할 때는 n-gram의 로컬 기억 장치에 더 의존하도록 설계하여 압축 효율을 극대화했다.

핵심 아이디어 이해하기

기존의 신경망 기반 압축은 미리 대규모 데이터로 학습된 거대한 모델을 사용해야 하므로 모델 파일 자체가 너무 커서 배포에 불리하거나, 압축 시점에 학습을 하더라도 연산량이 너무 많다는 한계가 있었다. 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와 같은 기존 사전 기반 압축기가 블록 단위 복사 기능을 통해 더 높은 효율을 내기 때문에, 대용량 파일에서는 경쟁력이 떨어집니다.

실무 활용

StateSMix는 GPU가 없는 환경이나 외부 모델 가중치를 공유하기 어려운 보안 환경에서 텍스트 데이터를 고효율로 압축해야 할 때 유용합니다. 특히 10MB 미만의 개별 문서 압축에서 기존 xz보다 뛰어난 성능을 보입니다.

임베디드 시스템 및 IoT 기기의 로그 데이터 압축
네트워크 대역폭이 제한된 환경에서의 실시간 텍스트 전송
사전 학습 모델 배포가 불가능한 폐쇄망 환경의 데이터 보관
개별 자연어 문서(위키피디아, 소스 코드 등)의 고효율 아카이빙

코드 공개 여부: 공개

코드 저장소 보기

키워드

SSM(상태 공간 모델)Mamba(맘바)무손실 압축(Lossless Compression)온라인 학습(Online Learning)n-gram(엔그램)

용어 해설

상태 공간 모델(State Space Model): — 시퀀스 데이터를 처리하기 위해 입력, 상태, 출력 사이의 선형 미분 방정식을 사용하는 아키텍처입니다. RNN의 재귀적 특성과 CNN의 병렬 학습 가능성을 결합하여 긴 시퀀스에서도 선형적인 연산 복잡도를 유지하며 효율적인 정보 유지가 가능합니다.
산술 부호화(Arithmetic Coding): — 데이터의 각 심볼에 확률 구간을 할당하고 전체 메시지를 하나의 실수 범위로 표현하는 무손실 압축 기법입니다. 심볼의 발생 확률이 높을수록 더 넓은 구간을 할당하여 적은 비트로 표현함으로써 엔트로피 한계에 가까운 압축률을 달성합니다.
엔그램(N-gram): — 텍스트 시퀀스에서 연속된 n개의 아이템(토큰 또는 문자)의 묶음을 의미합니다. 이전 n-1개의 토큰을 바탕으로 다음 토큰이 나타날 확률을 예측하는 통계적 모델링에 주로 사용되며, 로컬 패턴을 기억하는 데 매우 효과적입니다.
온라인 학습(Online Learning): — 데이터가 순차적으로 들어올 때마다 모델을 실시간으로 업데이트하는 학습 방식입니다. 별도의 사전 학습 데이터 없이 현재 처리 중인 데이터의 특성에 즉각적으로 적응할 수 있어, 특정 파일의 고유한 패턴을 파악하는 압축 작업에 유리합니다.
로짓 바이어스(Logit Bias): — 모델이 출력하는 확률 분포의 이전 단계인 로짓 값에 특정 수치를 더하거나 빼서 특정 토큰의 선택 확률을 인위적으로 조정하는 기법입니다. 본 논문에서는 n-gram 모델의 예측 결과를 SSM의 결과와 결합하는 용도로 사용됩니다.

코드 예제

python

niter = {
    8 chunks 1-10,
    4 chunks 11-30,
    2 chunks 31 +
}

학습 초기 단계에서 모델을 빠르게 수렴시키기 위해 적용된 Adam 최적화 반복 횟수 스케줄링 예시

text

delta_j = lambda * log(1 + c_j / alpha), c_j > 0

n-gram 카운트 정보를 기반으로 로짓 바이어스 값을 계산하는 수식

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드