[논문] COUNT: N-gram 카운팅을 통한 호스트 메모리 확장 및 학습 없는 적응
[Paper] The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation
핵심 요약
학습된 파라미터 없이 N-gram 빈도를 호스트 메모리에 저장하여 언어 모델의 성능을 효율적으로 확장하는 COUNT 기법을 소개합니다.
- COUNT 기법 — 학습된 파라미터 없이 N-gram 빈도를 호스트 메모리에 저장하여 효율성 극대화
- 성능 향상 — 1B 파라미터 모델에서 57% 더 많은 연산량을 투입한 것과 유사한 검증 손실 개선
- 메모리 최적화 — GPU 메모리 부족 시 백본 파라미터에 자원을 집중하여 성능 향상
- 학습 없는 적응 — 추론 시 도메인별 N-gram 테이블을 교체하여 추가 학습 없이 작업 적응 가능
The scaling relationship of compute and performance for language models (LMs) has been extensively studied; more recently, models are being designed to scale memory use, as well. In particular, adding embeddings for n-grams scales memory separately from compute use, because each pass only processes individual embedding rows, which are much smaller than the storage used for the whole embedding table. However, existing methods train embeddings from scratch, which needs expensive GPU memory for fast reading and writing. Meanwhile, the ordinary memory of the host system is cheap and underutilized by these training methods. To scale on host memory, we introduce COUNT (Counting Offline Unlocks N-gram Transformers), which builds n-gram embeddings without adding learned parameters by storing precomputed n-gram counts on host. We use the corpus frequencies to weight together the LM's own unembeddings of the next tokens most often following the current n-gram history. At a fixed compute budget, COUNT can be added as an additional input embedding for pretraining 1B-parameter LMs to improve validation loss as much as an LM with 57% more compute. If GPU memory is limited, using COUNT to let backbone parameters fill GPU memory gives better performance than splitting memory with additional learned n-gram embeddings. Beyond efficiency, COUNT allows switching in a different table of n-gram counts from a domain-specific corpus at inference time to adapt to a new task without any training.
-
arXiv : https://arxiv.org/abs/2610.20456
-
GitHub :
