Class-Partitioned VQ-VAE and Latent Flow Matching for Point Cloud Scene Generation

TL;DR

This paper introduces a class-partitioned VQ-VAE and latent flow matching approach for generating complex 3D point cloud scenes directly, improving scene plausibility and accuracy without external object databases.

Contribution

The paper proposes a novel class-partitioned VQ-VAE with class-aware codebook updates and a latent flow matching model for direct point cloud scene generation.

Findings

Achieves up to 70.4% reduction in Chamfer error on complex scenes.

Effectively decodes class-specific point clouds without external databases.

Reliable scene recovery demonstrated on complex living room scenes.

Abstract

Most 3D scene generation methods are limited to only generating object bounding box parameters while newer diffusion methods also generate class labels and latent features. Using object size or latent feature, they then retrieve objects from a predefined database. For complex scenes of varied, multi-categorical objects, diffusion-based latents cannot be effectively decoded by current autoencoders into the correct point cloud objects which agree with target classes. We introduce a Class-Partitioned Vector Quantized Variational Autoencoder (CPVQ-VAE) that is trained to effectively decode object latent features, by employing a pioneering $\textit{class-partitioned codebook}$ where codevectors are labeled by class. To address the problem of $\textit{codebook collapse}$, we propose a $\textit{class-aware}$ running average update which reinitializes dead codevectors within each partition. During inference, object features and class labels, both generated by a Latent-space Flow Matching Model (LFMM) designed specifically for scene generation, are consumed by the CPVQ-VAE. The CPVQ-VAE's class-aware inverse look-up then maps generated latents to codebook entries that are decoded to class-specific point cloud shapes. Thereby, we achieve pure point cloud generation without relying on an external objects database for retrieval. Extensive experiments reveal that our method reliably recovers plausible point cloud scenes, with up to 70.4% and 72.3% reduction in Chamfer and Point2Mesh errors on complex living room scenes.