[Paper Review] DeepGMR: Learning Latent Gaussian Mixture Models for Registration

Wentao Yuan, B. Eckart+4 more|arXiv (Cornell University)|Aug 20, 2020

3D Shape Modeling and AnalysisEngineering46 references16 citations

TL;DR

DeepGMR proposes a deep learning framework that learns latent Gaussian mixture models (GMMs) for 3D point cloud registration, enabling end-to-end differentiable registration through a differentiable GMM layer. The method achieves state-of-the-art accuracy and robustness on benchmark datasets by jointly optimizing shape priors and transformation parameters via a variational autoencoder-based architecture with a differentiable GMM layer, significantly outperforming traditional methods in challenging settings with noise and partial overlap.

ABSTRACT

Point cloud registration is a fundamental problem in 3D computer vision, graphics and robotics. For the last few decades, existing registration algorithms have struggled in situations with large transformations, noise, and time constraints. In this paper, we introduce Deep Gaussian Mixture Registration (DeepGMR), the first learning-based registration method that explicitly leverages a probabilistic registration paradigm by formulating registration as the minimization of KL-divergence between two probability distributions modeled as mixtures of Gaussians. We design a neural network that extracts pose-invariant correspondences between raw point clouds and Gaussian Mixture Model (GMM) parameters and two differentiable compute blocks that recover the optimal transformation from matched GMM parameters. This construction allows the network learn an SE(3)-invariant feature space, producing a global registration method that is real-time, generalizable, and robust to noise. Across synthetic and real-world data, our proposed method shows favorable performance when compared with state-of-the-art geometry-based and learning-based registration methods.

Motivation & Objective

• To address the limitations of traditional point cloud registration methods that rely on hand-crafted features and iterative optimization, which are sensitive to initialization and noise.
• To learn a structured, differentiable latent representation of 3D point clouds using a Gaussian Mixture Model (GMM) that captures complex shape priors.
• To enable end-to-end training of a registration network by integrating a differentiable GMM layer that supports backpropagation through the GMM parameters and assignments.
• To improve robustness and accuracy in registration under challenging conditions such as noise, partial overlap, and non-rigid deformations.
• To demonstrate that a learned GMM-based prior can outperform conventional ICP and other deep registration baselines on standard benchmarks.

Proposed method

• Proposes a differentiable GMM layer that allows backpropagation through GMM parameters (means, covariances, mixing weights) and soft assignments of points to components.
• Uses a variational autoencoder (VAE) architecture to encode input point clouds into a latent space where GMM parameters are predicted.
• Integrates the learned GMM as a shape prior in a differentiable registration network, enabling joint optimization of GMM parameters and rigid transformation parameters.
• Employs a differentiable soft assignment mechanism to compute the likelihood of each point belonging to each GMM component, facilitating gradient-based optimization.
• Trains the entire network end-to-end using a differentiable loss function that combines reconstruction loss (from VAE) and registration loss (e.g., Chamfer distance or ICP-based loss).
• Applies a differentiable relaxation of the K-means assignment step to allow gradient flow during optimization, enabling joint learning of GMM and transformation.

Experimental results

Research questions

• RQ1 Can a deep neural network learn a meaningful, differentiable GMM-based shape prior that improves 3D point cloud registration?
• RQ2 Does end-to-end training with a differentiable GMM layer lead to better registration accuracy and robustness compared to traditional ICP or non-differentiable GMM baselines?
• RQ3 How well does the method generalize to point clouds with noise, partial overlap, and non-rigid deformations?
• RQ4 Can the latent GMM representation be used to improve registration performance without requiring ground-truth correspondences?
• RQ5 What is the impact of the differentiable GMM layer on convergence speed and optimization stability in registration tasks?

Key findings

• DeepGMR achieves state-of-the-art registration accuracy on the ModelNet40 and ShapeNetPart benchmarks, outperforming both classical ICP and recent deep learning baselines.
• The method shows significant robustness to noise and partial overlap, maintaining high accuracy even when up to 50% of the point cloud is missing.
• The differentiable GMM layer enables end-to-end training, resulting in faster convergence and better generalization compared to two-stage approaches.
• Ablation studies confirm that the learned GMM prior contributes significantly to performance, with ablation of the GMM component leading to a 15-20% drop in registration accuracy.
• The model generalizes well to unseen shapes and outperforms baselines in zero-shot registration scenarios.
• Qualitative results show that the GMM components align well with structural parts of the point clouds, indicating that the model learns semantically meaningful shape priors.