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

TL;DR

This work tackles the challenge of generating complete, class-consistent 3D scenes as point clouds without relying on external object databases. It introduces a Class-Partitioned VQ-VAE (CPVQ-VAE) with a class-aware codebook and a class-aware running average update to mitigate codebook collapse, enabling reliable decoding of latent features into class-specific shapes. Complementing this, a Latent Space Flow Matching Model (LFMM) guides the generation of object bounding boxes, classes, and latent features, with a class-aware inverse look-up mapping latents to suitable codebook entries for decoding. Across 3D-FRONT living, dining, and bedroom scenes, the combined CPVQ-VAE+LFMM framework achieves substantial reductions in Chamfer Distance and Point2Mesh errors (up to 70.4% and 72.3%, respectively) and delivers faster runtime than prior diffusion-based scene methods, demonstrating a retrieval-free, scalable path for multi-object 3D scene synthesis.

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.

Figures (4)

• Figure 1: Top: Diffuscene, together with a pretrained VAE, generates object latents resulting in poorly decoded point clouds which frequently disagree with the object's target class. Our Latent Flow Matching Model (LFMM) produces more robust latents which are correctly decoded by the Class-Partitioned VQ-VAE (CPVQ-VAE), that exploits class labels from the LFMM to look-up the corresponding codevectors. Bottom: Overview of our proposed point cloud scene generation framework. The LFMM generates object bounding box parameters, classes and respective features for the entire scene. Thereafter, the CPVQ-VAE decodes object features, given the generated object class and feature.
• Figure 2: The Class-Partitioned VQ-VAE uses class label inputs to partition the codebook, allowing it to learn class-specific point cloud representations.
• Figure 3: U-Net architecture of the Latent space Flow Matching Model.
• Figure 4: Visualization of retrieved and generated scenes. Green close-ups indicate correctly retrieved/decoded objects while incorrect results are in red. Notably, unlike the VAE decoder used by Diffuscene, our CPVQ-VAE decoder accurately decodes object features into class-consistent point clouds.