Fast rigid rotational alignment of voxelized 3-D density maps in sliced Wasserstein distance

Oscar Mickelin; Yunpeng Shi; Amit Singer; Eric J. Verbeke

math.NAOct 29, 2026

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 ·

Abstract

We develop a fast method for the rigid alignment of voxelized three-dimensional density functions over $\mathsf{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 $\pm$ 0.002	0.118 $\pm$ 0.05	0.627 $\pm$ 0.03	4.796 $\pm$ 0.015	39.879 $\pm$ 0.483
Ramp-Filtered Sliced 2-Wasserstein	0.028 $\pm$ 0.001	0.127 $\pm$ 0.001	0.828 $\pm$ 0.003	6.342 $\pm$ 0.028	52.249 $\pm$ 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\times N\times 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 $\sigma=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\times N\times N=64\times 64\times 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\approx 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\pm 0.77$	$2.12\pm 0.77$	$2.14\pm 0.78$	$2.13\pm 0.76$	$2.12\pm 0.77$
FSW3D + local	$0.09\pm 0.04$	$0.13\pm 0.05$	$0.21\pm 0.09$	$0.22\pm 0.1$	$0.08\pm 0.04$
BOTalign	$3.52\pm 3.11$	$9.54\pm 19.18$	$16.4\pm 35.29$	$2.65\pm 1.19$	$3.51\pm 3.08$
BOTalign + local	$0.36\pm 0.48$	$0.21\pm 0.09$	$12.11\pm 35.94$	$0.22\pm 0.1$	$0.07\pm 0.03$
AlignOT	$78.04\pm 67.81$	$30.47\pm 43.86$	$124.67\pm 44.92$	$107.31\pm 56.86$	$119.77\pm 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\times N\times N=64\times 64\times 64

. Values are reported in degrees.

	FSW3D	FSW3D + local	BOTalign	BOTalign + local	AlignOT
$N=32$	$0.23\pm 0.01$	$0.38\pm 0.07$	$6.57\pm 0.14$	$12.25\pm 1.9$	$4.83\pm 0.04$
$N=64$	$0.95\pm 0.02$	$1.47\pm 0.19$	$37.35\pm 0.55$	$42.9\pm 1.27$	$4.81\pm 0.05$
$N=128$	$7.58\pm 0.08$	$10.94\pm 1.38$	$128.47\pm 0.87$	$142.2\pm 2.9$	$6.17\pm 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\times N\times N=128\times 128\times 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.

3D Shape Matching

See all 18 papers

Rubix: Global Correspondence-Free Point Set Alignment through Assignment Geometry
Oct 7, 2026Subhransu S. Bhattacharjee, Dylan Campbell, Rahul ShomePoint Cloud Registration Combinatorial Optimization
School of Computing, Australian National University
Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density
Aug 17, 2026Keyi Li, Yuval Kluger, Boris LandaOptimal Transport Entropy Regularization
Program of Applied & Computational Mathematics, Yale University · Department of Pathology, Yale University · Department of Electrical and Computer Engineering, Yale University
Min Generalized Sliced Gromov Wasserstein: A Scalable Path to Gromov Wasserstein
May 13, 2026Ashkan Shahbazi, Xinran Liu, Ping He +1Non-Rigid Shape Correspondence Optimal Transport
Department of Computer Science, College of Connected Computing, Vanderbilt University · Department of Electrical and Computer Engineering, Vanderbilt University

Oct 7, 2026J/K move · Enter open · S save

Rubix: Global Correspondence-Free Point Set Alignment through Assignment Geometry

Subhransu S. Bhattacharjee, Dylan Campbell, Rahul Shome

School of Computing, Australian National University

Point Cloud Registration Combinatorial Optimization

Open paper AskPDF arXiv

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_σ=\sum_i\bar x_i y_{σ(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\ge2

, answering Rote's rotation-assignment open problem. In exact arithmetic, assignment queries recover the polygon in

\mathcal O(n^5)

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.

Abstract

3D Shape Matching

Explore similar work

Rubix: Global Correspondence-Free Point Set Alignment through Assignment Geometry

Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density

Min Generalized Sliced Gromov Wasserstein: A Scalable Path to Gromov Wasserstein