Matching the Optimal Denoiser in Point Cloud Diffusion with (Improved) Rotational Alignment

2025-10-02

AI Summary

Diffusion models for molecular and protein point clouds—objects lacking intrinsic orientation—suffer from misalignment between rotational alignment procedures (e.g., Kabsch-Umeyama) and the underlying optimal denoiser, particularly under SO(3) symmetry. Method: We prove that Kabsch-Umeyama alignment corresponds to the zeroth-order approximation of the optimal denoiser under small-noise conditions, where the clean data follow a matrix Fisher distribution on SO(3); we further derive higher-order analytical approximations. Building upon this insight, we propose a denoising distribution learning framework grounded in SO(3) manifold statistics, integrating stochastic rotation augmentation and an improved alignment strategy. Contribution/Results: Experiments show that conventional alignment achieves near-theoretical optimality at critical noise levels—providing the first rigorous theoretical justification for existing practices. Moreover, our higher-order approximations significantly improve denoising accuracy in high-noise regimes, empirically validating the necessity of explicit manifold-aware modeling for diffusion-based point cloud generation.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Diffusion models are a popular class of generative models trained to reverse a noising process starting from a target data distribution. Training a diffusion model consists of learning how to denoise noisy samples at different noise levels. When training diffusion models for point clouds such as molecules and proteins, there is often no canonical orientation that can be assigned. To capture this symmetry, the true data samples are often augmented by transforming them with random rotations sampled uniformly over $SO(3)$. Then, the denoised predictions are often rotationally aligned via the Kabsch-Umeyama algorithm to the ground truth samples before computing the loss. However, the effect of this alignment step has not been well studied. Here, we show that the optimal denoiser can be expressed in terms of a matrix Fisher distribution over $SO(3)$. Alignment corresponds to sampling the mode of this distribution, and turns out to be the zeroth order approximation for small noise levels, explaining its effectiveness. We build on this perspective to derive better approximators to the optimal denoiser in the limit of small noise. Our experiments highlight that alignment is often a `good enough' approximation for the noise levels that matter most for training diffusion models.

Problem

• Analyzing rotational alignment effects in point cloud diffusion models
• Deriving improved approximators for optimal denoisers at small noise
• Studying SO(3) symmetry in molecular and protein point clouds

Innovation

• Uses matrix Fisher distribution over SO(3)
• Derives better approximators for small noise
• Improves rotational alignment in denoising