From Theory to Throughput: CUDA-Optimized APML for Large-Batch 3D Learning

TL;DR

This paper introduces CUDA-APML, a GPU-optimized, sparse implementation of APML that significantly reduces memory usage while maintaining accuracy for large-batch 3D point cloud learning tasks.

Contribution

It presents a novel sparse CUDA implementation of APML that scales near-linearly in memory, enabling efficient large-batch 3D learning without sacrificing gradient quality.

Findings

Reduces peak GPU memory by 99.9% on ShapeNet and MM-Fi datasets.

Matches dense APML accuracy within a small tolerance.

Enables large-batch 3D learning with efficient memory usage.

Abstract

Loss functions are fundamental to learning accurate 3D point cloud models, yet common choices trade geometric fidelity for computational cost. Chamfer Distance is efficient but permits many-to-one correspondences, while Earth Mover Distance better reflects one-to-one transport at high computational cost. APML approximates transport with differentiable Sinkhorn iterations and an analytically derived temperature, but its dense formulation scales quadratically in memory. We present CUDA-APML, a sparse GPU implementation that thresholds negligible assignments and runs adaptive softmax, bidirectional symmetrization, and Sinkhorn normalization directly in COO form. This yields near-linear memory scaling and preserves gradients on the stored support, while pairwise distance evaluation remains quadratic in the current implementation. On ShapeNet and MM-Fi, CUDA-APML matches dense APML within a small tolerance while reducing peak GPU memory by 99.9%. Code available at: https://github.com/Multimodal-Sensing-Lab/apml