16MB 고정 어휘 코드로 작은 언어모델의 임베딩 테이블을 대체한다

무엇인가

이 논문이 다루는 문제는 작은 언어모델에서 어휘 임베딩 테이블이 차지하는 몫이다. Gemma 3 270M은 파라미터 268M 중 168M, 즉 63%가 임베딩 테이블이고 모델 연산의 대부분을 하는 FFN은 71M에 불과하다. 양자화는 이 상황을 개선하지 않는다. BitNet은 임베딩과 출력 헤드를 16비트로 유지하고, GPTQ와 AWQ는 embed_tokens와 lm_head를 건드리지 않으며, llama.cpp는 출력 헤드를 6~8비트로 저장한다. 그래서 body를 아무리 압축해도 파일에 남는 것은 테이블이다. ternary body에 16비트 어휘를 붙인 Gemma 3 270M은 body 20MB에 어휘 335MB가 된다. 저자들은 Microsoft의 BitNet b1.58 2B4T 보고서가 메모리를 "0.4 GB (non-embedding)"으로 적어 테이블을 헤드라인 숫자에서 빼놓은 점도 지적한다.

어떻게 동작하나

제안 방법은 학습된 테이블을 규칙으로 만든 고정 테이블로 교체하는 것이다. 사전학습된 임베딩 행렬 E(Gemma 3 270M의 262,144 x 640)를 중심화하고 상위 512개 주성분 방향 W를 구한 뒤, 각 투영의 부호만 남긴다. code(t) = sign(W^T (E_t - mu))로 -1과 +1로 이루어진 512비트가 토큰마다 부여되고, 전체 어휘로 16.8MB이며 이후 절대 변하지 않는다. 입력 쪽에서는 토큰의 코드가 학습된 작은 선형 맵(512→640)을 통과해 body로 들어가고, 출력 쪽에서는 body의 최종 상태가 다른 맵(640→512)을 거쳐 모든 코드와 내적을 하고 sqrt(512)로 나눠 점수화된다. 두 맵을 합쳐 0.65M 파라미터다. 여기에 학습 데이터의 로그 유니그램 빈도로 만든 토큰별 고정 prior(1MB, 학습 파라미터 없음)를 점수에 더하는데 이를 seal이라 부른다. 512개 방향은 E 분산의 92.3%를 담고, WordSim-353에서 코드의 유사도는 인간 평가와 0.617 상관을 보여 원래 float 임베딩의 0.614와 사실상 같다. 각 토큰의 float 최근접 10개 중 67%가 코드 최근접 10개 안에 들어온다. 테이블이 E의 함수이므로 학습 후에도 새 단어를 조각들의 행 평균에 해당하는 코드로 즉시 추가할 수 있다는 것이 학습 테이블과의 구조적 차이다.

무엇과 다른가

실험은 Gemma 3 270M 아키텍처(18층, 폭 640, gated GELU FFN 2,048, QK-norm을 쓰는 256차원 query 4개와 KV 1개, global layer마다 local 512토큰 슬라이딩 윈도 5개, RMSNorm)를 무작위 초기화에서 처음부터 학습하는 방식이다. 데이터는 FineWeb-Edu 2.05B 토큰, 시퀀스 길이 1,024, 옵티마이저 스텝당 262,144 토큰, 7,820 스텝이며 세 arm이 같은 시드에서 같은 배치를 같은 순서로 받는다. Arm A는 학습된 tied 테이블(268.1M 파라미터), Arm B는 고정 테이블과 고정 seal을 쓰는 동일한 float body(101.0M), Arm C는 body의 모든 선형층을 BitNet b1.58의 BitLinear로 바꾸고 seal을 유니그램 시작점에서 학습시키는 구성(101.2M)이다. arm당 RTX 5090 한 장에서 14~16시간, 약 18달러가 들었고 평가까지 포함한 전체 연구비는 75달러 미만이었다.

어떻게 쓰나

결과에서 학습 테이블이 사는 것은 검증 손실 0.20 nats(퍼플렉시티 약 23% 낮음)이고, 그 대가로 학습 파라미터가 2.6배다. 벤치마크에서는 수익이 훨씬 적다. B는 LAMBADA에서 A보다 1.6포인트 높고 HellaSwag와 PIQA에서는 동률이며 ARC-Easy에서는 3.4포인트 낮다. 168M 학습 파라미터의 총수익이 벤치마크 하나라는 뜻이다. 손실 곡선을 보면 B가 200 스텝에서 이미 0.23 nats 앞서지만 A가 600 스텝에 역전하고 3,000 스텝 부근부터 0.20 격차가 끝까지 유지된다. 빈도별로 쪼개면 학습 테이블은 100~100,000회 등장한 중간 빈도 토큰에서 이기고, 10~99회 등장 토큰에서는 고정 테이블이 2.3 nats(B), 1.7 nats(C) 앞선다. 전체 텍스트의 52%를 차지하는 최빈 161개 토큰에서는 격차가 0.035 nats에 불과하다. Gemma 어휘 262,144개 중 52,782개는 이 2.05B 토큰에 아예 등장하지 않고 10회 이상 등장하는 토큰은 147,187개뿐이라, A에서는 나머지 행이 무작위이거나 무작위에서 거의 움직이지 않은 상태로 남는다.

전제와 한계

ternary 모델 C는 가중치당 1.58비트로 약 40MB(테이블 16.8MB, ternary body 20MB, 16비트 인터페이스 맵 1.3MB, prior 0.5MB)이고, 실제로 실행하는 파일은 body를 2비트로 패킹하고 맵·norm·스케일을 32비트로 유지해 46.2MB다. 비교 대상인 A의 16비트 GGUF는 543MB, B는 약 220MB가 된다. C는 LAMBADA와 PIQA에서 A와 동률이고 HellaSwag에서 0.9포인트, ARC-Easy에서 4.9포인트 뒤진다. 노트북 CPU에서 C는 62MB 메모리로 초당 615~740 토큰을 생성해, 같은 머신·같은 8스레드에서 돌린 A의 8-bit GGUF보다 5.3~6.4배 빠르고 메모리는 7.6배 적게 쓴다. A는 어떤 양자화에서도 446MB 아래로 내려갈 수 없는데, llama.cpp가 262,144 x 640 테이블을 16비트로 유지해 4비트 파일의 250MB가 어휘이기 때문이다. 저자들이 공개한 AVX2 CPU 커널은 PyTorch 모델을 재현해 64개 검증 토큰에서 모든 위치의 argmax가 일치하고 다음 토큰 손실이 4.495 대 4.483이었다.

어휘 확장 실험도 있다. PubMed에서 Gemma 토크나이저가 두 조각 이상으로 쪼개는 생의학 단어 10,000개를 골라 단일 토큰으로 추가했을 때, 300개 홀드아웃 논문에서 새 토큰이 전체 토큰의 3.2%를 대체하고 품질 비용은 3%였다. 같은 조건에서 학습 테이블은 43% 무너진다. 다만 prior 설정과 코드 구성(조각 평균, 첫 조각, 모델 출력에서 읽은 코드)을 여러 방식으로 시도해도 비용을 0으로 만들지 못했고, 저자들은 이를 개선이 아니라 "가능성"에 대한 주장으로 한정한다. 20개 언어 실험에서 B는 아랍어·불가리아어·스페인어·러시아어·베트남어에서 A보다 낫고 나머지 15개 언어에서는 나쁘다. 학습에 없던 토큰은 seal이 log(1/N)로 벌점을 주는데, 그 prior를 본 토큰의 중앙값으로 바꾸면 태국어·아랍어·러시아어·프랑스어·그리스어에서 B가 앞선다. 반면 새 제어 토큰을 문맥 예시로 배치하는 테스트(5-shot에서 B 88%, C 75%, A 73%, 1-shot에서는 A가 최고)와 트윗 이모지 예측(모두 2.6~5.7%로 우연 수준)은 기대와 달리 고정 테이블에 유리하지 않았다.

개발자 입장에서 이 논문이 유용한 지점은 온디바이스 소형 모델의 병목이 연산이 아니라 파일 크기와 메모리 대역폭일 때다. 특히 llama.cpp를 포함한 대부분의 양자화 파이프라인이 임베딩과 출력 헤드를 16비트로 남겨둔다는 사실은, 4비트로 압축한 모델이 왜 여전히 수백 MB인지를 설명해 준다. 다만 고정 코드가 Gemma 3 270M의 임베딩에서 파생되므로 테이블과 B·C 체크포인트에 Google의 Gemma 약관이 함께 적용되고, 라이선스는 비상업적 연구용이다. 또한 토크나이저가 이미 커버하는 단어를 사후에 어휘로 추가하는 모든 설정에서 품질 비용이 발생했으므로, 도메인 어휘 확장을 기대만큼의 이득으로 계산하면 안 된다.

저자들이 밝힌 한계는 분명하다. arm당 시드가 하나이고 오차막대는 학습 런이 아니라 테스트 질문에 대한 95% 부트스트랩 범위다. 2B 토큰은 101M 파라미터 arm에는 Chinchilla 최적이지만 A의 테이블에는 부족한 양이라, 토큰당 데이터가 늘면 손실 격차가 0.20보다 커지고 벤치마크 양상도 달라질 수 있다. C는 B와 세 가지(body, BitNet 스케줄, 학습 prior)에서 동시에 다르며, prior의 몫은 1,833 스텝에서 중단한 별도 실행으로 0.003~0.022 nats라고 측정했지만 스케줄의 몫은 분리하지 않았다. CPU 비교는 노트북 한 대에서 서로 다른 두 코드베이스(자체 커널 대 llama.cpp)를 쓴 것이고 B용 커널은 없으며, 커널의 prefill이 토큰 하나씩이라 생성 속도만 비교했다.