21M 모델에 6.4B 파라미터 룩업 테이블을 붙여봤습니다. 114M 밀집 모델과 대등하며 SSD에서 구동됩니다 (RX 9070)
I gave a 21M model a 6.4B-parameter lookup table. It matches a 114M dense model and runs with the table on an SSD (RX 9070)
핵심 요약
소형 모델에 거대한 룩업 테이블을 결합해 성능을 높이고, SSD 오프로딩으로 VRAM 사용량을 0.4GB까지 줄인 실험입니다.
- 모델 구조 — 21M 파라미터 모델에 6.4B 파라미터 룩업 테이블을 결합함
- 성능 효율 — 114M 밀집 모델과 대등한 성능을 보임
- SSD 오프로딩 — NVMe SSD를 활용해 VRAM 사용량을 0.4GB로 최소화함
- 기술적 도전 — Triton 커널을 직접 작성하여 다양한 GPU 환경에서 구동함
지난 몇 주 동안 취미로 연구 프로젝트 하나를 진행했는데, 방금 공개했다.
아이디어 자체는 새로운 게 아니다(product-key memory, Lample et al. 2019, 그리고 Meta의 "Memory Layers at Scale"). 모델한테 학습된 벡터로 가득 찬 거대한 테이블을 주고, 토큰 하나당 수백 개 정도만 읽게 만드는 방식이다. 이게 작은 모델에서 실제로 얼마나 가치가 있는지, 비용은 얼마나 드는지, 그리고 테이블이 꼭 VRAM에 올라가 있어야 하는지 궁금했다.
결과는 다음과 같다:
- 1,680만 행짜리 테이블을 가진 21M 모델(테이블 파라미터 64억 개, 토큰당 3,300만 개 사용)은 같은 5억 개의 위키피디아 토큰으로 학습시킨 114M 덴스 모델과 성능이 비슷하다.
- 테이블이 VRAM에 있을 필요가 없다. NVMe SSD에서 4비트 테이블을 메모리 맵으로 불러와도 내 RX 9070에서 VRAM 0.4GB만 쓰면서 초당 약 140토큰을 뽑아낸다. 다만 긴 프롬프트를 SSD에서 읽어올 때는 좀 느린데, 행 하나를 놓칠 때마다 4KB 페이지 전체를 불러와야 해서 그렇다.
- Triton 커널도 직접 짰다. 내 라데온 그래픽카드, MI350X, 그리고 H100/H200에서 수정 없이 그대로 돌아간다.
- 이미 완성된 모델(Qwen3.5-0.8B)에 테이블을 갖다 붙이는 건 효과가 없었다. 같은 연산량을 쓰는 작은 덴스 애드온보다 나을 게 하나도 없더라.
주의할 점: 모델이 아주 작고, 대규모 실행은 시드 하나로만 돌렸으며, 모델이 뱉는 텍스트는 위키피디아 영어처럼 유창하긴 한데 내용은 죄다 지어낸 거다. 실행할 때마다 성공 기준을 미리 적어뒀고, 안 된 것들도 전부 기록해 뒀다.
대부분 내 게이밍 PC에서 돌렸고, 큰 규모의 실행은 Runpod에서 70달러 정도 썼다. Claude Code랑 같이 만들었는데(커밋 보면 알 거다), 아이디어나 결정, 돈은 다 내 거다.
저장소: https://github.com/re133/sparse-memory-lm
단어를 클릭하면 모델이 어떤 테이블 항목을 읽는지 볼 수 있다: https://re133.github.io/sparse-memory-lm/explorer/
모델: https://huggingface.co/fechyy/sparse-memory-lm-B-16M
피드백 환영한다. 특히 내가 잘못 알고 있는 게 있다면 꼭 알려줘라. 혹시 남는 고성능 GPU 있는 사람 있으면 1B 규모로도 한번 돌려보고 싶다.
