优化3D点云学习的损失函数，大幅降低显存占用
From Theory to Throughput: CUDA-Optimized APML for Large-Batch 3D Learning
- 用稀疏GPU实现改进的APML，直接在COO格式上运算
- 在ShapeNet和MM-Fi上保持精度，显存减少99.9%
- 适合大批次3D点云训练，尤其受限于显存的场景
损失函数对构建高精度3D点云模型至关重要，但常见方法在几何保真度与计算成本间权衡。Chamfer Distance虽高效却允许多对一匹配，Earth Mover Distance虽能反映一对一传输但计算开销大。APML通过可微分的Sinkhorn迭代和解析导出的温度近似运输过程，但其密集形式内存复杂度为二次方。本文提出CUDA-APML，一种稀疏的GPU实现，通过阈值剔除无关分配，并在COO格式上直接执行自适应softmax、双向对称化和Sinkhorn归一化。该方法实现近线性内存增长，同时保留存储支撑上的梯度，而配对距离计算仍为二次复杂度。在ShapeNet和MM-Fi数据集上，CUDA-APML与密集版APML误差极小，峰值显存降低99.9%。代码已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。