임베딩은 과한 비용일까? 고정 토큰 코드로 소형 LM을 46MB로

배경

소형 언어 모델에서는 어휘를 저장하는 임베딩 테이블이 트랜스포머 본체보다 커질 수 있습니다. Gemma 3 270M의 경우 전체 2억6800만 개 파라미터 중 약 1억6800만 개가 임베딩 테이블에 들어갑니다. 다른 층만 양자화하면 임베딩은 더 높은 정밀도로 남는 경우가 많아, 최종 모델 파일에서 차지하는 비중도 커집니다.

Hugging Face Daily Papers에 소개된 연구는 이 구조가 반드시 필요한지 검토했습니다. 연구진은 학습 가능한 임베딩 행렬 대신 각 토큰에 512비트 고정 코드를 배정했습니다. 코드는 사전 학습된 임베딩을 주성분 투영한 뒤 값의 부호를 이진화해 생성합니다. 초기화가 끝나면 이 어휘 테이블은 학습 과정에서 업데이트되지 않습니다.

핵심 결과

• 학습 파라미터 감소: 연구진은 동일한 Gemma 형태의 모델을 같은 20억5000만 토큰으로 처음부터 학습했습니다. 학습 가능한 테이블은 검증 손실을 0.20 nat 개선했지만, 약 2.6배 많은 학습 파라미터를 요구했습니다.
• 대부분의 평가에서 경쟁력: 고정 테이블 모델은 LAMBADA, HellaSwag, PIQA에서 학습 가능한 테이블 모델을 앞서거나 동률을 기록했습니다. ARC-Easy에서는 학습 가능한 테이블이 더 나았습니다.
• 희귀 단어와 신규 토큰에 강함: 희귀 단어 평가에서 고정 방식은 1.7~2.3 nat 앞섰습니다. 학습이 끝난 뒤 1만 개 토큰을 추가했을 때 성능 비용은 약 3%였지만, 학습 가능한 테이블은 약 43% 하락했습니다.
• 온디바이스 배포 가능성: BitNet b1.58 삼진화 본체와 결합하면 전체 모델 크기는 약 46MB입니다. 노트북 CPU에서 초당 615740토큰을 생성하고 메모리 사용량은 약 62MB였습니다. llama.cpp의 부동소수점 모델 8비트 GGUF와 비교하면 속도는 약 56배, 메모리는 약 7분의 1이라고 보고됐습니다.

의미와 한계

이 결과는 소형 언어 모델이 대규모 학습 가능 어휘 테이블을 반드시 필요로 한다는 통념에 의문을 제기합니다. 고정 코드를 사용하면 저장 공간을 트랜스포머 본체에 배분할 수 있고, 희귀 단어나 학습 후 어휘 확장에서도 안정적인 동작을 기대할 수 있습니다. 스마트폰과 노트북에서 실행되는 오프라인 모델에 특히 적합한 방향입니다.

다만 임베딩이 모든 모델에서 불필요하다는 뜻은 아닙니다. 이번 실험은 작은 모델과 특정 학습 토큰 예산을 대상으로 했습니다. 모델과 데이터가 커질수록 학습 가능한 임베딩의 유연성이 중요해질 수 있으며, 고정 표현이 병목이 되는 시점도 아직 확인되지 않았습니다. 언어, 문자 체계, 분야, 희귀 문자별 성능을 평균 검증 손실과 별도로 분석할 필요도 있습니다.

후속 연구에서는 새 토큰의 코드를 어떤 방식으로 만들지, 투영을 다시 계산해야 하는지, 더 긴 학습에서 고정 테이블이 본체의 학습을 제한하는지를 검증해야 합니다. 현재로서는 임베딩을 전면 부정하는 결론이라기보다, 저비트 소형 모델을 위한 유력한 온디바이스 설계로 보는 편이 정확합니다.

출처: Hugging Face Daily Papers