VkSplat: High-Performance 3DGS Training in Vulkan Compute

以下は、論文「VkSplat: High-Performance 3DGS Training in Vulkan Compute」の詳細な技術的サマリーです。

1. 問題提起

高品質な新規視点合成のための 3D ガウシアンスプラッティング（3DGS）の急速な普及にもかかわらず、実用的な展開は 2 つの主要な制限によって妨げられています。

• パフォーマンスのボトルネック: 既存のトレーニングパイプラインは、特に大規模なシーンにおいて、速度とメモリ使用量の面で最適化が不十分です。主なボトルネックには、ソート中の誤った正のタイル交差、後方ラスタライゼーション中のアトミック競合、およびオプティマイザ（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 ベースのベースラインを上回ることを可能にする 5 つの主要な技術的最適化について詳述しています。

A. スキャンライン交差による完全なタイルカリング

• 問題: 既存の手法（例：[KKLD23], [RSP∗24]）は、実際にはガウシアンと重ならないタイルをチェックする誤った正のタイル交差を生成し、ソートおよびラスタライゼーションの時間を浪費しています。
• 解決策: VkSplat は、正確な交差を計算するためにスキャンライン定式化を使用します。ガウシアンを楕円として扱い、交差する行/列の閉形式の区間を計算することで、システムは正確なタイル - 深度ペアのリストを生成します。
• 影響: 個別のカリングパスを必要とせずに誤った正を排除し、VRAM 使用量とソートのオーバーヘッドを削減します。

B. 適応的ラスタライゼーション後方

• 問題: 標準的なピクセルごとの後方パスは、ガウシアン勾配を更新する際に高いアトミック競合を引き起こします。
• 解決策: 著者らは 2 つの戦略を実装し、シーンに基づいて最も高速なものを動的に選択するためにトンプソンサンプリングスケジューラを使用します。
  1. タイルごとの並列性: 1 つのタイルあたり 1 つのスレッドブロックを持ち、遅延を最小化するために動的に調整されたガウシアンバッチサイズ（$S = \sqrt{NP}$ に最適化）を使用します。
  2. 共有メモリ最適化: 2 パスアプローチ（前方パスで共有メモリに透過度を計算し、後方パスでガウシアンごとに勾配を蓄積）を採用し、スレッドの発散を最小化します。
• 影響: アトミック競合を大幅に削減し、スレッド利用率を向上させます。

C. 融合された投影後方とオプティマイザ

• 問題: 標準的な実装は勾配を個別に保存し、PyTorch の Adam を使用するため、不要なメモリフットプリントとカーネル起動が発生します。
• 解決策: 投影後方パスと Adam オプティマイザは単一のカーネルに融合されます。
  • スケールと不透明度の値および勾配は、カーネル内で直接 log/logit 空間間を変換されます。
  • 球面調和関数（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 倍の高速化を達成しました（例：デフォルトの増殖において 412 秒対 1384 秒）。
  • VRAM: メモリ使用量を**33%**削減しました（3.01 GiB 対 4.56 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 トレーニングなど、成熟したトレーナーに見られるいくつかの「生産準備完了」機能がありません。ただし、コアアーキテクチャはこれらの機能を受け入れるように設計されています。