VkSplat: High-Performance 3DGS Training in Vulkan Compute

"VkSplat: High-Performance 3DGS Training in Vulkan Compute" 논문에 대한 상세한 기술 요약은 다음과 같습니다.

1. 문제 제기

고품질의 새로운 뷰 합성을 위한 3D 가우시안 스플래팅 (3DGS) 의 빠른 채택에도 불구하고, 실제 배포는 두 가지 주요 제한 사항으로 인해 방해받고 있습니다:

• 성능 병목 현상: 기존 학습 파이프라인은 특히 대규모 장면의 경우 속도와 메모리 사용 측면에서 종종 최적화되지 않았습니다. 주요 병목 현상으로는 정렬 중 발생하는 잘못된 양의 타일 교차 (false-positive tile intersections), 역방향 래스터화 중 발생하는 원자적 경쟁 (atomic contention), 그리고 옵티마이저 (예: PyTorch 의 Adam) 의 비효율적인 메모리 처리 등이 있습니다.
• 하드웨어 종속성: 대부분의 최첨단 구현은 CUDA 생태계와 PyTorch 에 크게 의존하여 학습을 NVIDIA GPU 로 제한합니다. 이는 AMD, Intel 또는 모바일 GPU 를 가진 사용자의 접근성을 제한합니다.
• 분열: 크로스 벤더 API(예: Vulkan) 를 사용하려는 이전 시도들은 래스터화만 해결했거나 (엔드투엔드 학습이 아님) 학습 루프에 PyTorch 를 의존하여 CUDA 의존성을 완전히 제거하지 못했습니다.

2. 방법론

저자들은 VkSplat을 제시합니다. 이는 Vulkan Compute로 완전히 구현된 엔드투엔드 3DGS 학습 파이프라인입니다. 이 시스템은 크로스 벤더 호환성을 보장하면서도 높은 성능을 유지하기 위해 (Vulkan 를 대상으로 하는) Slang 쉐이딩 언어를 기반으로 구축되었습니다.

핵심 방법론은 특정 아키텍처 변경을 통해 3DGS 파이프라인을 완전히 다시 작성하는 것을 포함합니다:

• 완전한 Vulkan Compute: 초기화 및 밀도 증가 (densification) 에서 최적화에 이르기까지 전체 파이프라인이 CPU 측 개입이나 PyTorch 의존성 없이 GPU 에서 실행됩니다.
• Slang 백엔드: Slang 을 활용하면 코드가 여러 백엔드를 대상으로 할 수 있지만, 현재 구현은 Vulkan 에 집중되어 있습니다.

3. 주요 기술적 기여

이 논문은 VkSplat 이 CUDA 기반 베이스라인을 능가할 수 있게 하는 다섯 가지 주요 기술적 최적화를 상세히 설명합니다:

A. 스캔라인 교차를 통한 완전한 타일 커팅

• 문제: 기존 방법들 (예: [KKLD23], [RSP∗24]) 은 실제로 가우시안과 겹치지 않는 타일을 확인하는 잘못된 양의 타일 교차 (false-positive tile intersections) 를 생성하여 정렬 및 래스터화 시간을 낭비합니다.
• 해결책: VkSplat 은 **스캔라인 공식 (scan-line formulation)**을 사용하여 정확한 교차를 계산합니다. 가우시안을 타원체로 간주하고 교차하는 행/열에 대한 폐쇄형 구간을 계산함으로써, 시스템은 정확한 타일 - 깊이 쌍 목록을 생성합니다.
• 영향: 별도의 커팅 패스를 요구하지 않고 잘못된 양을 제거하여 VRAM 사용량과 정렬 오버헤드를 줄입니다.

B. 적응형 래스터화 역방향

• 문제: 표준 픽셀 단위 역방향 패스는 가우시안 기울기를 업데이트할 때 높은 원자적 경쟁을 유발합니다.
• 해결책: 저자들은 두 가지 전략을 구현하고 장면 기반으로 가장 빠른 것을 동적으로 선택하는 Thompson 샘플링 스케줄러를 사용합니다:
  1. 타일 단위 병렬화: 지연 시간을 최소화하기 위해 동적으로 조정된 가우시안 배치 크기 (최적화된 $S = \sqrt{NP}$) 를 가진 타일당 하나의 스레드 블록.
  2. 공유 메모리 최적화: 전방향 패스에서 공유 메모리에 투과율을 계산하고, 역방향 패스에서 가우시안당 기울기를 누적하는 2 패스 접근 방식으로 스레드 분기를 최소화합니다.
• 영향: 원자적 경쟁을 현저히 줄이고 스레드 활용도를 향상시킵니다.

C. 퓨즈드 프로젝션 역방향 및 옵티마이저

• 문제: 표준 구현은 기울기를 별도로 저장하고 PyTorch 의 Adam 을 사용하여 불필요한 메모리 풋프린트와 커널 실행을 초래합니다.
• 해결책: 프로젝션 역방향 패스와 Adam 옵티마이저가 단일 커널로 퓨즈됩니다.
  • 스케일과 불투명도의 값 및 기울기는 커널 내에서 로그/로그이트 공간 간에 직접 변환됩니다.
  • 구형 조화 (Spherical Harmonics, SH) 계수는 메모리 병합을 개선하기 위해 서브그룹 크기 (128 비트 값) 에 정렬된 열 기반 형식으로 저장됩니다.
• 영향: 기울기를 위한 중간 메모리 저장을 제거하고 커널 실행 오버헤드를 줄입니다.

D. 32 비트 타일 - 깊이 정렬

• 문제: Vulkan 의 네이티브 정렬은 종종 CUDA 의 최적화된 내장 함수보다 느립니다.
• 해결책: 64 비트 키 (32 비트 타일 ID + 32 비트 깊이) 대신 32 비트 키를 사용합니다. 저자들은 깊이 $z$를 FP32 부동소수점의 하위 23 비트 만변하는 범위로 매핑하고, 타일 ID 를 상위 비트에 패킹합니다.
• 영향: 깊이 정밀도나 이미지 품질을 희생하지 않고 Vulkan 하드웨어에서 효율적인 32 비트 정렬을 가능하게 합니다.

E. 완전히 퓨즈된 손실 기울기 평가

• 해결책: 단일 커널이 L1 및 SSIM 손실의 가중 합과 그 기울기를 직접 계산합니다.
• 최적화: 참조 이미지는 UINT8 RGBA 로 저장되어 FP32 변환을 방지하고, 알파 마스킹은 거의 오버헤드 없이 처리됩니다.
• 영향: 중간 축소 단계와 메모리 레이아웃 변환 (채널 마지막에서 채널 첫 번째로) 을 제거합니다.

4. 결과

저자들은 NVIDIA RTX 3090 에서 GSplat베이스라인 (주요 CUDA+PyTorch 구현) 과 비교하여 Mip-NeRF 360데이터셋의 7 개 장면에서 VkSplat 을 평가했습니다.

• 성능:
  • 속도: GSplat 베이스라인 대비 3.3 배 속도 향상을 달성했습니다 (예: 기본 밀도 증가의 경우 1384 초 대비 412 초).
  • VRAM: 메모리 사용량을 **33%**줄였습니다 (4.56 GiB 대비 3.01 GiB).
  • 세분화: VkSplat 은 모든 개별 파이프라인 단계 (프로젝션, 타일링, 래스터화, 손실, 옵티마이저) 에서 더 빨랐습니다.
• 품질:
  • 이미지 품질 지표 (PSNR, SSIM, LPIPS) 는 베이스라인과 동일하여 최적화가 충실도를 훼손하지 않았음을 확인했습니다.
• 크로스 벤더 호환성:
  • AMD Radeon RX 7800 XT에서 성공적으로 학습되었습니다.
  • AMD 는 전반적으로 느렸지만 (NVIDIA 의 575 초 대비 1201 초), 일관된 품질 지표와 VRAM 사용량을 보였습니다. 성능 격차는 주로 AMD 의 느린 호스트 - 장치 메모리 전송에 기인한 것으로, 추가적인 하드웨어별 튜닝의 여지가 있음을 시사합니다.

5. 의의

• 최초의 엔드투엔드 Vulkan 3DGS: 이는 CUDA 독점을 깨고 최첨단 성능을 달성한 최초의 완전 Vulkan 기반 3DGS 학습 파이프라인입니다.
• 하드웨어 무관성: PyTorch 와 CUDA 의존성을 제거함으로써 VkSplat 은 AMD, Intel 및 모바일 GPU 에서 3DGS 학습을 가능하게 하여 고품질 3D 재구성에 대한 접근성을 민주화합니다.
• 성능 리더십: 역설적으로, Vulkan 구현은 매우 최적화된 CUDA 베이스라인보다 성능이 뛰어납니다. 이는 신중한 저수준 GPU 최적화 (퓨즈드 커널, 정확한 커팅, 적응형 스케줄링) 가 비 NVIDIA 하드웨어에서도 이점을 얻을 수 있음을 증명합니다.
• 확장성: Slang 의 사용은 파이프라인이 Metal, DirectX 또는 WebGPU 와 같은 다른 백엔드로 쉽게 이식될 수 있음을 시사하여 생태계를 더욱 확장합니다.

한계점: 현재 구현은 노출 보정, 깊이 감독, 배치 처리, 멀티 GPU 학습과 같은 성숙한 학습기에서 발견되는 일부 "프로덕션 준비" 기능이 부족합니다. 그러나 핵심 아키텍처는 이러한 기능을 수용하도록 설계되었습니다.