APML: Adaptive Probabilistic Matching Loss for Robust 3D Point Cloud Reconstruction

Sasan Sharifipour, Constantino Alvarez Casado, Mohammad Sabokrou, Miguel Bordallo Lopez

Keywords: point cloud loss, deep learning, one-to-one matching, optimal transport, Sinkhorn algorithm, adaptive temperature, spatio-temporal point cloud, WiFi CSI, point cloud completion, point cloud generation, shape prediction, differentiable loss, low-density regions, probabilistic assignment

TL;DR: We propose APML, a differentiable and efficient loss for point cloud prediction tasks, approximating one-to-one matching using Sinkhorn iterations with adaptive temperature.

Abstract: Training deep learning models for point cloud prediction tasks such as shape completion and generation depends critically on loss functions that measure discrepancies between predicted and ground-truth point sets. Commonly used functions such as Chamfer Distance (CD), HyperCD, InfoCD and Density-aware CD rely on nearest-neighbor assignments, which often induce many-to-one correspondences, leading to point congestion in dense regions and poor coverage in sparse regions. These losses also involve non-differentiable operations due to index selection, which may affect gradient-based optimization. Earth Mover Distance (EMD) enforces one-to-one correspondences and captures structural similarity more effectively, but its cubic computational complexity limits its practical use. We propose the Adaptive Probabilistic Matching Loss (APML), a fully differentiable approximation of one-to-one matching that leverages Sinkhorn iterations on a temperature-scaled similarity matrix derived from pairwise distances. We analytically compute the temperature to guarantee a minimum assignment probability, eliminating manual tuning. APML achieves near-quadratic runtime, comparable to Chamfer-based losses, and avoids non-differentiable operations. When integrated into state-of-the-art architectures (PoinTr, PCN, FoldingNet) on ShapeNet benchmarks and on a spatio‑temporal Transformer (CSI2PC) that generates 3‑D human point clouds from WiFi‑CSI measurements, APM loss yields faster convergence, superior spatial distribution, especially in low-density regions, and improved or on-par quantitative performance without additional hyperparameter search. The code is available at: https://github.com/apm-loss/apml.

Primary Area: Deep learning (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 20417