SDFusion: Multimodal 3D Shape Completion, Reconstruction, and Generation

In plain words

Making a 3D object, such as a chair or a building, is hard for people who are not 3D artists. This paper builds a generator that accepts what a person can easily give it: a shape with parts missing, a photo, a sentence, or several of these at once. The user can also say how much each input should count.

The model works in two stages. First, an autoencoder squeezes each 3D shape, stored as a grid of distances to its surface, into a much smaller code. Second, a diffusion model learns to produce such codes. It starts from random noise and removes it step by step, and the decoded result is a shape. The inputs are turned into features and fed to the denoiser through attention. During training some inputs are sometimes left out, so the model can later run with any subset.

At generation time each input gets a weight, and the output follows the inputs with larger weights more closely. The authors also show that a 2D text-to-image model can supply the surface appearance of a generated shape, which the shape model itself does not produce.

TL;DR

• A latent diffusion model for 3D shapes: a 3D VQ-VAE compresses truncated signed-distance grids, and a diffusion model samples in that latent space §3.1 §3.2.
• Partial shapes, images and text condition the model through task-specific encoders, cross-attention and classifier-free guidance, with one weight per input §3.3 eq. 4.
• Author-reported: lower fidelity error (UHD) and higher diversity (TMD) than MPC and AutoSDF on shape completion tab. 1; best Chamfer distance and F-score on Pix3D tab. 2; the ShapeGlot evaluator prefers it over AutoSDF in 49% of comparisons against 36% tab. 3.
• Read "Paper vs code" before reimplementing: the released code differs from the paper in how multi-input weights, the image encoder and dropout are applied.

Problem

Conventional 3D generative models need direct 3D supervision such as point clouds, signed distance functions or voxels, and most of them work at low resolution, for example 32³ voxel grids and 2K-point clouds §1. Running a diffusion model directly on a high-resolution 3D grid is computationally very demanding §3.

The target use is interactive generation by people without 3D skills, who can supply partial geometry, an image, or a text description, and who should be able to weight these inputs §1. The method is shape-only; texture is added in a separate step §3.4 §5. The paper evaluates object-centric data and leaves whole scenes to future work §5.

Method

The model has two stages §3 fig. 2. A 3D VQ-VAE maps a shape to a small latent grid and back §3.1. A diffusion model is trained on these latents, and conditions enter through task-specific encoders and cross-attention §3.2 §3.3. A separate texturing step uses a pretrained 2D model §3.4 fig. 3.

1. Shape representation. A shape is a volumetric truncated signed distance field (T-SDF) X of size D × D × D §3. The paper does not give the truncation value; the code uses 0.2 (see Hyperparameters and Pitfalls).
2. Compression. A 3D VQ-VAE encodes X to a latent z with encoder E_φ, replaces z with its nearest entry in a codebook Z, and decodes with D_τ §3.1 eq. 1. Encoder, decoder and codebook are optimised jointly. The pre-training loss has a reconstruction term, a commitment term and the VQ objective, following van den Oord et al. §3.1. Training data are ShapeNet or BuildingNet §3.1. The latent dimension d is smaller than D; its value is not stated in the main text §3.1.
3. Latent diffusion. Latents are noised with a variance schedule over steps t = 1 … T. A time-conditioned 3D UNet ε_θ predicts the noise. Training uses the simplified noise-prediction objective of Ho et al. §3.2 eq. 2. Sampling starts from N(0, 1), denoises, and decodes with D_τ §3.2.
4. Conditions. Each modality i has its own encoder E_φi. The encoded features pass through dropout D, are combined by concatenation F, and enter the UNet through cross-attention §3.3 eq. 3. Partial shapes use blended diffusion (Avrahami et al.); single images use CLIP; text uses BERT §3.3.
5. Classifier-free guidance with weights. At inference the noise prediction is the unconditional prediction plus, for each modality, a weighted difference. The difference is between the prediction with only that modality present and the unconditional prediction, scaled by s_i §3.3 eq. 4. Larger s_i makes modality i more influential §3.3.
6. Texturing. The generated T-SDF is turned into a fixed density with VolSDF. A colour field F_θ(x, d) → c is trained on rendered views. The gradient comes from score distillation sampling (SDS) against a pretrained 2D diffusion model conditioned on a sentence S §3.4 eq. 5 fig. 3.

\mathbf{z} = E_{\varphi}(\mathbf{X}), \quad \text{and} \quad \mathbf{X}^{\prime} = D_{\tau}(\text{VQ}(\mathbf{z})) eq. 1

L_{\mathrm{simple}}(\theta) \coloneqq \mathbb{E}_{\mathbf{z}, \epsilon \sim N(0, 1), t} \left[ \lVert \epsilon - \epsilon_\theta(\mathbf{z}_t, t) \rVert^2 \right] eq. 2

L(\theta, \{\phi_i\}) \coloneqq \mathbb{E}_{\mathbf{z}, \mathbf{c}, \epsilon, t} \left[ \lVert \epsilon - \epsilon_\theta(\mathbf{z}_t, t, F\{D \circ E_{\phi_i}(\mathbf{c}_i)\}) \rVert^2 \right] eq. 3

\epsilon_\theta(\mathbf{z}_t, t, F\{E_{\phi_i} \, \forall i\}) = \epsilon_\theta(\mathbf{z}_t, t, \boldsymbol{\emptyset}) + \sum_i s_i \left( \epsilon_\theta\big(\mathbf{z}_t, t, F\{E_{\phi_i}(\mathbf{c}_i), E_{\phi_j}(\mathbf{c}_j) : \mathbf{c}_j = \boldsymbol{\emptyset} \; \forall j \neq i\}\big) - \epsilon_\theta(\mathbf{z}_t, t, \boldsymbol{\emptyset}) \right) eq. 4

\mathbb{E}_{t, \epsilon} \left[ \frac{\partial I(F_\theta(\cdot))}{\partial \theta} \, w(t) \, (\tilde{\epsilon}_\phi(\mathbf{z}_t, t, S) - \epsilon) \right] eq. 5

Symbols as printed in the paper: X is the T-SDF, z its latent, E and D the encoder and decoder, VQ the quantiser, ε_θ the denoiser, c_i the i-th condition, F the feature aggregation (concatenation), D (inside eq. 3) the dropout operator, ∅ a zero-filled condition, s_i the weight of modality i, and I the rendered image §3.1 §3.3 §3.4.

The training step below follows eqs. 1–3 and the code's per-modality dropout (see Paper vs code). The sampling step follows eq. 4.

# Training step, conditional latent diffusion  [§3.2] [§3.3] [eq. 2, 3]
# Input:  shape X as a T-SDF (D, D, D); one condition per modality
# Output: scalar loss

def training_step(shape, conditions, encoders, vq_vae, denoiser):
    # Compress the shape with the pretrained VQ-VAE  [§3.1] [eq. 1]
    z = vq_vae.encode(shape)                         # (d, d, d)
    # Sample a timestep and Gaussian noise  [§3.2]
    t = sample_timestep()
    noise = sample_normal_like(z)
    z_noisy = add_noise(z, noise, t)
    # Encode each condition; a dropped one is zero-filled  [§3.3]
    # Dropout rate is not stated in the paper
    features = []
    for c_i, encoder_i in zip(conditions, encoders):
        features.append(drop(encoder_i(c_i)))
    context = concatenate(features)                  # F  [§3.3]
    # Predict the noise, with the context through cross-attention  [§3.2]
    predicted = denoiser(z_noisy, t, context)
    # Simplified objective: squared error on the noise  [eq. 2]
    return mean_squared_error(predicted, noise)

# Sampling with per-modality guided weights  [§3.3] [eq. 4]
# Input:  conditions c_i with weights s_i; trained denoiser and decoder
# Output: decoded T-SDF shape

def sample_shape(denoiser, vq_vae, conditions, weights, steps):
    # Start from standard normal noise in the latent space  [§3.2]
    z = sample_normal((d, d, d))
    for t in reversed_timesteps(steps):
        # Unconditional prediction: every modality zero-filled  [eq. 4]
        uncond = denoiser(z, t, empty_context())
        guided = uncond
        for i, (c_i, s_i) in enumerate(zip(conditions, weights)):
            # Prediction with only modality i present  [eq. 4]
            only_i = denoiser(z, t, context_with_only(i, c_i))
            # Add the weighted difference  [eq. 4]
            guided = guided + s_i * (only_i - uncond)
        # Sampler update: not named in the paper; the code uses DDIM
        z = sampler_step(z, guided, t)
    return vq_vae.decode(z)                          # [§3.2]

Hyperparameters

Values marked "not stated" are absent from the paper's main text. The code column is what the released training and inference scripts use; it is a starting point, not a paper result.

Results

Author-reported throughout; the paper gives no error bars and no repeated runs.

Shape completion §4.1: ShapeNet at 64³ and BuildingNet at 128³, given the bottom half of the shape. UHD (fidelity, lower is better) measures the unidirectional Hausdorff distance from partial shapes to generated shapes. TMD (diversity, higher is better) is the average Chamfer distance among k = 10 samples. Baselines: MPC and AutoSDF.

Single-view reconstruction §4.2: Pix3D. The chair category uses the dataset's split; other categories use random splits made by the authors, so they are not comparable with other papers. Metrics are Chamfer distance (lower is better) and F-score at 1% (higher is better), with all methods aligned to 32³ voxels. The Chamfer distance unit is not stated.

Text-guided generation §4.3: ShapeGlot. A neural evaluator, trained by the authors, picks the target over a distractor given a description; a comparison counts as confused when the two confidences differ by at most 0.2.

Qualitative only: shape completion fig. 4; single-view fig. 5; text fig. 6; joint partial-shape and text, or partial shape and image fig. 7; weight control on the three modalities fig. 8; texturing with text prompts fig. 9. The paper reports no quantitative texture or multi-condition metric §4.4 §4.5. The introduction says the model uses less memory than AutoSDF's autoregressive model §1; no memory measurement appears in the main text.

Paper vs code

Official code: the repository linked from the project page, at the commit above, MIT licence. The paper's equations and tables were checked against it.

• Multi-input weights. Paper: eq. 4 applies a separate weight s_i to each modality §3.3. Code: the default multi-modal inference (mm_cls_free=False) scales the condition features by img_scale and txt_scale before concatenation and runs ordinary classifier-free guidance [code: models/sdfusion_mm_model.py#L464-L465] sdfusion_mm_model.py. Eq. 4 is implemented only in the mm_cls_free=True branch [code: models/sdfusion_mm_model.py#L527-L548] [code: models/networks/diffusion_networks/samplers/ddim.py#L189-L204]. There, the global guidance scale multiplies every w_i ddim.py, so the effective weight is scale × w_i.
• Image encoder for multi-modal input. Paper: CLIP is named for single-view reconstruction only §3.3; the multi-modal encoder is not named. Code: the multi-modal model uses a ResNet-18 with 512-d features, projected to the text width sdfusion_mm_model.py. The single-view model uses CLIP ViT-B/32, with its weights unfrozen [code: models/sdfusion_img2shape_model.py#L85-L88] sdfusion-img2shape.yaml.
• Dropout. Paper: task-specific encoders "with dropout" §3.3 [abstract]; no rate is given. Code: the multi-modal model keeps each modality independently with probability 0.5 sdfusion_mm_model.py. The single-view model has no training-time dropout; its forward pass uses the image features directly sdfusion_img2shape_model.py. Its unconditional branch at inference is an all-zero image through the same encoder sdfusion_img2shape_model.py [code: models/sdfusion_img2shape_model.py#L381].
• Shape completion. Paper: blended diffusion, applied to the bottom-half partial shape §3.3 §4.1. Code: the unconditional model fills the known region in latent space at each step, from the noised encoding of the input, and the rest is generated ddim.py demo_util.py. The guidance call in that path omits unconditional_conditioning, so the guidance scale has no effect there [code: models/sdfusion_model.py#L438] [code: models/networks/diffusion_networks/samplers/ddim.py#L187-L188].
• VQ-VAE loss. Paper: reconstruction, commitment and VQ losses, no weights given §3.1. Code: β = 1.0 is passed with legacy=False [code: models/networks/vqvae_networks/network.py#L68-L69]. The quantiser's default is legacy=True, which puts β on the other term [code: models/networks/vqvae_networks/quantizer.py#L18-L19] [code: models/networks/vqvae_networks/quantizer.py#L89-L95].
• Symbol reuse. The paper uses E_φ for the VQ-VAE encoder §3.1 and E_φi for condition encoders §3.3. The code keeps these in separate modules, so use the code's names when implementing.
• Texturing. Paper: VolSDF density, a colour field and SDS against a pretrained 2D diffusion model §3.4. Code: no implementation found in the repository (searched for score distillation, VolSDF, DreamFusion and density terms in the Python files and notebooks). The 2D model is not in the repository.
• BuildingNet resolution. Paper: 128³ §4.1. Code: the configs and launchers checked contain only 64 for the grid (see Hyperparameters). The BuildingNet pipeline was not run.
• Sampler. Paper: the sampler is not named §3.2. Code: DDIM with η = 0 [code: options/base_options.py#L54].

No discrepancies were checked for the Pix3D evaluation and the ShapeGlot evaluator, because the evaluator is not in the repository.

Pitfalls

• Unstated training details. The paper's main text omits the dropout rate, the codebook size, the latent dimension, the VQ loss weight, the SDF truncation, the optimiser, the training length, and all sampling settings. The code fills these in, and an implementation that copies only the paper will not match the results reported.
• Truncation. The T-SDF is "truncated", but the paper gives no value. The code clamps distances to ±0.2 snet_dataset.py. The same clamp applies to the BuildingNet loader buildingnet_dataset.py.
• Guidance scales are per model. The code uses 1.0, 3.0 and 3.0 for the single-view, text and unconditional models (see Hyperparameters). A single scale carried across tasks will not reproduce the paper.
• Costs of guidance. The multi-modal guided branch runs three denoiser passes per step by stacking the input three times [code: models/networks/diffusion_networks/samplers/ddim.py#L199-L201]. Default sampling is 100 DDIM steps, and the unconditional model uses 200 [code: models/sdfusion_model.py#L117]. The paper gives no timing or memory figures.
• Table 2, one value. In the arXiv LaTeX source, the SDFusion chamfer distance on Pix3D is printed as 1.852, but an older value of 1.185 remains in a comment with the note that it was "entered wrong" (tables/tab_img2shape.tex). The printed value is the one used on this card. The table also omits the IoU column that the source keeps in a comment.
• Small margins, no error bars. The UHD margin over AutoSDF on ShapeNet is 0.0010 (0.0557 against 0.0567) tab. 1. No variance or seed is reported.
• Chamfer units. Pix3D chamfer values (2.267 to 4.670) have no stated unit or scale tab. 2. Compare them only within the same table.
• ShapeGlot is an author-trained proxy. The evaluator is trained by the authors §4.3; its architecture and training are not described. The text calls the gain "large" §4.3, but the table gap is 13 points with 15% confused tab. 3.
• Symbol and naming traps. E_φ means different encoders in §3.1 and §3.3. "Dropout" means per-modality masking in the code, not a dropout layer.
• Third-party code. The repository vendors CLIP sources (external/clip), TBB and isosurface binaries (preprocess/isosurface). The MIT licence covers the repository's own files; the vendored code was not checked.
• Texture model. Texturing depends on an off-the-shelf 2D diffusion model §3.4. Its weights and licence are outside this card; check them before commercial use.

Lineage

• Builds on: a 3D VQ-VAE (van den Oord et al., "Neural discrete representation learning"); latent diffusion (Rombach et al., "High-resolution image synthesis with latent diffusion models"); the simplified diffusion objective (Ho et al.); classifier-free guidance (Ho and Salimans, arXiv:2207.12598); blended diffusion (Avrahami et al.); score distillation sampling (DreamFusion, Poole et al., arXiv:2209.14988); VolSDF (Yariv et al.). None of these has a card.
• Compared against: AutoSDF (autoregressive SDF generation, CVPR 2022), MPC (multimodal shape completion, ECCV 2020), Pix2Vox, and the ShapeGlot evaluation protocol.