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

AI summary

Overview

Research area: 3D point cloud learning, optimal transport loss functions, and GPU kernel optimization (CUDA), with an application focus on point cloud completion and cross-modal 3D generation.

Technical level: Advanced. The paper assumes familiarity with optimal transport, entropy-regularized Sinkhorn normalization, permutation-invariant point set losses, and sparse GPU kernels (COO format, prefix sums, atomic accumulation).

Scope: A one-sentence scope: the paper presents CUDA-APML, a sparse GPU reimplementation of the Adaptive Probabilistic Matching Loss (APML) that keeps transport-plan construction in sparse COO form, matching dense APML accuracy on ShapeNet and MM-Fi while cutting peak GPU memory by 99.9%.

What This Paper Is About

Training models that output unordered 3D point sets requires a loss that compares predicted points to reference points without depending on ordering. Chamfer Distance (CD) is cheap but allows many-to-one matches, while Earth Mover Distance (EMD) better reflects one-to-one transport but is expensive; APML sits between them by building a differentiable transport plan with Sinkhorn iterations and an analytically derived temperature. APML's problem is memory: it stores dense N × M pairwise matrices, and prior work reported close to 300 GB of GPU memory when training FoldingNet on ShapeNet-55. This paper's goal is to keep APML's behavior while never materializing those dense tensors.

Key Contributions

1. A sparse COO formulation of APML. CUDA-APML thresholds negligible assignments after adaptive softmax and stores only the surviving index pairs in coordinate (COO) format, using the union support of the row-wise and column-wise directions rather than a dense N × M plan.

2. On-device sparse operators. The implementation constructs row-wise and column-wise adaptive-softmax assignments in two GPU passes per direction (a counting pass plus an exclusive prefix scan, then a write/normalize pass), symmetrizes by concatenating COO streams, sorting on a 64-bit key (for example key = i·M + j), and reducing duplicates by averaging, and then runs Sinkhorn column and row scaling directly on the sparse support using atomic-add accumulation kernels.

3. Gradients without dense intermediates. The loss is evaluated as a sum over nonzero entries only, and the backward pass differentiates through the same sparse graph, so pruned entries are treated as zero and contribute nothing to forward or backward computation.

4. An empirical scaling and efficiency study. The paper reports peak GPU memory across point counts on ShapeNet-55 with FoldingNet, synthetic measurements of nonzero COO entries (500 trials per point count, up to N = 262,144) and per-sample loss-side memory, and benchmark accuracy on ShapeNet-34 and MM-Fi against CD, DCD, HyperCD, InfoCD, and dense APML. Code is released at https://github.com/Multimodal-Sensing-Lab/apml.

Main Findings

• Accuracy is preserved. On ShapeNet-34 with FoldingNet, Dense-APML and CUDA-APML both reach an F1-score of 0.20 and EMD × 100 of 9.5, outperforming CD (F1 0.11, EMD × 100 27.7), HyperCD (F1 0.19, EMD × 100 16.5), and InfoCD (F1 0.18, EMD × 100 22.3).

• Cross-modal generation improves slightly over dense APML. On MM-Fi with CSI2PC, CUDA-APML reports EMD × 100 of 16.1 versus 16.3 for dense APML and 25.7 for DCD. F1-score is not reported for the MM-Fi rows in the table.

• Large memory reductions on ShapeNet-55. With FoldingNet at batch size 32, peak GPU memory falls from 50 MB to 0.39 MB at 1k points, from 1.6 GB to 2.6 MB at 8k points, from 18.5 GB to 8.6 MB at 32k points, and from 68 GB to 17.1 MB at 65k points. The abstract, introduction, and conclusion state the reduction as 99.9%, up to 99.9%, and more than 99.9% respectively.

• Near-linear memory scaling in the stored support. The plan is stored as 12·nnz bytes using 32-bit indices and 32-bit values, giving overall memory of O(B(nnz + N + M)) up to constant factors from sorting and scan workspaces, versus O(BNM) for the dense cost matrix and plan.

• The distance evaluation is still quadratic. The authors state explicitly that the adaptive-softmax stage still scans all pairs, so pairwise distance evaluation remains O(BNMd); speedups come from reduced memory traffic and from restricting Sinkhorn updates and loss evaluation to nnz entries.

• Qualitative coverage improves. In the shown MM-Fi sample, Chamfer-based losses produce more isolated points and gaps in the human silhouette, while APML and CUDA-APML produce more coherent coverage with fewer missing regions.

Methodology in Plain English

APML works by comparing each predicted point with each reference point, converting the distances into soft matching probabilities, and then iterating a normalization procedure (Sinkhorn) that pushes the plan toward one-to-one matching. The temperature controlling the sharpness of those probabilities is computed analytically from the gap between the smallest and second-smallest distances in a row or column, which sets a minimum probability mass p_min and removes the need to hand-tune a regularization strength.

The key observation this paper exploits is that after that sharpening, most entries of the similarity matrix are close to zero. CUDA-APML therefore applies a pruning threshold τ to the unnormalized similarities, keeps only pairs above it, and stores them as (row, column, value) triples. Because the number of kept entries per row and column is data-dependent, the GPU kernel runs in two passes: the first pass scans each row (or column), computes distances on the fly, tracks the minimum and second minimum to derive the temperature, counts entries passing the threshold, and converts those counts into write offsets with a prefix sum. The second pass rescans, recomputes the kept similarities, writes the triples, and normalizes them within each row or column. Row-wise and column-wise plans are then merged by concatenating them, sorting on a packed 64-bit key, and averaging duplicates.

Sinkhorn normalization then runs entirely on the sparse array: each iteration divides every stored value by the sum of values sharing its column (plus a stability constant), then by the sum of values sharing its row, implemented as an accumulation kernel using atomic adds into dense 1D buffers of length N or M followed by a scaling kernel. The final loss is just a weighted sum of distances over the stored pairs, and backpropagation follows the same sparse path. Default hyperparameters are τ = 10⁻⁸, L_iter = 10, and ε_stab = 10⁻⁸, with p_min set as in the original APML paper.

Evaluation uses ShapeNet-34 and ShapeNet-55 with standard partial-to-complete pairs and official splits, PCN (2,048-point partial input and 16,384-point reference), and MM-Fi, which pairs WiFi CSI with LiDAR-based 3D human point clouds under the CSI2PC manual_split protocol. Losses are integrated into FoldingNet, PCN, PoinTr, and CSI2PC. Metrics are CD in ℓ1 and ℓ2 variants, EMD × 100, F1-score at threshold 0.01, wall-clock time per epoch, and peak GPU memory; each configuration is trained three times with different seeds and the mean is reported. All experiments use PyTorch 2.5 with CUDA 12.8 on NVIDIA V100 (32 GB) and RTX 4070 GPUs, Adam at learning rate 10⁻⁴ and batch size 32.

Why This Matters

Impact on research. Transport-based supervision has been limited by memory rather than by theory: the paper notes that APML's dense formulation previously required close to 300 GB to train FoldingNet on ShapeNet-55, which constrains batch size, point count, and the number of feasible training runs. Showing that the same objective can be computed on a sparse support while preserving gradients on the stored entries makes one-to-one transport losses a practical option in higher-resolution regimes, and connects point cloud learning to the broader trend of removing dense intermediates shown by kernels such as FlashAttention and KeOps.

Real-world applications (as framed by the paper):

• Robotics, where 3D point cloud models support perception and manipulation.
• AR and VR, where accurate geometric reconstruction affects rendered scene fidelity.
• Wireless human sensing, where MM-Fi-style pipelines infer 3D human point clouds from WiFi CSI rather than from cameras.
• Digital twin and mapping pipelines, which involve large scenes and repeated updates, and structured-light, depth camera, and LiDAR scanning more generally.

Industry relevance. The paper frames the cost of quadratic memory as more than a technical detail: it constrains batch size and point count, reduces the number of feasible runs under fixed compute budgets, and raises energy and monetary cost during GPU training and tuning. A 99.9% reduction in peak loss-side memory can make transport supervision viable on hardware where it previously was not, and the released CUDA implementation lets practitioners adopt it without reimplementing the kernels.

Future Directions

• Remove the remaining quadratic term. The authors identify the adaptive-softmax stage's full pairwise scan as the main limitation and propose restricting the candidate support via neighborhood pruning or approximate nearest neighbors.

• Scale to larger settings. The stated next step is to evaluate the method on larger scene-level and digital twin update settings, where point counts and repeated updates are the binding constraint.

• Results beyond the reported configurations. The evaluation methodology lists PCN and PoinTr backbones and CD ℓ1/ℓ2 metrics, but the provided tables report F1-score and EMD × 100 only for FoldingNet on ShapeNet-34 and CSI2PC on MM-Fi; results for the other backbones, for CD ℓ1/ℓ2, and for per-epoch runtime are not reported in the content shown, leaving room for a fuller ablation.

• Interaction between pruning and accuracy. The threshold τ is fixed at 10⁻⁸ by default, and the paper does not report a sensitivity study over τ or over the number of Sinkhorn iterations, which would clarify how much pruning is safe before geometric fidelity degrades.

Target Audience

This paper is most useful to researchers and engineers building 3D point cloud generative or completion models who are memory-constrained during training, and to practitioners working on optimal-transport losses who want a drop-in sparse GPU implementation of APML. It is also relevant to GPU kernel developers interested in sparse reformulations of quadratic operators, and to applied teams in wireless sensing and digital twins who need cross-modal 3D pipelines that fit within fixed compute budgets. Readers without background in optimal transport or CUDA will need to consult the original APML paper and Sinkhorn literature first.

Authors’ 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