Fast rigid rotational alignment of voxelized 3-D density maps in sliced Wasserstein distance
Authors: Oscar Mickelin, Yunpeng Shi, Amit Singer, Eric J. Verbeke
Organizations: Yau Mathematical Sciences Center, Tsinghua University, Beijing 100084, China · Department of Mathematics, University of California at Davis, Davis CA, USA · Department of Mathematics, Princeton University, Princeton, NJ, USA · Program in Applied and Computational Mathematics, Princeton University, Princeton, NJ, USA · Center for Statistics and Machine Learning, Princeton University, Princeton, NJ, USA · Premium Research Institute for Human Metaverse Medicine, The University of Osaka, Japan
We develop a fast method for the rigid alignment of voxelized three-dimensional density functions over SO(3) using sliced optimal transport. Our approach efficiently transforms each density map into a tensor of one-dimensional quantile functions, taken from central slices over the unit sphere, such that the Euclidean distance of the transformed volumes approximates the sliced 2-Wasserstein distance. Our formulation is therefore compatible with existing, fast, voxel-based Euclidean alignment algorithms. We provide theoretical analysis and numerical experiments demonstrating that the sliced 2-Wasserstein distance is a robust metric for comparing three-dimensional shapes and for computing their rigid alignment.
Figures & tables
Figure 1: Visualization of the transforms used to compute the 3-D sliced Wasserstein distance.
Metric
32
64
128
256
512
Sliced 2-Wasserstein
0.023 ± 0.002
0.118 ± 0.05
0.627 ± 0.03
4.796 ± 0.015
39.879 ± 0.483
Ramp-Filtered Sliced 2-Wasserstein
0.028 ± 0.001
0.127 ± 0.001
0.828 ± 0.003
6.342 ± 0.028
52.249 ± 0.619
Table 1: Quantile tensor transform runtime. Computational runtime for transforming voxel data into the quantile tensor used to compute the the sliced 2-Wasserstein and ramp-filtered sliced 2-Wasserstein distance for increasing N×N×N volumes. Values are reported in seconds as the mean and standard deviation of n=5 trials.
Figure 2: Translation equivariance of the sliced 2-Wasserstein distance in 3-D. A reference 3-D volume (standard 3-D Gaussian, with σ=5 ), is shifted in voxel increments and the distance is computed in Euclidean, sliced 2-Wasserstein and ramp-filtered sliced 2-Wasserstein distance between the reference and shifted volume. Each metric is normalized by its maximum value for visualization.
Figure 3: Rotational alignment of the ModelNet10 data using different metrics. Models are discretized on a N×N×N=64×64×64 voxel grid and are aligned to a reference model in Euclidean, sliced 2-Wasserstein and ramp-filtered sliced 2-Wasserstein distance. (Top) Example alignment of the first 11 models from the chair category. The reference model is shown in purple and the aligned models are shown in gold. (Bottom) Summary of alignment results for 10 different categories. Each dataset contains n≈100 models, and the plots show the number of correctly aligned models as a function of increasing radians under the geodesic distance in eq. 24 .
Figure 4: Cryo-EM volumes from the EMDB [ 35 ] used for the homogeneous alignment experiments.
EMD-80134
EMD-2677
EMDB-8724
EMD-3342
EMD-1717
FSW3D
2.12±0.77
2.12±0.77
2.14±0.78
2.13±0.76
2.12±0.77
FSW3D + local
0.09±0.04
0.13±0.05
0.21±0.09
0.22±0.1
0.08±0.04
BOTalign
3.52±3.11
9.54±19.18
16.4±35.29
2.65±1.19
3.51±3.08
BOTalign + local
0.36±0.48
0.21±0.09
12.11±35.94
0.22±0.1
0.07±0.03
AlignOT
78.04±67.81
30.47±43.86
124.67±44.92
107.31±56.86
119.77±36.97
Table 2: Alignment accuracy results. The geodesic distance is computed between the estimated rotation and ground truth rotation for n=50 random rotations of 5 different cryo-EM volumes from the EMDB [ 35 ] , and is reported as the mean and standard deviation. The volumes are first downsampled to N×N×N=64×64×64 . Values are reported in degrees.
FSW3D
FSW3D + local
BOTalign
BOTalign + local
AlignOT
N=32
0.23±0.01
0.38±0.07
6.57±0.14
12.25±1.9
4.83±0.04
N=64
0.95±0.02
1.47±0.19
37.35±0.55
42.9±1.27
4.81±0.05
N=128
7.58±0.08
10.94±1.38
128.47±0.87
142.2±2.9
6.17±0.25
Table 3: Timing results for EMDB-1717 at increasing number of voxels, N={32,64,128} . Reported values are the mean and standard deviation of n=50 trials, and are shown in seconds.
Figure 5: Distance and alignment of heterogeneous cryo-EM data. (a) Cryo-EM volumes of size N×N×N=128×128×128 with simulated motion (see [ 42 ] for more details about the data). The distances are computed between the initial volume and the subsequent volumes with increasing state of motion. (b) Rotational alignment of the final volume to the initial volume under each metric. Orthogonal projection plots show the difference between the initial and final volume after rotational alignment, projected onto each orthogonal plane. Alignment in both sliced Wasserstein distances gives similar results which are distinct from the alignment in Euclidean distance.
Procrustes-Wasserstein alignment jointly estimates a matching and rotation without supplied correspondences, but alternating minimization can stop at suboptimal solutions. Rubix solves the equally weighted planar problem globally under squared Euclidean loss. Each matching σ of two centered n-point sets defines a complex correlation zσ=∑ixˉiyσ(i). Their convex hull is the permutation polygon: supporting vertices give optimal matchings at fixed rotations, and the farthest vertex gives the global alignment. We prove the sharp bound of n(n−1) vertices for n≥2, answering Rote's rotation-assignment open problem. In exact arithmetic, assignment queries recover the polygon in O(n5) operations. Assignment-based bounds extend the approach to three-dimensional rotations and partial matching at a supplied translation through branch-and-bound. On timed MPEG-7 shape pairs, Rubix attains every numerical reference value in 12 ms on average, 50 times faster than a rotation grid at the same accuracy. Its distances improve gravity-aligned matching of real 3D scans, shape retrieval and noisy crystal classification over alternating minimization.
Subhransu S. Bhattacharjee, Dylan Campbell, Rahul Shome
School of Computing, Australian National University
Dataset alignment is a central step in data analysis across science and engineering, where the goal is to match observations between datasets. Entropic Optimal Transport (EOT) offers a computationally tractable framework for this task by encoding cross-dataset affinities in a transport plan. However, when two datasets are sampled from geometrically similar low-dimensional structures with substantially different sampling densities, the EOT plan may match points by relative sampling density rather than geometric proximity, yielding geometrically misleading correspondences. To address this issue, we propose a density-reweighted EOT framework in which the influence of sampling density on the transport plan can be discounted to a desired degree, ranging from standard EOT to alignment driven purely by underlying geometry. Under suitable regularity conditions, we establish convergence of the reweighted EOT plan to a family of population-level plans whose dependence on sampling density is made explicit. Through simulations, we show that our approach recovers geometrically faithful correspondences, improving over related EOT-based frameworks when datasets exhibit substantial sampling density disparity.
Keyi Li, Yuval Kluger, Boris Landa
Program of Applied & Computational Mathematics, Yale University · Department of Pathology, Yale University · Department of Electrical and Computer Engineering, Yale University
We propose min Generalized Sliced Gromov--Wasserstein (min-GSGW), a sliced formulation for the Gromov--Wasserstein (GW) problem using expressive generalized slicers. The key idea is to learn coupled nonlinear slicers that assign compatible push-forward values to both input measures, so that monotone coupling in the projected domain lifts to a transport plan evaluated against the GW objective in the original spaces. The resulting plan induces a GW objective value, and min-GSGW minimizes this cost directly in the original spaces. We further show that min-GSGW is rigid-motion invariant, a crucial property for geometric matching and shape analysis tasks. Our contributions are threefold: 1) we introduce generalized slicers into the sliced GW framework, 2) we construct a slicing-based efficient GW transport plan; and 3) we develop an amortized variant that replaces per-instance optimization with a learned slicer for unseen input pairs. We perform experiments on animal mesh matching, horse mesh interpolation, and ShapeNet part transfer. Results show that min-GSGW produces meaningful geometric correspondences and GW objective values at substantially lower computational cost than existing GW solvers.
Ashkan Shahbazi, Xinran Liu, Ping He +1
Department of Computer Science, College of Connected Computing, Vanderbilt University · Department of Electrical and Computer Engineering, Vanderbilt University