3D Shape Generation and Completion through Point-Voxel Diffusion

In plain words

A 3D shape is stored here as a cloud of points in space. The model learns to make such clouds by starting from random scatter and removing noise in many small steps, the same idea that image diffusion models use to turn static into a picture. During training it sees real shapes with known amounts of noise added, and learns to predict that noise.

The same network can also complete a shape. The points you already have stay fixed, and the model fills in the missing points while the rest of the cloud is denoised around them. Because every run starts from different random noise, you get several plausible completions of the same partial shape rather than one averaged answer.

The model combines points with a coarse voxel grid, so the network can see local structure as well as individual points. Pure voxels and pure point clouds were tried and did worse.

TL;DR

• Denoising diffusion (DDPM-style) on 3D point clouds, with a point-voxel CNN as the denoiser §3 §1.
• One trained network does unconditional generation and conditional completion; completion changes only the training objective and the sampling loop §3.2.
• Trained with a plain L2 loss on the predicted noise eq. 6; sampling runs T = 1000 steps app. H.2.
• Author-reported: best 1-NN on all six generation cells (ShapeNet Airplane, Chair, Car; CD and EMD) tab. 1. For completion it wins on EMD but not CD tab. 2.
• Code is public under MIT and matches the paper in most places, but several settings differ (see Paper vs code).

Problem

• Voxel models are memory-hungry at high resolution, and voxels are binary, which does not suit a Gaussian diffusion process §1.
• Point-cloud generators usually map a latent vector, produced by a deterministic encoder, to a point set. For completion this gives one output, not a distribution over completions §1.
• Plain diffusion on pure voxels produced noisy shapes. Plain diffusion on pure points with a PointNet++-style permutation-invariant network produced no visible shapes §4.1.
• Completion is practically useful when only a single depth view is available §4.2.

Method

The shape is a set of N = 2048 points with xyz coordinates §3. The denoiser is a point-voxel CNN in the style of PVCNN, modified from PointNet++ by replacing the PointNet part with point-voxel convolution §3 app. H.1.

1. Forward process. Gaussian noise is added over T steps with a schedule β_1..β_T. Each step is q(x_t | x_{t-1}) = N(sqrt(1-β_t) x_{t-1}, β_t I) §3.1 eq. 2. With α_t = 1 - β_t and ᾱ_t = product of α_s, the closed form is x_t = sqrt(ᾱ_t) x_0 + sqrt(1-ᾱ_t) ε app. D.

2. Training objective. The variational lower bound on the likelihood is maximised §3.1 eq. 3. Because the forward process is fixed, this reduces to maximum likelihood over the reverse transitions §3.1 eq. 4. Each reverse step is matched to the Gaussian posterior q(x_{t-1} | x_t, x_0) §3.1 eq. 5 app. D. Reparameterised to predict the noise, the loss becomes an L2 distance between true and predicted noise §3.1 eq. 6 app. E.

3. Generation step. Points are sampled by running the reverse chain from t = T down to 1, with mean given by the predicted noise and fresh Gaussian noise added at each step, using σ_t² = β_t §3.1 eq. 7. The paper says this variance choice "works well" empirically §3.1. The derivation is in app. F.

4. Shape completion. The partial shape z_0 (M points) is fixed for all timesteps. Only the free points x̃_t are diffused §3.2 eq. 8. The reverse transition for the free points is conditioned on z_0 §3.2 eq. 8. The training objective is the same bound, conditioned on z_0 §3.2 eq. 9. The loss is the L2 noise error on the free points only, with the partial points masked out of the error §3.2 eq. 10.

5. Input at every step. The partial and free points are concatenated and passed through the network. The output is kept only for the free points, and the partial points are reinserted before the next step §3.2 app. H.1.

6. Time conditioning. A sinusoidal timestep embedding of dimension 64 is passed through a small MLP and concatenated with point features at the start of each set-abstraction and feature-propagation module app. H.1.

7. Same network for both tasks. Generation and completion use the same architecture. Only the objective (and, for completion, the 200 fixed points) differ §3.2 app. H.1.

Block-level architecture, from app. H.1:

The output head is an MLP(64, 3) app. H.1. Every PVConv block uses 3x3x3 convolutions, GroupNorm(8), Swish and Dropout(0.1) app. H.1.

Pseudocode for one training step (generation, eq. 6):

# Training step for shape generation  [§3.1] [eq. 6]
# Input:  points (B, N, 3) with N = 2048; T = 1000
# Output: scalar loss

def training_step(points, alpha_bar):
    # Draw one timestep per shape, uniformly over the schedule
    # [code: train_generation.py#L434]
    steps = uniform_int(0, T, size=B)
    # Sample Gaussian noise with the same shape as the points  [§3.1]
    noise = gaussian(shape=points.shape)
    # Diffuse the points in closed form to timestep steps  [app. D]
    noisy = (
        sqrt(alpha_bar[steps]) * points
        + sqrt(1 - alpha_bar[steps]) * noise
    )
    # Predict the noise with the point-voxel network  [app. H.1]
    predicted = model(noisy, steps)
    # Mean squared error between true and predicted noise  [eq. 6]
    return mean((noise - predicted) ** 2)

Pseudocode for generation (eq. 7, variance as stated in the paper):

# Sample one shape by ancestral sampling  [§3.1] [eq. 7] [app. F]
# Input:  alpha, beta, alpha_bar of length T; variance[t] (see Pitfalls)
# Output: points (N, 3)

def sample_shape(model):
    # Start from a standard Gaussian over N points  [§3.1]
    points = gaussian(shape=(N, 3))
    # Walk the reverse chain from the last step to the first  [§3.1]
    for step in reversed(range(T)):
        # Predict the noise at this step  [eq. 7]
        predicted = model(points, step)
        # Mean of the reverse transition  [eq. 7]
        mean = (
            points
            - beta[step] / sqrt(1 - alpha_bar[step]) * predicted
        ) / sqrt(alpha[step])
        # Add noise on every step except the last  [eq. 7]
        if step > 0:
            noise = gaussian(shape=(N, 3))
            points = mean + sqrt(variance[step]) * noise
        else:
            points = mean
    return points

Pseudocode for completion (fixed partial shape, eq. 10):

# One completion training step  [§3.2] [eq. 10]
# Input:  partial (M, 3) with M = 200; free (1848, 3); T = 1000
# Output: scalar loss on the free points

def completion_training_step(partial, free):
    # Draw a timestep for this sample  [code: train_completion.py#L430]
    step = uniform_int(0, T)
    # Sample noise for the free points only  [§3.2]
    noise = gaussian(shape=free.shape)
    # Diffuse the free points; the partial points stay clean  [§3.2]
    noisy_free = (
        sqrt(alpha_bar[step]) * free
        + sqrt(1 - alpha_bar[step]) * noise
    )
    # Predict noise for all points, keep only the free part  [§3.2]
    predicted = model(concat(partial, noisy_free), step)[free_part]
    # L2 loss on the free points  [eq. 10]
    return mean((noise - predicted) ** 2)

Hyperparameters

Values in the paper are cited to the paper. Where the paper says nothing, the code's value is given and labelled.

Results

Author-reported. All results use the authors' own evaluation implementation §4.1. Generation uses ShapeNet Airplane, Chair and Car with 2048 points per shape, compared with r-GAN, l-GAN (CD and EMD), PointFlow, SoftFlow, DPF-Net, Shape-GF and a voxel-diffusion baseline (Vox-Diff) §4.1. Lower is better for 1-NN.

Generation, 1-NN (%), lower is better tab. 1

Source: tab. 1.

Voxel baseline, 1-NN (%) app. A

Full coverage (COV, %, higher is better) and MMD (lower is better) metrics, with CD and EMD distances app. A:

Bold marks the best value per column in the source app. A. PVD is bold on Airplane COV (CD and EMD), Chair COV (CD and EMD), Chair MMD EMD and Car COV EMD. It ties PointFlow on Airplane MMD CD at 0.2243. Baselines win the other columns: PointFlow on Chair MMD CD and Car MMD CD, SoftFlow on Airplane MMD EMD, Shape-GF on Car MMD EMD and Car COV CD.

Shape completion, from single-view depth on GenRe renderings (CD ×10³, EMD ×10²; lower is better) tab. 2

Setup: 200 partial points per view, 20 views per shape, all 20 partial shapes evaluated per ground-truth sample §4.2. Shape-GF was left out because its completions were not realistic after Langevin sampling §4.2.

Multi-modal completion on PartNet (TMD ×10², higher is better; MMD ×10³, lower is better; CD distance) tab. 3

Setup: 1024 given points, cGAN and KNN-latent protocol §4.3.

Paper vs code

Code: MIT License, Linqi Zhou, 2022 at code_commit 9747265a. The repository is linked from the project page the paper cites. Code paths below are relative to that repository.

• Generation variance. The paper sets σ_t² = β_t §3.1. The training code's default variance is fixedsmall, the posterior variance β̃_t, and the sampler adds noise with that variance train_generation.py train_generation.py. The test script's sampler has the same switch, use_var test_generation.py.
• Airplane schedule. The paper interpolates β from 1e-5 to 0.008 over the first 90% of steps, then holds it app. H.2. The code's warm0.1 schedule ramps over the first int(0.1 * T) = 100 steps (10%), then holds train_generation.py train_generation.py. Only the airplane setting is overridden for generation in main.
• Chair and Car β_T. The paper sets β_T = 0.01 for Car and Chair app. H.2. The default in both scripts is β_T = 0.02, and no category override sets 0.01 train_generation.py train_completion.py. Running the scripts without flags therefore uses 0.02.
• Completion schedule for Airplane. train_completion.py has no Airplane override, so the completion script uses the linear 0.02 schedule for every category unless flags are passed train_completion.py. The paper reports completion on Airplane, Chair and Car tab. 2.
• Feature-propagation attention never turns on. The paper's architecture table has attention in FP 2 app. H.1. In both models the FP attention flag requires c < len(fp_blocks) - 1, but fp_blocks holds only the PointNet FP module when the flag is computed (p = 0), so the condition is always false pvcnn_generation.py pvcnn_generation.py. The completion model has the same guard with c % 2 == 0 pvcnn_completion.py. Net effect: no FP block uses attention in the released code, although the paper's table lists one. The paper says the same architecture is used for both tasks §3.2 app. H.1.
• Global attention. Both models add a global attention layer after SA 4 when attention is enabled pvcnn_generation.py pvcnn_generation.py. The paper's architecture table has no such layer app. H.1.
• SA 4 widths. The code's SA 4 MLP output is (256, 256, 512) train_generation.py. The paper's SA 4 column lists no output size app. H.1.
• Time embedding frequencies. The code uses frequencies 10000^(-k/(d/2 - 1)) for k = 0..d/2-1 pvcnn_generation.py. The paper's formula is ω_k = 1/10000^(2k/d) app. H.1. With d = 64 the two exponents differ.
• Clipping. When clip_denoised is set, the predicted x_0 is clamped to [-0.5, 0.5] train_generation.py. Evaluation and the test script disable clipping train_generation.py test_generation.py. The paper does not mention clipping, and the ±0.5 range is not stated.
• Optimiser details. The code uses Adam with β1 = 0.5 and an exponential learning-rate decay of 0.998 per epoch train_generation.py train_generation.py train_generation.py train_generation.py. The paper states only the learning rate app. H.3.
• Test-time de-normalisation. The test script multiplies samples by a per-shape standard deviation and adds a per-shape mean from the dataset before evaluation test_generation.py. The paper does not describe this step.
• Noise for t = 0. The training loop pre-draws noise for every example and redraws it only for t ≠ 0 train_generation.py train_generation.py. For t = 0 the same fixed noise is reused across epochs. The paper does not describe this.
• Vox-Diff. The voxel-diffusion baseline is not in the repository. Only PVD is provided.

The following agree between paper and code: the L2 loss on predicted noise eq. 6, T = 1000, the Adam learning rate of 2e-4, the 200 partial points for completion, and the 2048-point shape size.

Pitfalls

• Variance choice. The paper's σ_t² = β_t and the code's posterior variance β̃_t are both in the repository. The paper does not say which produced the reported numbers. A reimplementation should pick one and state it; the difference is in the sampler's noise scale, and it is not quantified in the paper.
• Airplane schedule. The paper's 90% ramp and the code's 10% ramp give different noise levels for most of the first 90% of steps. An implementer who follows the paper will not reproduce the released airplane schedule.
• Chair and Car β_T. Running the public scripts without flags uses β_T = 0.02, not the paper's 0.01.
• Different metrics per task. Generation is scored with 1-NN §4.1. Completion is scored with CD and EMD distances only §4.2. Numbers from the two tables cannot be compared.
• CD versus EMD in completion. PVD is worse than the baselines on CD in all three categories tab. 2. The authors argue EMD is the better measure because it enforces matching density, and they note CD is blind to visual inferiority §4.2. A project that optimises CD should expect the baselines to win.
• Mean-shape failure. When the input view leaves much of the object unseen, baselines output a mean shape, which scores well on CD. PVD outputs a plausible but different shape instead, which scores worse on CD §4.2 fig. 5.
• PartNet point counts. The text says 1024 points are given and 2048 are generated §4.3. It also says TMD is computed on the 1024 free points only §4.3. The two statements are compatible only if the output is 1024 given plus 1024 new points; the implementation should check this against the released data split.
• Diversity metric scope. TMD for PVD is computed on free points only, while baselines are scored on a subsampled 1024-point set §4.3. The comparison is not like-for-like.
• Multi-modal completion data. Real-scan tables come from random views of ShapeNet meshes, not from real scans, because GenRe has no table data §4.3.
• Sampling cost. Sampling runs T = 1000 network evaluations. The authors say 100 steps is not enough to build shapes app. H.2. No wall-clock time, GPU count or training duration is stated.
• Naive point or voxel diffusion. The authors report that pure-point diffusion with PointNet++ fails to generate visible shapes and pure-voxel diffusion is noisy §4.1. The repository does not contain these baselines, so the claim cannot be re-run from the code.
• Controlled completion. Latent interpolation at time T with weight λ, between two partial encodings, is described as a qualitative feature app. G. No numbers are given.
• Undisclosed training budget. Batch size, epochs, seed and training time are not in the paper. The repository defaults are the only values available train_generation.py.

Lineage

• Builds on: the DDPM denoising formulation (Ho et al., 2020, as cited in §1 §3.1), and the point-voxel CNN of Liu et al. (2019) §3 §1.
• Compared against: r-GAN, l-GAN, PointFlow, SoftFlow, DPF-Net and Shape-GF for generation and completion §4.1 §4.2; cGAN and KNN-latent for PartNet multi-modal completion §4.3; a voxel-diffusion baseline (Vox-Diff), which is the authors' own §4.1.
• Superseded by: not stated in this card's source.