Apple 실리콘을 위한 초고속 차원 축소와 시각화 라이브러리 mlx‑vis

초록

mlx‑vis는 Apple Silicon의 Metal GPU와 MLX 프레임워크만을 이용해 UMAP, t‑SNE, PaCMAP, LocalMAP, TriMap, DREAMS, CNE, MMAE 등 8가지 차원 축소 알고리즘과 NNDescent 기반 k‑NN 그래프 구축을 GPU 전용으로 구현한다. 70 K Fashion‑MNIST 데이터를 M3 Ultra에서 2‑4.7초 내에 임베딩하고, 800프레임 애니메이션을 1.4초에 렌더링한다. 의존성은 MLX와 NumPy뿐이며, GPU‑네이티브 렌더러가 하드웨어 H.264 인코딩을 활용한다.

상세 분석

mlx‑vis는 Apple Silicon의 통합 메모리와 Metal GPU를 완전히 활용하도록 설계된 최초의 차원 축소·시각화 라이브러리이다. 핵심 설계는 모든 연산을 MLX 텐서로 표현하고, @mx.compile를 통한 JIT 컴파일로 GPU 커널을 자동 융합한다는 점에 있다. 이를 통해 기존 CPU‑기반 구현에서 발생하던 파이썬 루프와 데이터 복사 오버헤드를 근본적으로 제거한다.

첫 단계인 NNDescent는 거리 계산을 ‖a‑b‖² = ‖a‖² + ‖b‖² – 2 aᵀb 형태의 행렬 곱으로 전환하고, mx.argpartition을 이용해 부분 정렬만 수행한다. 이는 70 K 포인트에 대해 500 epoch 동안 평균 3 GB 이하의 GPU 메모리와 40‑80 W 전력만을 소모한다는 실험 결과와 일치한다.

각 차원 축소 알고리즘은 원 논문의 수식과 업데이트 규칙을 그대로 구현했지만, GPU 친화적으로 변형하였다. 예를 들어 UMAP는 Gauss‑Newton 최적화를 직접 구현해 scipy 의존성을 없앴으며, t‑SNE는 FFT 기반 O(n log n) 반발력 계산을 MLX에서 수행한다. LocalMAP은 동적 로컬 그래프 재샘플링을 mx.argsort 기반 후보 선택으로 대체해 파이썬 루프를 제거했고, CNE와 MMAE는 손실 계산을 정적 함수로 컴파일해 연산자를 한 번에 실행하도록 했다.

시각화 파이프라인은 전통적인 matplotlib 대신 GPU에서 원형 스플래팅(circle‑splatting) 방식을 채택한다. 각 포인트는 반경 R 내 픽셀 오프셋을 선형 가중치 w = max(0, 1‑r/R) 로 계산하고, mx.array.at