Adaptive Probabilistic Matching Loss (APML)

• APML is a differentiable loss that enables soft one-to-one matching between predictions and ground truth using temperature-controlled adaptive softmax.
• It leverages Sinkhorn normalization to produce an approximately doubly-stochastic transport plan for stable and efficient geometric matching.
• Empirical results demonstrate that APML achieves faster convergence and improved fidelity compared to traditional losses like Chamfer Distance and Earth Mover’s Distance.

Adaptive Probabilistic Matching Loss (APML) is a fully differentiable, analytically parameterized surrogate for one-to-one assignment between structured predictions and ground truth in geometric settings. Originating in robust 3D point cloud comparison and registration, APML leverages a temperature-controlled probabilistic assignment mechanism refined with Sinkhorn normalization. This produces an approximately doubly-stochastic transport plan that enables efficient, stable optimization and accurate geometric matching. APML has been validated across 3D completion, cross-modal reconstruction, robust registration, and generative modeling benchmarks, showing advantages in convergence, fidelity, and adaptivity over classical losses such as Chamfer Distance and exact Earth Mover’s Distance (Sharifipour et al., 9 Sep 2025, Sharifipour et al., 17 Dec 2025, Barron, 2017).

1. Mathematical Formulation

Let $\hat{X} \in \mathbb{R}^{B\times N\times d}$ be a batch of predicted point clouds and $X \in \mathbb{R}^{B\times M\times d}$ ground truth. For each batch $b$, construct the cost matrix

$$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$

Convert $D_b$ to a bidirectional similarity matrix using a temperature $T$,

$S_b = \exp(-D_b/T).$

The key innovation is adaptive per-row/per-column temperature scheduling: for softmax probabilities $P$, enforce that the minimum cost in each row/column is assigned at least $p_{\min}$ mass. For a row vector $c$,

$X \in \mathbb{R}^{B\times M\times d}$0

where $X \in \mathbb{R}^{B\times M\times d}$1 is the minimal cost gap plus a small $X \in \mathbb{R}^{B\times M\times d}$2, with $X \in \mathbb{R}^{B\times M\times d}$3 the number of entries in the row or column.

Bidirectional assignment proceeds via symmetrized row- and column-adaptive softmaxes, averaged to form $X \in \mathbb{R}^{B\times M\times d}$4. This is refined for $X \in \mathbb{R}^{B\times M\times d}$5 rounds using alternating normalization (the Sinkhorn–Knopp algorithm), producing approximately doubly-stochastic $X \in \mathbb{R}^{B\times M\times d}$6: $X \in \mathbb{R}^{B\times M\times d}$7 Repeat for $X \in \mathbb{R}^{B\times M\times d}$8 iterations. The loss per batch instance,

$X \in \mathbb{R}^{B\times M\times d}$9

is averaged over the batch.

2. Algorithmic Structure and CUDA Optimization

A typical APML pipeline consists of: (1) pairwise distance computation, (2) adaptive softmax in both directions, (3) initial symmetrization, (4) Sinkhorn normalization, and (5) loss computation. All steps are differentiable, supporting full backpropagation.

CUDA-APML escalates scalability, storing only non-negligible assignment matrix entries above a threshold $b$0. Adaptive softmax and Sinkhorn are implemented directly on sparse COO representations. Pairwise distance evaluation still requires $b$1, but memory utilization scales near-linearly in the number of effective assignments. On ShapeNet-55 with $b$2 points, CUDA-APML peaks at $b$3 MB RAM versus $b$4 GB for dense (Sharifipour et al., 17 Dec 2025).

CUDA-APML pseudocode processes each row/column independently for minima and softmax, concatenates and symmetrizes supports, and performs sparse Sinkhorn rounds. Gradients flow exclusively through the retained sparse support, retaining exact differentiability on the effective assignment graph.

3. Comparison to Alternative Losses

Loss	Geometry	Runtime	Differentiability	Coverage Issues
Chamfer	Many-to-one	$b$5	Non-smooth	Clumping, poor in sparse regions
EMD	One-to-one	$b$6	Non-smooth	High cost, covers sparsity well
APML	Soft one-to-one	$b$7	Fully differentiable	Improves coverage and spatial uniformity

APML achieves comparable runtime to Chamfer while eliminating nearest-neighbor index selection (thus, all steps remain smooth) and enforcing soft, transport-like correspondences that better capture global structure (Sharifipour et al., 9 Sep 2025, Sharifipour et al., 17 Dec 2025).

4. Integration and Empirical Effects in Deep Architectures

APML is a drop-in replacement for Chamfer loss in any model that requires point-wise matching, including PoinTr, FoldingNet, PCN, and the CSI2PC WiFi-to-3D transformer. No changes to core architecture or dataflow are needed beyond the loss layer. Only the $b$8 parameter, which determines temperature sharpness, is exposed; in practice, $b$9 works robustly.

Empirically, APML produces faster convergence (e.g., FoldingNet on ShapeNet-55 attains near-final F1 in $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$020 epochs, compared to 150 for CD). It yields lower Earth Mover’s Distance by 15–81% and superior point distribution, especially in sparse reconstruction zones, as observed in both ShapeNet and cross-modal MM-Fi settings (Sharifipour et al., 9 Sep 2025, Sharifipour et al., 17 Dec 2025).

5. Theoretical Foundations and Generalizations

APML belongs to the broader class of generator-matching losses with Bregman divergence structure. This extends to any differentiable assignment framework with time- or condition-dependent weighting schedules. The generator-matching framework admits APML-style matching as an instantiation, supported by theorems guaranteeing that, for any strictly positive reweighting, the minima of the loss coincide with the correct matching generator (Billera et al., 20 Nov 2025). The time- and state-dependent Bregman form is: $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$1 where $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$2 is the time sampling density, $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$3 is any positive schedule, $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$4 is the Bregman divergence, and $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$5 the learned assignment.

Any adaptive, learned, or error-driven $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$6 preserves the correctness of the optimal generator, enabling flexible stability/speed tradeoffs and targeting difficult instances during optimization (Billera et al., 20 Nov 2025).

6. Memory, Efficiency, and Practical Considerations

Dense APML is memory-bound due to the $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$7 cost of assignment and transport matrices. The sparsification introduced in CUDA-APML—pruning assignments below threshold $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$8—enables large-scale applications with nearly linear memory scaling and only negligible loss in accuracy (error bound of order $$D_b \in \mathbb{R}^{N\times M}, \quad (D_b)_{ij} = \| \hat{X}_{b,i} - X_{b,j} \|_2.$$9 EMD/$D_b$0 F1). The adaptive softmax remains the memory bottleneck; future improvements may use pre-filtering (k-NN search or locality-sensitive hashing) to restrict assignment support prior to softmax (Sharifipour et al., 17 Dec 2025).

APML consistently outperforms or matches existing losses in EMD and F1 metrics across 3D point cloud tasks, with a wall-clock overhead of $D_b$1–$D_b$2 and a peak RAM usage $D_b$3–$D_b$4 that of CD at high cardinalities, though $D_b$5 or more of entries are pruned in sparse implementation.

7. Extensions and Broader Implications

APML’s formulation generalizes beyond classic pairwise matching, allowing deployment in flow matching, diffusion models, and robust registration. Its differentiable, adaptive, and probabilistic construction offers a unified mechanism for loss design in problems where classical assignment or robust regression is needed. The theoretical foundation ensures that gradients, optimization trajectories, and minima are preserved under a broad class of adaptive schedule choices, including neural-learned weighting and error-adaptive updates. This positions APML as a versatile tool in robust geometric and probabilistic modeling frameworks (Sharifipour et al., 9 Sep 2025, Billera et al., 20 Nov 2025, Sharifipour et al., 17 Dec 2025, Barron, 2017).