AutoSDF: Shape Priors for 3D Completion, Reconstruction and Generation

In plain words

A 3D shape can be stored as a grid of distance values, which is large. This paper first squeezes each small block of that grid into a short code taken from a fixed list of 512 codes, so a whole shape becomes a small 8 by 8 by 8 grid of list entries. A transformer then learns which combinations of entries make a plausible chair, table or lamp.

The unusual part is the order. The transformer is trained with the entries shown in random order, so at test time any set of known blocks can be fixed first and the rest filled in. A chair seen only from the back leg can still be completed. For a photo or a sentence, a second small network predicts which codes are likely at each position, and these predictions are combined with the shape model. The idea is an easy-to-state prior that the paper uses for three tasks, each judged against its own specialised baselines.

TL;DR

The method has two parts: a patch-wise VQ-VAE that maps a 64^3 truncated SDF to an 8^3 grid of 512 codes, and a transformer that models that grid with random generation orders. Completion samples the transformer conditioned on the observed codes. Image and text inputs add per-position naive conditionals, which are mixed with the prior using a weight alpha. All quantitative results are author-reported, on subsets, and not all headline claims hold in the tables (see Results and Pitfalls).

Problem

• Many 3D tasks are under-determined from one view or a partial scan, so a prior over plausible shapes is needed §1.
• Point-cloud completion baselines either give one accurate answer or many diverse ones, rarely both §2.
• A transformer over raw 3D voxels is too expensive: attention cost grows quadratically with input size, so the shape must first be compressed to a small discrete grid §3.
• Earlier non-sequential autoregressive models (on pixels) cannot synthesise high-resolution outputs for this reason §2.
• Specialised models trained for one task each (for example image-to-shape) must be retrained per task and cannot reuse a shared shape prior §4.

Method

The model has three learned parts: a discrete shape code (P-VQ-VAE), a non-sequential transformer prior over that code, and task-specific conditional networks. Each task reuses the first two and trains only a small conditional network §3.

1. Input. Shapes are normalised to the cube [-1, 1]^3 and their SDF is sampled on a 64^3 grid. The signed distance is then truncated at 0.2 to give a T-SDF app. VQ-VAE Training.
2. Patch-wise encoder E. The 64^3 volume is split into 512 non-overlapping 8^3 patches. Each patch is encoded on its own into one 256-dimensional vector, so each code depends only on local context §3.1.
3. Vector quantisation. Each vector is replaced by its nearest entry in a learned codebook of 512 entries with dimension 256. The result is an 8^3 grid Z of indices §3.1 eq. 1.
4. Decoder D. A decoder maps the whole grid of code vectors back to a 64^3 T-SDF. Decoding is joint, so it can use global context, while encoding is local §3.1.
5. VQ-VAE loss. Reconstruction, codebook and commitment terms. The paper writes the reconstruction term as a negative log-likelihood and does not state the commitment weight app. VQ-VAE Training.
6. Random-order transformer prior. A transformer models p(Z) by factorising it over a random permutation of the 512 grid positions. Each factor is conditioned on the already-known codes and on the location to predict. Training maximises the log-likelihood over sampled random orders §3.2 eq. 2.
7. Shape completion. Observed partial shapes are encoded to codes at their positions. The transformer then fills the remaining positions by sampling one at a time in a random order, and the decoder turns the completed grid into a shape §3.3 eq. 3 §4.1.
8. Naive conditionals. For image or text input C, a separate network predicts a per-position distribution over the 512 codes. A ResNet-18 (image) or BERT (text) encoder is followed by linear layers and 3D up-convolutions producing an 8^3 grid §3.3 app. Naive Condition Predictions.
9. Product with the prior. The conditional distribution is approximated by the prior multiplied by the per-position conditionals. The paper raises the prior to the power 1 - alpha and the conditional to the power alpha. The paper gives alpha = 0.75 for images and 0.5 for language §3.3 app. Naive Condition Predictions.
10. Baselines for the conditional tasks. The language task is compared against a joint model (JE) that predicts the codes directly, and against Text2Shape. Single-view reconstruction is compared against Pix2Vox and ResNet-based direct decoders §4.2 §4.3.

# P-VQ-VAE encoding: local patches, one code each  [§3.1] [eq. 1]
# Input:  T-SDF volume, shape (64, 64, 64), values clipped to [-0.2, 0.2]
# Output: grid of 512 codebook indices, shape (8, 8, 8)

def encode_shape(volume):
    # Split the volume into 512 non-overlapping 8x8x8 patches  [app. VQ-VAE]
    patches = split_into_patches(volume, patch_size=8)  # (512, 1, 8, 8, 8)
    # Encode each patch on its own, with no context from other patches  [§3.1]
    latents = encoder(patches)                          # (512, 256)
    # Replace each latent with its nearest of the 512 codebook entries  [§3.1]
    indices = nearest_codebook_index(latents, codebook)  # (512,)
    return indices.reshape(8, 8, 8)

# P-VQ-VAE decoding: joint decoding of the whole grid  [§3.1]
# Input:  grid of codebook indices, shape (8, 8, 8)
# Output: T-SDF volume, shape (64, 64, 64)

def decode_shape(indices):
    # Look up the 256-dim code vector of each grid cell  [§3.1]
    code_vectors = codebook[indices]                    # (8, 8, 8, 256)
    # The decoder sees the whole grid, so it can use global context  [§3.1]
    return decoder(code_vectors)                        # (64, 64, 64)

# Completion by sampling the random-order prior  [§3.2] [§4.1] [eq. 2] [eq. 3]
# Input:  observed (index, location) pairs, transformer prior
# Output: one completed index grid of shape (8, 8, 8)

def complete_shape(observed, prior):
    # Start from the codes that are already known  [§3.3] [eq. 3]
    known = list(observed)
    # Visit the unknown locations in a random order  [§3.2] [eq. 2]
    for location in random_order(unobserved_locations(observed)):
        # Distribution over 512 entries given known codes and this location  [§3.2]
        logits = prior(known, location)                 # (512,)
        # Sample a code from the distribution  [§3.3]
        code = sample_categorical(logits)
        known.append((code, location))
    return grid_from(known)

# Conditioning with a naive per-position term  [§3.3] [app. Naive Condition Predictions]
# Input:  prior logits (512,), naive log-probabilities (512,), weight alpha
# Output: mixed log-probabilities (512,) to sample from

def mix_prior_and_condition(prior_logits, naive_log_probs, alpha):
    # Prior raised to (1 - alpha), conditional raised to alpha, in log space  [app.]
    prior_log_probs = log_softmax(prior_logits)         # (512,)
    mixed = (1 - alpha) * prior_log_probs + alpha * naive_log_probs
    return log_softmax(mixed)                           # renormalise

The naive-conditional form in the last block is the paper's, written in log space. The released code uses a different input convention and top-k truncation, described under Paper vs code.

Hyperparameters

Values the paper states, followed by values the code releases. Training optimiser, learning rate, batch size and number of epochs are not stated in the paper.

Results

Author-reported. Shape completion (Table 1) uses held-out chairs from ShapeNet, with a subset chosen for compute reasons; the paper does not state its size §4.1 tab. 1. UHD is the completion fidelity (lower is better), TMD the diversity (higher is better). Units are not stated in the table.

Single-view reconstruction (Table 2) compares on ShapeNet renderings (Choy et al. voxel targets) and on Pix3D chairs, where predictions are downsampled to 32^3 for scoring §4.2. The paper says only that the method performs favourably across almost all categories; the table reports dataset averages §4.2 tab. 2.

Language-guided generation (Table 3) uses a 2-way neural evaluator with about 83% accuracy. Each description gets two generations; an outcome is "confused" when the evaluator's confidence differs by at most 0.2. The test subset is not stated §4.3 tab. 3.

Qualitative results (completions, multiple shapes per input, language examples) are shown in figures described in §4.1, §4.2 and §4.3. No figures are reproduced here.

Paper vs code

Official code: yccyenchicheng/AutoSDF at 4e372f1d. The owner is the first author's account. The paper's project page links it as "Code". The repository has no LICENSE file, so its code may not be quoted or reused; this card cites it by path and line only. Searched: project page, GitHub name search for "AutoSDF" (three repositories: the official one and two unrelated forks, which were not used).

• Sampling truncation. The paper samples from the full distribution at each step and does not mention top-k §3.2. The code keeps only the top-k logits before sampling: k = 30 for unconditional and completion calls rand_tf_model.py demo_shape_comp.ipynb demo_shape_comp.ipynb, k = 10 in the language and single-view demos demo-lang-conditional.ipynb demo_single_view_recon.ipynb.
• Samples per input. The paper generates k = 10 completions per partial shape §4.1. The completion demo uses a batch of 9 demo_shape_comp.ipynb. The paper's count and the demo's count are not the same protocol.
• Image alpha. The paper gives alpha = 0.75 for images app. Naive Condition Predictions. The single-view demo uses 0.7 demo_single_view_recon.ipynb, while the function default is 0.75 rand_tf_model.py.
• Mixing form. The code mixes (1 - alpha) times the transformer's log-softmax with alpha times the conditional term directly, then applies softmax and samples rand_tf_model.py rand_tf_model.py. The conditional term's form (log-probability or probability) depends on the caller. I did not verify what the image and language callers pass.
• Reconstruction loss. The paper writes the reconstruction term as -log p(X | Z) app. VQ-VAE Training. The code uses the mean absolute error (L1) between input and reconstruction losses.py losses.py, with codebook weight 1.0 pvqvae_snet.yaml.
• Adversarial and perceptual terms. The config sets a discriminator start step and a discriminator weight, and perceptual weight 0.0 pvqvae_snet.yaml. The discriminator code is commented out losses.py, and the perceptual term is switched off because its weight is 0 losses.py. The paper describes neither term.
• Commitment weight. The paper's VQ objective has unit weights app. VQ-VAE Training. The code's quantiser takes beta = 1.0, so the two terms are unit-weighted in effect. The code comment records a historical bug where beta was applied to the other term quantizer.py. With beta = 1 this has no effect.
• VQ batch size. The launcher passes batch size 6 train_pvqvae_snet.sh, while the config sets 12 pvqvae_snet.yaml. Which value is used depends on option parsing, which I did not check. The paper does not state a batch size.
• Optimiser and schedules. The paper states none of these; the code uses Adam with step decays pvqvae_model.py rand_tf_model.py. Values are listed in Hyperparameters.
• Patch encoding. The code splits the volume into 8^3 patches, applies a 3-down-sampling encoder to each, and folds the 8^3 grid back pvqvae_model.py pvqvae_model.py. This matches the paper.
• Truncation. The code clamps SDF values to the truncation threshold of 0.2 snet_dataset.py. This matches the paper.
• Single-view encoder. The code builds the image branch on a pretrained ResNet-18 resnet2vq_net.py. This matches the paper.

Pitfalls

• Optimiser and training schedule are absent from the paper. The learning rate, batch size, epoch count and gradient settings come only from the launchers and configs. Reimplementing from the paper alone leaves them open.
• Random-order training is central but underspecified. The paper says orders are sampled every iteration and that training is parallel with a causal-style mask app. Autoregressive Modeling. The released code's order sampling was not verified in this card.
• Top-k truncation changes the distribution. Sampling from the top 10 or 30 logits is not in the paper but is used by the released demos. Results with and without it may differ. Treat the choice as a hyperparameter and state it.
• Alpha is task-specific and tuned. The paper uses 0.75 for images and 0.5 for language, reporting that 0.5 worked better for language app. Naive Condition Predictions. Other tasks have no stated value.
• Table 1 does not support "favourable on both protocols". The text says the method performs favourably on fidelity and diversity §4.1. On the octant protocol its UHD (0.0599) is worse than MPC (0.0579) and PoinTr (0.0536). Its TMD advantage is large (0.0693 against 0.0376 for MPC).
• Baselines differ. The paper trains MPC separately per scenario, while its own method and PoinTr use one model §4.1. PoinTr's TMD is N/A. The paper also gives the baselines extra points within the truncation range for fairness §4.1.
• Language comparison is not controlled. The JE baseline uses a fixed raster order and only chair data, while the proposed method uses random orders and the full training set app. Naive Condition Predictions. The comparison therefore mixes several changes at once.
• Evaluator is weak. The neural listener is 2-way with about 83% held-out accuracy and was retrained with a different setup from the original app. Experimental Details. Confusion rates in Table 3 are judged by it.
• Subset evaluation. Shape completion is run on a subset of held-out chairs; the size is not stated §4.1. Table 2 reports averages without per-category values in the main text.
• Memory and time. Decoding 512 tokens sequentially at test time, with top-k and k = 10 samples, is slower than a single feed-forward pass. The paper gives no runtime numbers.
• Code is unlicensed. The released repository has no LICENSE file, so its code cannot be reused. The card uses it only for the points above.

Lineage

• Builds on: VQ-VAE (van den Oord et al., 2017) for the discrete latent, adapted here to a patch-wise encoder §3.1.
• Builds on: the non-sequential pixel transformer of Tulsiani and Gupta, which factorises over arbitrary observed sets; this paper moves it to the discrete 3D latent §2.
• Compared against: MPC and PoinTr for completion §4.1; Pix2Vox, ResNet2TSDF, ResNet2Voxel and Image-JE for single-view reconstruction §4.2; Text2Shape and a joint encoder-decoder for language §4.3.
• Evaluation uses the ShapeGlot dataset and neural listener of Achlioptas et al. for language §4.3.