COUNT는 호스트 RAM의 n-gram 카운트로 GPU 학습 비용을 줄인다
The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation
무엇인가
언어모델 스케일링은 오랫동안 연산량과 성능의 관계로 논의돼 왔고, 최근에는 메모리를 따로 키우는 설계가 주목받는다. MoE처럼 일부 파라미터만 활성화하는 방식이나 입력 임베딩을 늘리는 방식이 대표적인데, n-gram 임베딩은 한 번의 순전파에서 임베딩 테이블 전체가 아니라 조회된 행 하나만 처리하면 되기 때문에 연산량과 메모리 사용량을 분리해 키울 수 있다. 문제는 기존 방식이 임베딩을 처음부터 학습한다는 점이다. 가중치와 옵티마이저 상태를 빠르게 읽고 써야 해서 값비싼 GPU 메모리를 차지하고 백본 파라미터와 자리를 다툰다. 반면 호스트 시스템의 일반 RAM은 바이트당 저렴하고 이런 학습 파이프라인에서 거의 쓰이지 않고 남아 있다.
어떻게 동작하나
이 논문이 제안하는 COUNT(Counting Offline Unlocks N-gram Transformers)는 학습 파라미터를 하나도 추가하지 않고 n-gram 임베딩을 만든다. 미리 계산해 얼려둔 테이블 T_K가 저장된 n-gram 히스토리마다 다음 토큰 상위 K개와 그 점수를 담고 있다. 각 위치에서 현재 문맥의 가장 긴 저장 접미사를 찾아 조회하면 상위 K개 다음 토큰과 점수가 나오는데, 이 점수로 언어모델 자신의 언임베딩 행렬 B의 해당 행들을 가중합해 밀집 벡터 f_t를 만든다. 수식으로는 f_t = (1/J) Σ_j Σ_k S(w_k|h_t,j) B[w_k]이고, 이를 토큰 임베딩 e_t에 더해 첫 트랜스포머 블록의 입력 h_t^(0) = f_t + e_t로 쓴다. 점수는 stupid backoff로 계산한다. 관측된 continuation은 카운트 비율로, 없으면 0.4를 곱해 더 짧은 히스토리로 백오프하며 재정규화 없이 그대로 쓴다. 카운트 테이블은 읽기 전용이라 GPU 메모리를 차지하지 않고 호스트 RAM에 상주하며, 오프라인에서 Dolma 3 전체 코퍼스의 6차 n-gram까지 세어 만든다.
무엇과 다른가
실험은 OLMo-core 래더의 60M~1B 백본과 토큰 대 파라미터 비율 D/N ∈ {10, 20, 40}에서 이뤄졌고, allenai/dolma3_mix-6T-1025-7B로 사전학습하면서 같은 코퍼스로 n-gram 카운트를 뽑았다. 370M·600M·760M·1B 백본에는 797GB 테이블을, 60M·100M·190M에는 더 빡빡한 임계값으로 만든 17.1GB 테이블을 붙여 백본이 커질수록 테이블 용량도 함께 키웠다. 비교 대상은 표준 트랜스포머와 학습형 n-gram 임베딩(LongCat에 구현된 Over-Encoding, 바이그램·트라이그램에 차수당 해시 테이블 4개, 테이블 폭 d/8, 입력 임베딩 대 백본 파라미터 비율 1:1)이다. 평가는 Paloma 기반 11개 검증 소스의 평균 토큰 교차 엔트로피(LM 검증 손실), 74개 다운스트림 태스크 구성의 평균 BPB(태스크 손실), 41개 객관식 구성의 평균 정확도로 요약한다.
어떻게 쓰나
결과부터 보면, 같은 연산 예산에서 COUNT는 LM 손실을 46~89% 더 많은 연산을 쓴 것과 맞먹게 개선했고, 태스크 손실은 6~56% 더 많은 연산에 해당하는 개선을 냈다. 다운스트림 정확도도 거의 모든 예산 구간에서 표준 트랜스포머를 앞섰다. 학습 파라미터 수를 축으로 놓은 파레토 비교에서는 프런티어가 대부분 COUNT를 따라간다. 학습형 n-gram 임베딩이 비슷한 성능을 내는 지점도 있지만, 그 대가는 GPU 메모리에 상주해야 하는 추가 학습 파라미터다. 출력 단계에서 카운트 분포를 융합하는 late fusion과 비교하면 COUNT의 입력 단계 융합이 범용 성능 개선이 더 컸고, 사전학습 중 late fusion을 켜는 것은 추론 시점에만 켜는 것보다 나은 점이 없었다.
전제와 한계
COUNT 테이블은 학습된 임베딩과 달리 공유 어휘에 대한 다음 토큰 점수를 저장하고 같은 언임베딩 기저로 임베딩을 만들기 때문에 추론 시점에 테이블을 갈아끼울 수 있다. 수학(83.398GB, Dolmino Math 10.66B 토큰), 생의학(81.723GB, 95.39B 토큰), 법률(80.151GB, 12.34B 토큰) 도메인 테이블을 만들어 원래의 797.487GB 일반 테이블과 점수를 S_λ = (1-λ)S_general + λS_domain으로 보간했다. 세 도메인 모두 손실이 줄어드는 λ=0.25에서 수학과 법률은 손실·정확도가 개선됐지만 생의학은 결과가 엇갈렸다. 참고로 D/N=20에서 일반 테이블만 쓴 COUNT가 표준 1B 트랜스포머보다 태스크 손실을 0.010 BPB 낮추는데, 이는 적당한 도메인 가중으로 얻는 개선과 비슷한 크기다. 배포 환경의 RAM이 학습 때보다 작을 때를 위한 테이블 축소도 실험했는데, 459.222·142.037·68.458·17.114GB 테이블로 바꿔도 1B 백본에서는 저장량을 10분의 1 이상 줄이면서 표준 트랜스포머보다 나은 성능을 유지했다.
처리량은 1B 백본 체크포인트를 8개 H200에서 돌려 최적화 스텝 50~100의 중앙값 TPS로 쟀다. COUNT는 표준 트랜스포머 학습 처리량의 93.1~97.8%, 프롬프트 프리필의 95.5~95.8%, 자기회귀 디코딩의 95.9~96.7%를 유지했다. 이를 위해 컴파일된 멀티스레드 호스트 조회, GPU가 이전 배치를 처리하는 동안 다음 조회를 미리 가져오는 프리페치, 후보 벡터를 따로 저장하지 않고 모아서 합치는 fused embedding reduction을 썼다. 학습형 n-gram 임베딩 베이스라인도 조회된 행만 갱신하는 희소 옵티마이저로 최적화했는데, 190M 백본에 17.1GB 테이블을 붙인 경우 전체 테이블 갱신은 표준 대비 4.4% 처리량에 그친 반면 조회 행만 갱신하면 96.3%까지 회복됐다.
개발자 입장에서 이 논문이 주는 실용적 선택지는 두 가지다. 첫째, GPU 메모리가 빠듯할 때 학습형 n-gram 임베딩으로 GPU 메모리를 쪼개 쓰는 대신 COUNT로 백본 파라미터가 GPU 메모리를 채우게 하는 편이 성능이 좋다. 둘째, 도메인 데이터가 적거나 재학습 예산이 없을 때 도메인 코퍼스로 만든 카운트 테이블을 추론 시점에 끼워 넣는 방식으로 적응을 시도할 수 있다. 도입 전에 확인할 것은 세 가지다. 배포 노드의 호스트 RAM이 테이블을 감당하는지(학습 때 쓴 테이블을 그대로 못 쓸 수 있다), 조회·프리페치 구현이 실제 서빙 지연에 얼마를 더하는지, 그리고 도메인 테이블을 만들 코퍼스의 라이선스와 민감도다.
저자들이 직접 밝힌 전제와 한계도 분명하다. 카운트 테이블은 원문서가 아니라 집계 통계라서 공유가 덜 민감하고 더 압축적이지만, 희귀 문맥에 대한 고차 카운트는 여전히 원문 조각을 드러낼 수 있으므로 다른 파생 통계와 같은 주의를 기울여 공개해야 한다고 적었다. 테이블 용량은 백본 크기에 맞춰 키우는 설계이고, 학습 시점의 호스트 RAM이 배포 시점보다 클 수 있다는 전제도 깔려 있다. 도메인 적응 효과는 수학·법률에서 뚜렷했지만 생의학에서는 엇갈린 결과가 나왔다. 실험 규모는 60M~1B 백본과 D/N 10~40에 한정돼 있어, 그보다 큰 모델에서 같은 이득이 유지되는지는 이 논문이 답하지 않는다.