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

Conference: AAAI2026
arXiv: 2601.12391
Code: ddsediri/CPVQ-VAE-LFMM
Area: 3D Vision
Keywords: point cloud generation, scene generation, VQ-VAE, flow matching, codebook partitioning

TL;DR¶

Proposes Class-Partitioned VQ-VAE (CPVQ-VAE) and Latent Flow Matching Model (LFMM) to realize the first pure point cloud scene generation method that does not require external database retrieval, reducing Chamfer Distance by 70.4% on complex living room scenes.

Background & Motivation¶

Existing 3D scene generation methods (e.g., Diffuscene) usually only generate the bounding box parameters and latent features of objects when generating complex, multi-category, and multi-object scenes, and then retrieve the nearest object meshes from a predefined database using L2 distance. This pipeline has two critical bottlenecks:

Unreliable generated latent codes: The object latents directly generated by diffusion models often conflict with the target category, causing the VAE-decoded point cloud shapes to severely mismatch the object classes (e.g., decoding a chair into a table shape).
Reliance on external databases: The retrieval paradigm inherently limits the diversity of generation, making it impossible to produce new object shapes that do not appear in the training set.

The authors observe that Diffuscene has an extremely high decoding failure rate in complex scenes (such as living rooms), while methods like ATISS completely abandon generating object latents and only use object sizes for retrieval. This indicates that scene generation requires a new architecture that can both guarantee category consistency and achieve reliable decoding.

Core Problem¶

How to directly generate point cloud objects containing the correct categories and shapes to compose complete 3D scenes without relying on an external object database? Specifically, the following challenges need to be addressed:

How to guarantee that the generated latents are consistent with the target categories.
How to alleviate the codebook collapse problem of VQ-VAE under the class-partitioned setting.
How to efficiently sample the complete scene layout from noise.

Method¶

Overall Architecture¶

The system is divided into two stages: (1) LFMM generates the bounding box parameters (translation, rotation, size), category vector, and a 32-dimensional latent feature for each object in the scene; (2) CPVQ-VAE maps the generated classes and features to codebook entries via category-aware inverse lookup, which is then decoded into a point cloud.

CPVQ-VAE: Class-Partitioned Vector Quantized VAE¶

The codebook of traditional VQ-VAEs is unlabeled, and the quantization process searches for the nearest neighbor across the entire codebook. The core improvement of CPVQ-VAE is partitioning the codebook by category:

The total number of codevectors is \(N_K = N_c \times N_q\), where \(N_c\) is the number of categories, and \(N_q\) is the number of codevectors allocated per category.
An indicator function \(\mathbf{1}(c,k)\) is introduced during quantization to only search for the nearest neighbor within the corresponding category partition.
This ensures that each quantized feature \(z^{q_c}\) belongs to the correct category.

The training loss consists of three terms: Chamfer Distance reconstruction loss (weight \(\lambda_{CD}=10\)), codebook alignment loss, and commitment loss, which are consistent with standard VQ-VAEs but operate on the class-partitioned codebook.

Category-Aware Running Average Update¶

To solve codebook collapse (where a large number of codevectors become "dead" entries during training), the authors propose a category-aware reinitialization strategy:

Usage Tracking: The usage rate of each codevector is tracked via exponential moving average \(U_s^k = \gamma U_{s-1}^k + \frac{1-\gamma}{B} u_s^k\) (\(\gamma=0.99\)).
Category-Aware Anchor Selection: The closest encoding in the same category within the mini-batch is selected as the anchor for each codevector.
Decay-Based Update: A decay coefficient \(\alpha_s^k = \exp(-\frac{10 U_s^k N_q}{1-\gamma} - \epsilon)\) is calculated, where underutilized codevectors receive higher reinitialization weights.

The key difference is that, unlike the category-agnostic approach of Zheng & Vedaldi (2023), both anchor selection and updates are conducted within the respective category partition.

LFMM: Latent Flow Matching Model¶

Flow matching based on optimal transport is employed to transfer Gaussian noise to clean scene layouts in the latent space:

The intermediate state represents linear interpolation: \(x_t = (1-t)x_0 + tx_1\)
The objective is to learn a constant velocity field \(v_\theta(x_t; t, \mathcal{F}_p) = x_1 - x_0\)
The loss is weighted separately for different attributes (translation, rotation, size, class, features).
The network architecture is a U-Net conditioned on the room floor plan \(\mathcal{F}_p\).
In inference, sampling is performed using the Euler method with \(N_{\hat{t}}=100\) steps.

Category-Aware Inverse Lookup¶

During inference, the generated 32-dimensional feature \(\hat{F}\) needs to be mapped back to a 128-dimensional codebook entry. By maximizing the cosine similarity between the truncated codevectors (first 32 dimensions) and \(\hat{F}\), while restricting the search inside the target category partition using the indicator function, the best codebook entry is found and then decoded into the final point cloud by the decoder.

Key Experimental Results¶

Evaluated on the 3D-FRONT dataset, which contains three scene types: Living Room (2338/587), Dining Room (2071/516), and Bedroom (5668/224).

Point Cloud Generation Quality (CD/P2M \(\times 10^3\), lower is better):

Method	Living Room CD	Living Room P2M	Dining Room CD	Dining Room P2M	Bedroom CD	Bedroom P2M
Diffuscene	30.63	29.87	30.60	29.49	45.01	44.88
LFMM + VAE	24.65	23.41	2.66	2.62	4.24	3.63
LFMM + CPVQ-VAE	9.06	8.27	2.38	2.17	2.46	2.06

Living Room: Compared to Diffuscene, CD is reduced by 70.4% and P2M by 72.3%.
Compared to the LFMM+VAE variant, CD is reduced by 63.2% and P2M by 64.7%.

Running Efficiency: Inference takes 0.892s, which is 90.3% faster than Diffuscene (9.153s).

Ablation Study (Bedroom):

Variant	VAE	VQ-VAE	Class Partitioning	Running Average Update	CD	P2M
V1	✓				4.24	3.63
V2		✓			36.27	33.93
V3		✓	✓		5.00	4.17
V4		✓	✓	✓	2.46	2.06

V2 (VQ-VAE without partitioning) suffers from severe codebook collapse and performs the worst. Adding partitioning (V3) significantly improves results but still lags behind VAE. Adding the running average update (V4) comprehensively outperforms all other variants.

Highlights & Insights¶

First Pure Point Cloud Scene Generation Method: Directly decodes and generates complete point cloud objects without requiring retrieval from an external object database.
Clever Class-Partitioned Codebook Design: Implements a concise partitioning mechanism via indicator functions, ensuring category consistency during decoding.
Category-Aware Codebook Maintenance: The running average update mechanism effectively mitigates codebook collapse while remaining category-aware.
Flow Matching Replaces Diffusion: Reduces sampling steps from thousands in diffusion to 100 steps, achieving an order of magnitude speedup.
Clear Ablation Study Design: The V1-V4 ablation chain completely demonstrates the contribution of each component.

Limitations & Future Work¶

Limited Point Cloud Density: The current autoencoder only processes sparse point clouds with 2025 points, which limits the quality of mesh reconstruction.
Quantization Error: Vector quantization inevitably introduces error, slightly reducing the diversity of retrieved objects.
Fixed Scene Types: The model needs to be trained separately for each scene type (living room, dining room, bedroom).
Lack of Comparison with More Baselines: DeBaRa and CASAGPT were not compared due to the absence of open-source code.
Latent Feature Truncation: Truncating the 128-dimensional codebook entries to 32 dimensions for training LFMM may cause information loss, the impact of which is not fully analyzed.

Method	Generation Method	Generates Point Cloud	Requires External Database	Requires Category Condition
ATISS	Autoregressive	No	Yes	No
Diffuscene	Diffusion	Yes (Poor quality)	Yes	No
DeBaRa	EDM Diffusion	No	Yes	Yes
SDF Methods	Diffusion	No (Voxel)	No	Yes (+ Scene graph)
Ours	Flow matching	Yes (Good quality)	No	Automatically Generated

The core advantage of this work lies in simultaneously generating both category labels and volumetric features. CPVQ-VAE leverages the generated category labels for conditional decoding, forming a closed loop.

Insights & Inspiration¶

Codebook Partitioning Concept is Transferable: The idea of partition-based codebooks can be applied to any scenario requiring multi-category discrete representations (e.g., image tokenizers, speech synthesis).
Potential of Flow Matching in 3D Generation: Compared to diffusion, the linear interpolation path in flow matching exhibits a better efficiency-quality tradeoff on 3D structured data.
Category Consistency is a Core Challenge in Scene Generation: This observation is equally applicable to other modalities (such as text-to-3D, interior design generation).

Rating¶

Novelty: 8/10 — The class-partitioned codebook and category-aware update mechanisms are original contributions, and applying flow matching to scene generation is also novel.
Experimental Thoroughness: 7/10 — Comprehensive evaluations on three scene types and clear ablations are provided, but comparisons with methods like DeBaRa are lacking.
Writing Quality: 8/10 — The method description is clear, mathematical derivations are complete, and the illustrations are intuitive.
Value: 7/10 — Advances the transition of point cloud scene generation from the retrieval paradigm to a pure generation paradigm, though restricted by point cloud density.