Papers·4일 전
IBM Research, 단일 GPU에서 GMM 20배 가속하는 Flash-GMM 커널 공개 — ANN 검색에서 k-means 대체 가능
IBM Research 팀이 단일 GPU pass로 GMM(Gaussian Mixture Model)을 20배 빠르게 계산하는 fused Triton 커널 Flash-GMM을 공개했습니다. 책임 행렬(responsibility matrix)을 GPU 메모리에 전부 올리지 않도록 설계해 기존 대비 100배 이상 큰 데이터셋도 학습 가능합니다. IVF coarse quantizer에 적용한 결과, soft GMM 클러스터링이 k-means를 대체해 동일 recall 목표에서 거리 계산을 최대 1.7배 줄이거나, 동일 계산량에서 recall@10을 2~12 포인트 개선했습니다. 한계로는 GMM의 EM 수렴이 k-means보다 느릴 수 있다는 점이 있으며, 코드는 오픈소스로 공개되었습니다.
IBM Research가 단일 GPU에서 GMM을 20배 빠르게 계산하는 fused Triton 커널을 내놨습니다. ANN 검색에서 k-means를 대체할 가능성도 열었네요.
핵심 결론
- 속도 — 기존 GMM 구현 대비 20배 가속, 단일 GPU에서 처리 가능한 데이터셋 크기가 100배 이상 증가.
- ANN 적용 — IVF coarse quantizer에 soft GMM 클러스터링을 적용해 k-means 대비 동일 recall에서 거리 계산 1.7배 감소, 또는 동일 계산량에서 recall@10 2~12 포인트 향상.
방법
- Fused kernel — Triton으로 작성된 fused 커널로, E-step과 M-step을 하나의 GPU pass로 처리하며 책임 행렬을 메모리에 구체화하지 않음.
- 소프트 할당 — GMM의 soft assignment로 경계 벡터를 여러 클러스터에 할당해 검색 정확도를 높임.
한계·조건
- 수렴 속도 — EM 알고리즘 특성상 k-means보다 수렴이 느릴 수 있어, 일부 시나리오에서는 학습 시간이 더 걸릴 수 있습니다.
- 코드 — 오픈소스로 공개되어 재현 가능.
편집자 한 줄
GMM의 soft clustering이 ANN 검색에서 실질적 이득을 보여준 점이 흥미롭습니다. 다만 EM 수렴 문제는 실제 배포 시 고려해야 할 요소네요.
- #gmm
- #ann
- #triton
- #ibm-research
- #kernel-fusion
IBM Research