GPU용 초고속 희소 스케치 설계와 구현

초록

본 논문은 GPU에서 희소 Johnson‑Lindenstrauss 변환(SJLT)의 적용 속도를 크게 향상시키기 위해 스케치 구조와 CUDA 커널을 공동 설계한다. 블록‑퍼뮤테이션 방식으로 희소성을 제어한 BlockPerm‑SJLT를 제안하고, 이를 효율적으로 실행하는 FlashSketch 커널을 구현한다. 이 설계는 블록 수준의 규칙적인 메모리 접근을 가능하게 하여 전역 원자 연산을 최소화하고, 이론적 OSE 보장을 유지하면서 실험적으로 1.7배 정도의 평균 가속을 달성한다.

상세 분석

FlashSketch 논문은 GPU 아키텍처의 메모리 계층 구조와 병렬 실행 모델을 깊이 고려한 희소 스케치‑커널 공동 설계(co‑design) 사례를 제시한다. 기존 SJLT와 OSNAP은 열당 고정된 비영(非零) 개수 s만을 갖는 완전 무작위 희소 행렬을 사용한다. 이러한 무작위성은 이론적 보장은 뛰어나지만, GPU에서의 메모리 접근 패턴을 불규칙하게 만들어 전역 메모리 대역폭 활용을 저해한다. 논문은 이 문제를 해결하기 위해 입력 차원을 M개의 블록으로, 출력 차원도 동일하게 M개의 블록으로 나눈 뒤, 각 출력 블록이 κ개의 입력 블록과만 연결되도록 블록‑레벨 퍼뮤테이션 집합 {π₁,…,π_κ}을 샘플링한다. 여기서 퍼뮤테이션은 서로 간에 엣지‑디스조인트(edge‑disjoint) 조건을 만족해 동일한 입력‑출력 블록 쌍이 중복되지 않으며, 이는 블록‑레벨 그래프가 κ‑정규(biregular)임을 보장한다.

블록 내부에서는 기존 SJLT와 동일하게 각 열에 s개의 비영을 배치하고, 부호는 ±1/√s 로 정규화한다. 이렇게 하면 블록 내부에서는 충분한 혼합(mixing)이 유지되면서, 블록 간에는 퍼뮤테이션에 의해 제한된 연결만 존재하므로 각 출력 블록은 하나의 CUDA 스레드‑블록에 독점적으로 할당될 수 있다. 결과적으로 FlashSketch 커널은 입력 타일을 공유 메모리로 스트리밍하고, 공유 메모리 원자 연산을 이용해 출력 타일을 누적한 뒤, 전역 메모리로 한 번만 쓰는 구조를 구현한다. 전역 원자 연산이 사라짐에 따라 메모리 충돌과 직렬화 비용이 크게 감소하고, 공유 메모리의 높은 대역폭을 활용해 메모리 대역폭 효율을 극대화한다.

이론적 측면에서는 블록‑퍼뮤테이션 구조가 “neighborhood‑coherence”라는 새로운 파라미터를 도입한다. κ가 클수록 각 출력 블록이 더 많은 입력 블록을 섞어주어 OSE 보장이 강화되지만, 동시에 블록 간 연결이 늘어나면서 메모리 접근 패턴이 복잡해져 성능이 저하될 수 있다. 논문은 독립적인 무작위 퍼뮤테이션 모델을 가정하고, κ가 증가할수록 블록‑코히어런스 μ_blk이 감소함을 보이며, 이는 기존 블록‑다이어그램 SJLT의 보장과 일관된다. 따라서 κ는 GPU 효율성과 스케치 정확도 사이의 트레이드오프를 조정하는 튜닝 파라미터가 된다.

실험에서는 RTX 4090을 이용해 다양한 RandNLA 작업(최소제곱, 저‑랭크 근사, 회귀)과 데이터 귀속 파이프라인인 GraSS에 FlashSketch를 적용하였다. 비교 대상은 cuBLAS 기반의 dense Gaussian, cuSPARSE 기반의 일반 SJLT, 그리고 기존 CountSketch 구현이다. 결과는 전반적으로 1.5 ~ 2.0배의 속도 향상을 보이며, 특히 스케치 차원 k가 작을 때(예: k = 64 ~ 512) 가장 큰 이득을 얻었다. 품질 측면에서는 Gram 행렬 근사 오차가 기존 희소 스케치와 거의 동일하거나 약간 개선되는 수준으로, OSE 보장을 실험적으로도 확인했다.

이 논문은 희소 스케치 설계에 하드웨어 친화적 구조를 도입함으로써, 이론적 보장을 포기하지 않으면서도 GPU에서 실질적인 가속을 달성한 점이 가장 큰 공헌이다. 또한 블록‑퍼뮤테이션이라는 개념은 다른 종류의 랜덤 변환(예: 서브샘플드 Hadamard)이나 다른 하드웨어(예: TPU)에도 확장 가능성을 시사한다.