Probabilistic Implicit Scene Completion

Network Architecture. The autoencoder is composed of two modules, encoder $g_{\phi}$ and decoder $f_{\omega}$ to convert between the sparse voxel embedding and implicit fields, which is depicted in Fig. 2. The encoder $g_{\phi}$ generates the sparse voxel embedding $s$ from the coordinate-distance pairs $P=\{(p,d_{p})|p\in\mathbb{R}^{3},d_{p}\in\mathbb{R}\}$, where $p$ is a 3D coordinate and $d_{p}$ is the (signed) distance to the surface, $s=g_{\phi}(P)$. The decoder $f_{\omega}$ regresses the implicit field value $d_{q}$ at the 3D position $q\in\mathbb{R}^{3}$ given the sparse voxel embedding $s$, $d_{q}=f_{\omega}(q,s)$.

The encoder $g_{\phi}$ is a a local PointNet (Qi et al. (2016)) implemented with MLP of 4 blocks with hidden dimension 128, as in ConvOcc (Peng et al. (2020)). Given a set of coordinate-distance pair $P$, we find the nearest cell $c_{p}$ for each coordinate-distance pair $(p,d_{p})$. Then we make the representation sparse by removing the coordinate-distance pair if the nearest cell $c_{p}$ do not contain the surface. For the remaining pairs, we normalize the coordinate $p$ to the local coordinate $p^{\prime}\in[-1,1]^{3}$ centered at the nearest voxel. Then we use the local coordinate to build a vector in $\mathbb{R}^{4}$ for coordinate-distance pair $(p^{\prime},d_{p})$ which serves as the input to the PointNet for the encoder $g_{\phi}$. The local PointNet architecture employs average pooling to the point-coordinate pairs that belong to the same cell.

The architecture of decoder $f_{\omega}$, inspired by IFNet (Chibane et al. (2020a)) is depicted in Fig. 6. As shown on the left, the decoder $f_{\omega}$ can be decomposed into the sparse convolution $f_{\omega_{1}}$ that extracts hierarchical feature, followed by trilinear interpolation of the extracted feature $\Psi$, and final MLP layers $f_{\omega_{2}}$. The sparse convolution layer $f_{\omega_{1}}$ aggregates the information into multi-level grids in $n$ different resolutions, i.e., $f_{\omega_{1}}(s)=F_{1},F_{2},...,F_{n}$, Starting from the original resolution $F_{1}$, $F_{k}$ contains a grid that is downsampled $k$ times. However, in contrast to the original IFNet, our grid features are sparse, since they only store the downsampled occupied cells of sparse voxel embedding $s$. The sparse representation $F_{k}=\{(c,o_{c},e_{c})|c\in\mathds{Z}^{3},o_{c}\in\{0,1\},e_{c}\in\mathbb{R}^{L}\}$ is composed of occupancy $o_{c}$ and $L$-dimensional feature $e_{c}$ for each cell, like sparse voxel embedding. The grid points $c$ that are not occupied are considered as having zero features. We use $n=3$ levels, with feature dimension $L=32$ for all experiments.

The multi-scale features are then aggregated to define the feature for a query point $q$: $\Psi(f_{\omega_{1}}(s),q)=\sum_{k}\psi_{k}(F_{k},q)$. Here $\psi_{k}(F_{k},q)\in\mathbb{R}^{L}$ is the trilinear interpolation of features in discrete grid $F_{k}$ to define the feature at a particular position $q$, similar to IFNet. We apply trilinear interpolation at each resolution $k$ then combine the features.

Lastly, the MLP layer maps the feature at the query point $q$ to the implicit function value $\hat{d}_{q}$, $f_{\omega_{2}}:\mathbb{R}^{L}\rightarrow\mathbb{R}$. The MLP is composed of 4 ResNet blocks with hidden dimension 128 with $\tanh$ as the final activation function. The final value $\hat{d}_{q}$ is the truncated signed distance within the range of $[-1,1]$ except for 3DFront (Fu et al. (2021)) which does not have a watertight mesh. We use unsigned distance function for 3DFront as in NDF (Chibane et al. (2020b)) by changing the codomain of the decoder to $[0,1]$ with sigmoid as the final activation function. Thus, the decoder can be summarized as $f_{\omega}(s,q)=f_{\omega_{2}}(\Psi(f_{\omega_{1}}(s),q))$.

We would like to emphasize that the autoencoder does not use any dense 3D convolutions, and thus the implementation is efficient, especially when reconstructing the sparse voxel embedding. Since the decoding architecture only consists of sparse convolutions and MLPs, the memory requirement increases linearly with respect to the number of occupied voxels. In contrast, the requirement for dense 3D convolution increases cubic with respect to the maximum size of the scene.

Reconstruction Results. Fig. 7 shows a reconstruction of a chair (left) using signed distance fields and a scene in 3DFront (right) using unsigned distance fields. The figure shows that our autoencoder can faithfully reconstruct a continuous surface using the sparse voxel embedding. However, the usage of unsigned distance fields (Chibane et al. (2020b)) representation tends to generate relatively thicker reconstructions compared to that of signed distance fields. Without a clear distinction between the interior/exterior, the unsigned representation makes it difficult to scrutinize the zero-level isosurface and thus results in comparably thick reconstruction. Thus, the flexibility of being able to represent any mesh with various topology comes at a cost of expressive power, in contrast to signed distance fields which can only represent closed surfaces. Our work can be further improved by utilizing more powerful implicit representation techniques, such as SIREN (Sitzmann et al. (2020)), but we leave it as future work.

Hyperparameters and Implementation Details. We train the autoencoder with Adam (Kingma & Ba (2015)) optimizer with learning rate 5e-4 and use batch size of 3, 1, and 4 for ShapeNet scene (Peng et al. (2020)), 3DFront (Fu et al. (2021)), and ShapeNet (Chang et al. (2015)) dataset, respectively. The autoencoder for scene completion is pretrained with all classes of ShapeNet assuming that the local geometric details for scenes are represented with the patches from individual objects.

Network Architecture. Following GCA (Zhang et al. (2021)), we implement the transition model by using a variant of U-Net (Ronneberger et al. (2015)). The input to the sparse convolution network is the sparse voxel embedding $s^{t}$, and the output feature size of the last convolution is $K+1$, where we use generalized sparse convolution for the last layer to obtain the occupancy (with dimension 1) and the latent code (with dimension $K$) of the neighborhood. Unlike GCA, our implementation does not need an additional cell-wise aggregation step.

$s^{0}$ Conditioned Model. The $s^{0}$ conditioned model is a variant of our model, where we employ the use of encoder $g_{\phi}$ to estimate the initial latent codes of initial state $s^{0}$ and condition every transition kernel on $s^{0}$. Given incomplete point cloud $R=\{c_{1},c_{2},...,c_{M}\}$ with $c_{i}\in\mathbb{R}^{3}$, the variant constructs the initial $s_{0}$ by encoding a set of coordinated-distance pair $P_{R}=\{(c,0)|c\in R\}$, where the distance is set to 0 for every coordinate in $R$. Since the point cloud is sampled from the surface of the original geometry, this is a reasonable estimate of the latent code.

The transition kernel $p_{\theta}$ is conditioned on $s^{0}$ by taking a sparse voxel embedding $\hat{s}^{t}$ of latent code dimension $2K$ that concatenates the current state $s^{t}$ and $s^{0}$ as the input. Specifically, we define $\hat{s^{t}}$ to be $\hat{s}^{t}=\{(c,o_{c},z_{c})|o_{c}=o^{s^{t}}_{c}\vee o^{s^{0}}_{c},z_{c}=[z^{s^{0}}_{c},z^{s^{t}}_{c}]\in\mathbb{R}^{2K},c\in\mathbb{Z}^{3}\}$, where the $o^{s^{t}}_{c},o^{s^{0}}_{c},z^{s^{t}}_{c},z^{s^{0}}_{c}$ are the occupancies and latent codes of $s^{t}$ and $s^{0}$. The occupancy is set to one if either $s^{t}$ or $s^{0}$ is occupied and the latent code is set to be the concatenation of latent codes of $s^{0}$ and $s^{t}$. Thus, the only modification to the neural network $p_{\theta}$ is the size of input feature dimension, which is increased from $K$ to $2K$.

Hyperparameters and Implementation Details We train our model with Adam (Kingma & Ba (2015)) optimizer with learning rate 5e-4. The transition kernel is trained using infusion schedule $\alpha^{t}=0.1+0.005t$ and standard deviation schedule $\sigma^{t}=10^{-1-0.01t}$. We use the accelerated training scheme introduced in GCA (Zhang et al. (2021)) by managing the time steps based on the current step. We use neighborhood size $r=2$, with metric $d$ induced by $L_{1}$ norm, except for the ShapeNet (Chang et al. (2015)) experiment on $64^{3}$ voxel resolution which uses $r=3$. For transition model training, we utilize maximum batch size of 32, 16 and 6, for ShapeNet scene, 3DFront, ShapeNet dataset, respectively. We use adaptive batch size for the transition model, which is determined by the total number of cells in the minibatch for efficient GPU memory usage. Lastly, we use mode seeking steps $T^{\prime}=5$ for all experiments.

Visualizations and Testing. All surface visualizations are made by upsampling the original voxel resolution 4 times, where we efficiently query from the upsampled grid only near the occupied cells of the final state $s^{T+T^{\prime}}$. Then, the mesh is extracted with marching cubes (Lorensen & Cline (1987)) with no other postprocessing applied. From the $4\times$ upsampled high resolution grid, we extract cells with the distance value below 0.5 which serve as the point cloud to compute the metrics for testing.

We use the official implementation or the code provided by the authors by contact with default hyperparameters unless stated otherwise. All results are obtained by training on our dataset. The output point cloud used for testing, is obtained by sampling from mesh with the same resolution if possible, unless stated otherwise. For GCA (Zhang et al. (2021)) we use the same hyperparameters as ours for fair comparison. The mesh is created by using marching cubes (Lorensen & Cline (1987)) with the voxel resolution. We use the volume $64^{3}$ encoder for ConvOcc (Peng et al. (2020)) which achieves the best results in the original ShapeNet scene dataset (Peng et al. (2020)). For IFNet (Chibane et al. (2020a)), we use the ShapeNet128Vox model as used in the work of Siddiqui et al. with occupancy implicit representation. cGAN (Wu et al. (2020)) is tested by 2,048 generated point cloud.

ShapeNet Scene. ShapeNet scene (Peng et al. (2020)) contains synthetic rooms that contain multiple ShapeNet (Chang et al. (2015)) objects, which has randomly scaled floor and random walls, normalized so that the longest length of the bounding box has length of 1. The number of rooms for train/val/test split are 3750/250/995. We sample SDF points for training our model. The input is created by iteratively sampling points from the surface and removing points within distance 0.1 for each instance. We run 20 iterations and each removal is revoked if the aggregated removal exceeds that of the guaranteed rate. Since we are more interested in the reconstruction of objects, the walls and floors are guaranteed to have minimum rate of 0.8, regardless of the defined minimum rate. Lastly, we sample 10,000 points and add normal noise with standard deviation 0.005. For computing the metrics at test time, we sample 100,000 points.

3DFront. 3DFront (Fu et al. (2021)) is a large-scale indoor synthetic scene dataset with professionally designed layouts and contains high quality 3D objects. We collect rooms that contain multiple objects, where 10% of the biggest rooms are removed from the data to discard rooms with immense size for stable training. Thus, 7044/904/887 rooms are collected as train/val/test rooms. The incomplete point cloud is obtained by randomly sampling with density 219 points/$\text{m}^{2}$ and then removed and noised as in ShapeNet scene, but with removal sphere distance proportional to the object size and noise of standard deviation of 1cm. During evaluation, we sample points proportionally to the area for each ground truth surface with density of 875 points/$\text{m}^{2}$ to compute the metrics.

ShapeNet. We use chair, sofa, table classes of ShapeNet (Chang et al. (2015)), where the signed distance values of the query points are obtained by using the code of DeepSDF (Park et al. (2019)). The shapes are normalized to fit in a bounding box of $[-1,1]^{3}$. The partial observation is generated by extracting points within a sphere of radius 0.5, centered at one of the surface points. Then we sample 1,024 points to create a sparse input and compare the metrics by sampling 2,048 points for a fair comparison against cGAN (Wu et al. (2020)).

We derive the variational lower bound for cGCA (Eq. (11)) below. The lower bound is derived by the work of Sohl-Dickstein et al. (2015) and we include it for completeness.

The second term of the right hand side can be converted into KL divergence as following:

Thus,

is derived.

We derive the decomposition of KL divergence (Eq. (12) ) below.

Thus, Eq. (12) is derived.

In this section, we claim that $\mathcal{L}_{\mathrm{final}}$ can be replaced by $\mathcal{L}_{T-1}$, which allows us to train using a simplified objective. First, we show that $\mathcal{L}_{\mathrm{final}}=\mathbb{E}_{q_{\theta}}[\log p_{\theta}(x|s^{T-1})]=\infty$ due to the usage of Dirac distribution and introduce a non-diverging likelihood by replacing the Dirac function $\delta_{0}(z_{c})$ with indicator function $\mathds{1}[z_{c}=0]$. Recall,

Thus, if any cell $c\in\mathcal{N}(s^{T-1})$ is unoccupied, the likelihood diverges to $\infty$ since $z^{x}_{c}=0$. This disables us to use negative log-likelihood as a loss function, since all the gradients will diverge.

This problem can be easily resolved by substituting $\delta_{0}(z^{x}_{c})$ with indicator function $\mathds{1}[z^{x}_{c}=0]$. Both functions have non-zero value only at $z_{c}\neq 0$, but the former has value of $\infty$ while the latter has value $1$. While $\mathds{1}[z^{x}_{c}=0]$ is not a probability measure, i.e. $\int_{z^{x}_{c}}\mathds{1}[z^{x}_{c}=0]\neq 1$, replacing $\delta_{0}(z^{x}_{c})$ will have the effect of reweighting the likelihood at $z^{x}_{c}=0$ from $\infty$ to $1$, with the likelihood at other values $z^{x}_{c}\neq 0$ unchanged. Thus, $\mathds{1}[z^{x}_{c}=0]$ is a natural replacement of $\delta_{0}(z^{x}_{c})$ for computing a valid likelihood.