LaGeM: A Large Geometry Model for 3D Representation Learning and Diffusion

## Abstract

Abstract This paper introduces a novel hierarchical autoencoder that maps 3D models into a highly compressed latent space. The hierarchical autoencoder is specifically designed to tackle the challenges arising from large-scale datasets and generative modeling using diffusion. Different from previous approaches that only work on a regular image or volume grid, our hierarchical autoencoder operates on unordered sets of vectors. Each level of the autoencoder controls different geometric levels of detail. We show that the model can be used to represent a wide range of 3D models while faithfully representing high-resolution geometry details. The training of the new architecture takes 0.70x time and 0.58x memory compared to the baseline.
We also explore how the new representation can be used for generative modeling. Specifically, we propose a cascaded diffusion framework where each stage is conditioned on the previous stage. Our design extends existing cascaded designs for image and volume grids to vector sets.

## 1 Introduction

Diffusion models are currently the best-performing models for image, video, and 3D object generation.
For 3D object generation, there are two main branches of research. The first branch, pioneered by Dreamfusion (Poole et al., ), aims to lift 2D diffusion models to 3D model generation. The advantage of this method is that it can benefit from the large-scale 2D datasets used for training 2D diffusion models and it sparked a lot of follow-up work (Poole et al., ; Wang et al., ; Lin et al., ; Chen et al., ; Wang et al., ; Qian et al., ; Tang et al., ; Yi et al., ; Wang & Shi, ; Liu et al., ; Long et al., ; Zheng et al., ; Li et al., ; Ho et al., ; Xu et al., ).
The second branch tackles the training on 3D datasets directly. The advantage of this method is that it is more direct and leads to faster inference times (Mittal et al., ; Yan et al., ; Zhang et al., ; Zeng et al., ; Zheng et al., ; Hui et al., ; Zhang et al., ; Siddiqui et al., ; Chen et al., ; ). Our work sets out to contribute to this second branch of methods.

Among these 3D native generation methods, 3DShape2VecSet (Zhang et al., ) (or VecSet for short) has been proven to be an effective method to encode 3D geometry. It proposed an autoencoder to find an efficient representation for 3D models as a set of vectors. Because of the high reconstruction quality and compactness of the latent space, the method alleviates the difficulty of training 3D generative models. Some other works (Zhao et al., ; Cao et al., ; Dong et al., ; Petrov et al., ; Zhang et al., ; Zhang & Wonka, ) follow the VecSet representation.
We noticed that VecSet’s expressiveness is limited by the number of latent vectors. It is overfitting on smaller datasets like ShapeNet and is unable to scale to larger datasets. To improve the expressiveness, we need to scale up the latent size and the training dataset. The straightforward way is to employ hundreds of GPUs for training which is expensive (Zhang et al., ). Thus, our goal is to reduce the training cost in terms of time and memory consumption while achieving similar or even better autoencoding quality.

Figure: Figure 1: Autoencoders. We show different autoencoder architectures here, including AE (AutoEncoder), U-Net, VAE (Kingma, ), NVAE (Vahdat & Kautz, ), VecSet (Zhang et al., ) and the proposed LaGeM. VAE and NVAE are for image data, while VecSet and LaGeM are for geometry (distance function) data. In the top row, VAE and VecSet are using a single scale latent to represent the data. Both NVAE and LaGeM use multi-scale latents to represent data. All the previous works VAE, NVAE, and VecSet apply KL divergence in the bottleneck to regularize the latent space, while in this work, we apply standardization in the bottleneck.

In the image domain, NVAE (Vahdat & Kautz, ) extended the design of the variational autoencoder (VAE) (Kingma, ) to a hierarchical VAE based on the design of the U-Net. The latent space of the NVAE is a multi-scale latent grid and the reconstruction quality of the images from the NVAE improves a lot over the VAE. An illustration of the architectures can be found in Fig. 1.
We draw inspiration from the design of the NVAE and design a multi-scale latent VecSet representation, called *LaGeM*. We train our architecture on a large-scale geometry dataset Objaverse (Deitke et al., ) and improve training time by 0.7 and memory consumption by 0.58 compared to VecSet.

| Latents | Controlling |
| --- | --- |
| Level 3 | Main Structure |
| Level 2 | Major Details |
| Level 1 | Minor Details |

Additionally, we also propose a cascaded generative model for the hierarchical latent space. We generate the latent VecSet from the lower resolution level to the highest resolution level stage-by-stage. In each stage, we use the previously generated latents as conditioning information. As a result, this enables control over the level of detail of the generated geometry.

We summarize our contributions as follows:

- •
We propose a hierarchical autoencoder architecture with faster training time and low memory consumption. The latent space is composed of several levels.
- •
The model is capable of training on large-scale datasets like objaverse.
- •
We propose a cascaded diffusion model to generate 3D geometry in the hierarchical latent space. This enables control of the level of detail of the generated model.

**Table 1: Geometric Latent Representation and Generation.**
| Method |  | Learning Method | Latent Rep | Hierarchies |
| --- | --- | --- | --- | --- |
| ShapeFormer (Yan et al., ) |  | AutoEncoder | Sparse Volume | Single |
| 3DILG (Zhang et al., ) |  | AutoEncoder | Irregular Grid | Single |
| LION (Zeng et al., ) |  | AutoEncoder | Latent Points | Multi |
| TriplaneDiffusion (Shue et al., ) |  | AutoDeocder | Planes | Single |
| SDFusion (Cheng et al., ) |  | AutoEncoder | Volume | Single |
| 3DShape2VecSet (Zhang et al., ) |  | AutoEncoder | VecSet | Single |
| HyperDiffusion (Erkoç et al., ) |  | Per-Object Optimization | Network Weight | Single |
| XCube (Ren et al., ) |  | AutoEncoder | Sparse Volume | Multi |
| Mosaic-SDF (Yariv et al., ) |  | Per-Object Optimization | Irregular Grid | Single |
| 3DTopia-XL(Chen et al., ) |  | Per-Object Optimization | Irregular Grid | Single |
| OctFusion (Xiong et al., ) |  | AutoEncoder | Sparse Volume | Multi |
| LaGeM\faGem[regular] (Ours) |  | AutoEncoder | VecSet | Multi |

## 2 Related works

We show an overview of latent 3D generative models in Table 1, particularly focusing on the type of latent space used.

### 2.1 Learning Methods

Usually, a learning method is required to convert 3D geometry to latent space. 1) One way to do this is to convert 3d geometry to latent space with a per-object optimization method, e.g. (Erkoç et al., ; Yariv et al., ). For larger datasets, this approach is very time-consuming. 2) Alternatively, auto-decoder, e.g., DeepSDF (Park et al., ), jointly optimize the latent space for all objects in the dataset. However, as there is no encoder, new objects cannot be mapped to latent space easily. 3) Therefore, a commonly used framework is the auto-encoder. The optimization is efficient because it is performed jointly for all objects in the dataset, and new objects not in the training set can be quickly encoded using the encoder. Thus, we also build on this approach.

### 2.2 Latent Representations

Early methods used regular grids (Yan et al., ; Cheng et al., ) as the latent representation because of their simple structure. We can easily use convolutional layers to process volume data. To represent high-quality geometric details, we need large-resolution volumes. This makes the training even more difficult because of the $O(n^{3})$ complexity. A way to solve this problem is to introduce sparsity (Ren et al., ) to the representation like octrees (Xiong et al., ) or sparse irregular grids (Zhang et al., ; Yariv et al., ). Both structures have the potential to represent high-quality 3D models, but generating irregular structures explicitly is difficult for diffusion models. Different from the above mentioned approaches, 3DShape2VecSet (Zhang et al., ) is proposed to solve the reconstruction problem without using any sparse structures. The representation is easy to use. In this paper, we investigate how to improve the VecSet representation. Compared to Zhang et al. (), our goal is to obtain an even higher-quality latent space by introducing Level of Latents (LoL).

### 2.3 Cascaded Generation

In the field of image generation, there are multiple cascaded diffusion models,e.g., (Ho et al., ; Saharia et al., ). In the 3D domain, some works (Zeng et al., ; Ren et al., ) also modeled geometries with hierarchical latents and proposed 3D generative models using cascaded diffusion models. Our work encodes 3D geometry into hierarchical VecSets. Thus, it is straightforward to consider cascaded latent diffusion to train generative models in our latent space.

Figure: Figure 2: Pipeline. We proposed a U-Net-style transformer for the autoencoding. In this way, we obtain a hierarchical latent space, which contains several levels of latents. To train the generative diffusion models in the latent space, we propose the cascaded latent diffusion models.
Refer to caption: x1.png

## 3 Methodology

### 3.1 Background of VecSet Representations

The VecSet (Zhang et al., ) representation converts a dense point cloud to a latent vector set ${\mathcal{Z}}=\{{\mathbf{z}}_{1},{\mathbf{z}}_{2},\dots,{\mathbf{z}}_{M}\}$ with ${\mathbf{z}}\in{\mathbb{R}}^{D}$ so that an occupancy/distance function ${\mathcal{O}}({\mathbf{p}})$ can be recovered from the vector set. The simplified network is illustrated in Fig. 3.

#### Encoding.

The process first downsamples the 3D input point cloud ${\mathcal{P}}^{\text{Input}}=\{{\mathbf{p}}_{i}\}_{i=1,\dots,N}$ with furthest point sampling (FPS), ${\mathcal{P}}=\mathrm{FPS}({\mathcal{P}}^{\text{Input}},r)$,
where $r$ is the down-sampling ratio, and ${\mathcal{P}}$ is a low-resolution version of ${\mathcal{P}}^{\text{Input}}$. Then ${\mathcal{P}}^{\text{Input}}$ is converted to an unordered set with cross-attention

$$ $\mathrm{CA}(Q=\mathrm{PE}({\mathcal{P}}),K=\mathrm{PE}({\mathcal{P}}^{\text{ Input}}),V=\mathrm{PE}({\mathcal{P}}^{\text{Input}}))={\mathcal{X}}=\{{\mathbf {x}}\in{\mathbb{R}}^{C}\}_{i=1,2,\dots,M},$ (1) $$

where PE is a positional embedding function (Zhang et al., ) and $\mathrm{CA}(\cdot,\cdot,\cdot)$ is a cross-attention module. We also write $\mathrm{CA}({\mathcal{P}},{\mathcal{P}}^{\text{Input}})$ for short.
Here, the positional embedding used to project a 3D coordinate ${\mathbf{p}}\in{\mathbb{R}}^{3}$ to the high dimensional space ${\mathbb{R}}^{C}$ is omitted for simplicity.
To obtain a highly compressed latent space, the vectors in ${\mathcal{X}}$ are further compressed to a lower-dimensional space ${\mathbb{R}}^{D}$ where $D\leq C$ (Feature to Latent, or FtoL in short),

$$ $\mathrm{FtoL}({\mathcal{X}})={\mathcal{Z}}=\{{\mathbf{z}}\in{\mathbb{R}}^{D}\} _{i=1,2,\dots,M}.$ (2) $$

This compression step is also regularized by KL divergence.

#### Decoding.

Each latent vector in ${\mathcal{Z}}$ is first converted back to feature space ${\mathbb{R}}^{C}$ (Latent to Feature, or LtoF in short),

$$ $\mathrm{LtoF}({\mathcal{Z}})={\mathcal{X}}^{\prime}=\{{\mathbf{x}}^{\prime}\in {\mathbb{R}}^{C}\}_{i=1,2,\dots,M}.$ (3) $$

The features ${\mathcal{X}}^{\prime}$ are fed into a series self-attention layers to obtain final occupancy/distance function representations ${\mathcal{F}}$,

$$ $\mathrm{SAs}({\mathcal{X}}^{\prime})={\mathcal{F}}=\{{\mathbf{f}}\in{\mathbb{R }}^{C}\}_{i=1,2,\dots,M},$ (4) $$

where $\mathrm{SAs}(\cdot)$ is implemented using several self-attention layers. Now we can decode a continuous function. For a continuous coordinate in the space ${\mathbb{R}}^{3}$, we have

$$ ${\mathcal{O}}(\mathbf{p})=\mathrm{FC}\left(\mathrm{CA}(\mathbf{p},{\mathcal{F} })\right)\in{\mathbb{R}}.$ (5) $$

See Table 2 for more details on $\mathrm{FtoL}(\cdot)$ and $\mathrm{LtoF}(\cdot)$.

Figure: Figure 3: Geometry Autoencoder. The design from VecSet (Zhang et al., ) can be seen as a special case of the proposed LaGeM network with only one level.

Figure: Figure 4: LaGeM architecture. We show an illustration with 3 levels of latents.

**Table 2: Regularization in the Bottleneck. We compare the proposed regularization and VAE. We do not need an explicit loss to regularize the latent space.**
|  | Features to Latents (FtoL) | Latent Loss | Latents to Features (LtoF) |  |
| --- | --- | --- | --- | --- |
| VAE | $\bm{\mu}=\mathrm{FC}_{\mu}({\mathbf{x}})$ | ${\mathbf{z}}=\bm{\mu}+\bm{\sigma}\odot\bm{\epsilon}$ | KL Divergence | ${\mathbf{x}}^{\prime}=\mathrm{FC}_{\text{up}}({\mathbf{z}})$ |
| $\bm{\sigma}=\mathrm{FC}_{\sigma}({\mathbf{x}})$ |  |  |  |  |
| Ours | $\bar{{\mathbf{z}}}=\mathrm{FC}_{\text{down}}({\mathbf{x}})$ | $\displaystyle{\mathbf{z}}=\frac{\bar{{\mathbf{z}}}-\mathrm{E}[\bar{{\mathbf{z}<br>}}]}{\sqrt{\mathrm{Var}[\bar{{\mathbf{z}}}]}}$ | - | ${\mathbf{x}}^{\prime}=\mathrm{FC}_{\text{up}}({\mathbf{z}}\odot\bm{\gamma}+\bm<br>{\beta})$ |

### 3.2 Hierarchical VecSet

The complexity of the self-attention layers in Eq. 4 is $O(M^{2})$, i.e., quadratic in the number of input vectors. This severely affects the training time when $M$ is large. However, to represent high-quality geometry details, we usually need a large $M$. This is making training a large VecSet network more challenging (for example $M=2048$ in CLAY (Zhang et al., )). Motivated by the design of the U-Net and NVAE (Vahdat & Kautz, ), we propose a new network. Specifically,
in the design of the U-Net (see an illustration in Fig. 1), image feature grids are downsampled to lower resolutions where some convolution blocks are applied, and then upsampled to the original resolution. In this way, we can avoid performing convolutional layers in high resolution images (which can be time-consuming). We transferred this idea to the VecSet representations. Two necessary building blocks are operations to down-sample and up-sample a VecSet. Inspired by the design of 3DShape2VecSet (Zhang et al., ) (an illustration can be found in Fig. 3), we interpret the cross attention in the encoder part as a down-sampling operator. Similarly, we can also use it for up-sampling. The resulting network is shown in Fig. 4.

We have $L$ levels in the U-Net-style transformer, where we number the levels from one (highest resolution) to $L$ (lowest resolution). For notational convenience, we denote the input point cloud as level 0. In the $i$-th level, we first obtain a lower resolution of the point clouds in the $(i-1)$-th level, $\mathrm{FPS}({\mathcal{P}}_{i-1},r_{i-1})={\mathcal{P}}_{i}$
where ${\mathcal{P}}_{0}$ is the input point cloud. We use cross attention to compress the feature set $\mathrm{CA}({\mathcal{P}}_{i},{\mathcal{P}}_{i-1})={\mathcal{X}}_{i}$.
Different from previous approaches, we propose a new to way regularize the latent space,

$$ $\mathrm{FtoL}({\mathcal{X}}_{i})=\mathrm{ZeroMeanAndUnitVariance}(\mathrm{FC}_ {\text{down}}({\mathcal{X}}_{i}))={\mathcal{Z}}_{i},$ (6) $$

where we normalize each vector in the set to have zero mean and unit variance $({\mathbf{z}}-\mathrm{E}[{\mathbf{z}}])/\sqrt{\mathrm{Var}[{\mathbf{z}}]}$ (It is often called *standardization* in machine learning which is used to standardize the features present in the data in a fixed range.). The motivation behind this design is that diffusion starts with Gaussian noise which also has zero mean and unit variance. In this way, we enforce both our latent space and the initial Gaussian noise to have similar properties.
To map the latents back to features, we first scale and shift latents back ${\mathbf{z}}\odot\bm{\gamma}+\bm{\beta}$ (both $\bm{\gamma}$ and $\bm{\beta}$ are learnable parameters like in Layer Normalization (Lei Ba et al., )),

$$ $\mathrm{LtoF}({\mathcal{Z}}_{i})=\mathrm{FC}_{\text{up}}(\mathrm{ScaleAndShift }({\mathcal{Z}}_{i}))={\mathcal{X}}^{\prime}_{i}.$ (7) $$

Unlike KL divergence in a VAE, we do not need an explicit loss term for the latent space. See Table 2 for a comparison between the proposed regularization and commonly used KL divergence in VAEs.

Figure: Figure 5: Multiresolution Features

Inspired by the down-sampling usage of cross attention in Zhang et al. (), we generalize it to *resampling*. Here we use it as *upsampling for unordered set* ${\mathcal{F}}_{i}$.
Before feeding the features to self attention layers, we first upsample features ${\mathcal{F}}_{i+1}$ from lower resolution levels and apply self attentions,

$$ $\mathrm{SAs}(\mathrm{CA}({\mathcal{X}}^{\prime}_{i},{\mathcal{F}}_{i+1}))={ \mathcal{F}}_{i}.$ (8) $$

The query function in Eq. 5 is changed to

$$ ${\mathcal{O}}(\mathbf{p})=\mathrm{FC}\left(\left[\mathrm{CA}(\mathbf{p},{ \mathcal{F}}_{1})|\cdots|\mathrm{CA}(\mathbf{p},{\mathcal{F}}_{L})\right] \right)\in{\mathbb{R}},$ (9) $$

where $[\cdot|\cdot|\cdots|\cdot]$ is the symbol for concatenation. This means we are using features from all levels to build the final (occupancy) function representation (Fig. 5).

### 3.3 Diffusion

Cascaded Diffusion (Ho et al., ) proposed a method for generating high-resolution images. The method is composed of several stages, where each stage is a conditioned diffusion model. Motivated by this, we propose a cascaded latent diffusion model. In Cascaded Diffusion, images generated from the previous stage are used as a condition in the next stage. We build a cascaded latent diffusion model based on Cascaded Diffusion. Formally, the optimization goal (for our three-level implementation) is as follows,

$$ $\displaystyle\min_{D_{3}}\left\|D_{3}(\tilde{{\mathcal{Z}}}_{3}(t),t,{\mathcal {C}}\phantom{,{\mathcal{Z}}_{3},{\mathcal{Z}}_{2}})-{\mathcal{Z}}_{3}\right\|,$ (10) $\displaystyle\min_{D_{2}}\left\|D_{2}(\tilde{{\mathcal{Z}}}_{2}(t),t,{\mathcal {C}},{\mathcal{Z}}_{3}\phantom{,{\mathcal{Z}}_{2}})-{\mathcal{Z}}_{2}\right\|,$ $\displaystyle\min_{D_{1}}\left\|D_{1}(\tilde{{\mathcal{Z}}}_{1}(t),t,{\mathcal {C}},{\mathcal{Z}}_{3},{\mathcal{Z}}_{2})-{\mathcal{Z}}_{1}\right\|,$ $$

where $D_{i}$ is a denoising network, $t$ represents timestep or noise level, $\tilde{{\mathcal{Z}}}_{i}(t)$ is the noised version (at timestep $t$) of the latent, $\mathcal{C}$ is optional condition information (e.g., text, images, or categories). The network design is based on DiT (Peebles & Xie, ). To generate latents ${\mathcal{Z}}_{i}$, we need latents from previous stages ${\mathcal{Z}}_{>i}$. For diffusion-based image super-resolution methods, this is often done by bilinearly interpolating small images and concatenating them with denoising networks’ inputs. As shown in the previous section, we use cross attention for resampling (both down-sampling and upsampling). Here we also utilize cross attention to upsample a latent set. Specifically, assuming we are training a denoising network for ${\mathcal{Z}}_{2}$, the input of the network is $\tilde{{\mathcal{Z}}}_{2}(t)$,

$$ $\mathrm{CA}(\tilde{{\mathcal{Z}}}_{2}(t),{\mathcal{Z}}_{3}).$ (11) $$

Similarly, for ${\mathcal{Z}}_{1}$,

$$ $\mathrm{CA}(\mathrm{CA}(\tilde{{\mathcal{Z}}}_{1}(t),{\mathcal{Z}}_{3}),{ \mathcal{Z}}_{2}).$ (12) $$

In this way, we are gathering information from previous stages.
See Fig. 6 for an illustration about the pipeline.

Figure: Figure 6: Cascaded Latent Diffusion.

**Table 3: Running Statistics of LaGeM. When using a small number (512) of latent vectors, our model uses 0.87x time and 0.66x memory during training. For larger models (2k latent vectors), the advantage is even more significant (0.7x time and 0.58x memory).**
|  | VecSet | LaGeM | VecSet | LaGeM | VecSet | LaGeM |
| --- | --- | --- | --- | --- | --- | --- |
| Batch Size | 64 | 8 | 4 |  |  |  |
| Self Attn Layers | 24 | 8/8/8 | 24 | 8/8/8 | 24 | 8/8/8 |
| Attn Channels | 512 | 512/512/512 | 1k | 1k/1k/1k | 1k | 1k/1k/1k |
| # Parameters (M) | 106.13 | 125.15 | 424.24 | 499.85 | 424.24 | 499.85 |
| # Latent Vectors | 512 | 32/128/512 | 2k | 128/512/2k | 2k | 128/512/2k |
| # Latent Channels | 8 | 32/16/8 | 64 | 64/32/16 | 64 | 64/32/16 |
| Training Memory (M) | 56,125 | 37,055 (0.66$\times$) | OOM | 53,791 (-) | 54,543 | 31,662 (0.58$\times$) |
| Training Iteration (sec) | 0.6481 | 0.5658 (0.87$\times$) | - | 0.7714 (-) | 0.6902 | 0.4853 (0.70$\times$) |

## 4 Experiments

### 4.1 Autoencoding Model

The main autoencoding experiment is trained on Objaverse (Deitke et al., ). Models are zero-centered and normalized into the unit sphere. Since most 3D models in this dataset are not watertight, we use ManifoldPlus (Huang et al., ) to make all meshes watertight. Due to failures of modeling loading and conversion, we obtained around 600k watertight models for training. The three levels of latents are $128\times 64$, $512\times 32$, and $2048\times 16$ (where 64, 32, and 16 are channels of the latents). Some other hyperparameters of the network can also be found in Table 3. We name the model as LaGeM-Objaverse. We also apply the method to ShapeNet, where the train split is taken from (Zhang et al., ). Since ShapeNet is a relatively small and easy dataset compared to Objaverse, we choose smaller latents which are $32\times 32$, $128\times 16$, and $512\times 8$. The model is named as LaGeM-ShapeNet. Both models are compared against VecSet (Zhang et al., ). We use Chamfer distance and F-score as the metrics. The results are shown in Table 4. Like (Zhang et al., ), we first compare the results on the largest categories (which have several thousand training samples) in ShapeNet and then all categories. We can see that, LaGeM-ShapeNet has almost the same number of parameters as VecSet, but with much shorter training time and less training memory. The quantitative results (averaged over all ShapeNet categories) are also better than VecSet’s. While for LaGeM-Objaverse, there is a large improvement in both training cost and quantitative results. The quantitative results show an improvement of almost 50 percent averaged across the complete dataset in terms of the metric Chamfer. This demonstrates that LaGeM-Objaverse has good generalization ability. This can also be seen in Fig. 7. The results of LaGeM-Objaverse are good on small categories of ShapeNet. In previous works (Zhang et al., ), this is nearly impossible because of limited training samples.

**Table 4: Evaluation on ShapeNet. We compare our results to VecSet (Zhang et al., ) trained on ShapeNet. If we train our model on ShapeNet and evaluate on ShapeNet our model is slightly better than VecSet. When our model is trained on Objaverse and evaluated on ShapeNet, we can see a very large improvement. Note that it is difficult to scale VecSet to Objaverse training.**
|  | Chamfer $\downarrow(\times 100)$ | F-Score $\uparrow(\times 100)$ |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | VS | LaGeM($\Delta$) | VS | LaGeM($\Delta$) |  |  |  |  |  |  |
|  | ShapeNet | Objaverse | ShapeNet | Objaverse |  |  |  |  |  |  |
| table | 2.46 | 2.48 | 0.02 | 2.09 | -0.37 | 99.94 | 99.97 | 0.02 | 99.96 | 0.02 |
| car | 5.99 | 5.89 | -0.10 | 4.36 | -1.63 | 89.85 | 90.31 | 0.46 | 92.15 | 2.30 |
| chair | 2.92 | 2.89 | -0.03 | 2.01 | -0.91 | 96.40 | 96.49 | 0.09 | 99.91 | 3.51 |
| airplane | 1.78 | 1.81 | 0.03 | 1.58 | -0.21 | 99.50 | 99.48 | -0.02 | 99.78 | 0.29 |
| sofa | 2.64 | 2.63 | -0.01 | 2.25 | -0.39 | 98.92 | 99.04 | 0.11 | 99.60 | 0.67 |
| rifle | 1.78 | 1.77 | -0.01 | 1.44 | -0.34 | 99.88 | 99.88 | -0.01 | 99.94 | 0.06 |
| lamp | 4.36 | 4.44 | 0.08 | 2.37 | -2.00 | 96.78 | 97.18 | 0.39 | 99.43 | 2.64 |
| mean (selected) | 3.13 | 3.13 | 0.00 | 2.30 | -0.83 | 97.33 | 97.48 | 0.15 | 98.68 | 1.36 |
| mean (all) | 4.68 | 4.63 | -0.04 | 2.42 | -2.26 | 93.25 | 93.47 | 0.23 | 98.93 | 5.68 |

**Table 5: Generalization on Various Datasets. Our trained model is capable of doing inference on several existing datasets. It can be applied on non-watertight datasets like ABO and pix3d even the model is trained on watertight datasets. Note that models from ShapeNet are not watertight originally. We use the watertight version processed by (Zhang et al., ). The metric for ShapeNet-test is different from Table 4. It is because here we show metrics averaged over all objects instead of categories.**
|  |  |  | Chamfer $\downarrow(\times 100)$ | F-Score $\uparrow(\times 100)$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset | # Meshes | Manifold | VS | LaGeM($\Delta$) | VS | LaGeM($\Delta$) |  |  |
| Thingi10k (Zhou & Jacobson, ) | 10k | Yes | 4.52 | 2.99 | -1.53 | 92.75 | 97.19 | 4.44 |
| ABO (Collins et al., ) | 8k | No | 4.91 | 3.66 | -1.26 | 92.52 | 94.91 | 2.39 |
| ShapeNet (Chang et al., )-test | 2k | Yes | 3.25 | 2.33 | -0.92 | 97.41 | 99.49 | 2.08 |
| EGAD (Morrison et al., ) | 2k | Yes | 3.27 | 2.82 | -0.45 | 99.02 | 99.76 | 0.74 |
| GSO (Downs et al., ) | 1k | Yes | 3.78 | 2.35 | -1.43 | 94.70 | 99.54 | 4.84 |
| pix3d (Sun et al., ) | 700 | No | 6.53 | 6.02 | -0.50 | 87.25 | 87.96 | 0.71 |
| FAUST (Bogo et al., ) | 100 | Yes | 2.10 | 1.31 | -0.79 | 99.58 | 99.90 | 0.32 |

Figure: Figure 7: Generalization on ShapeNet. Our results are better than VecSet in all categories. On small categories, the results of VecSet are not stable because of limited training samples. In contrast, our trained model also performs well in these categories.
Refer to caption: extracted/5895005/images/shapenet.png

Figure: Figure 8: Qualitative Results on ShapeNet. We show autoencoding results on ShapeNet. We use VecSet as the baseline. Our model is capable of reconstructing detailed geometry, especially thin structures.

Figure: Figure 9: Qualitative Results on Thingi10k. Our model can even preserve highly detailed geometry in CAD models.

Figure: Figure 10: Qualitative Results on FAUST and GSO. Results of VecSet are over-smoothed, while our method can preserve sharp details.

To further prove the generalization ability of LaGeM-Objaverse, we also test the autoencoding on various datasets, including Thingi10k (Zhou & Jacobson, ), ABO (Collins et al., ), EGAD (Morrison et al., ), GSO (Downs et al., ), pix3d (Sun et al., ) and FAUST (Bogo et al., ). The objects from these datasets vary from daily objects, CAD models, human models, and synthetic objects. The quantitative results can be found in Table 5. We again use VecSet’s model as the baseline. From the metrics, we can see that LaGeM-Objaverse is able to represent different kinds of objects with highly detailed geometry and sharp features. Note that, even for non-watertight meshes, the model is still able to do reconstruction. Visual results of the method can be found in Fig. 8, Fig. 9, Fig. 10.

### 4.2 Generative Model

We conducted two generative experiments, one is on ShapeNet with categories as the condition, and the other one is unconditional generation on Objaverse-10k.
For ShapeNet, the denoising networks of the 3 levels have 12 self-attention blocks with 768 channels. We trained the model for around 200 hours with 4 A100 GPUs. The results are shown in Fig. 11.
For Objaverse-10k, due to limited training GPU resources, we select a subset of 10k models from Objaverse and train the unconditional generative model. There are 24 self-attention blocks with 768 channels in all stages of the latents. The model is trained on 16 A100 GPUs for around 100 hours. See Fig. 12 for some unconditional generation results.

#### Controllability of the Latents.

We verify that different levels of latents control different levels of detail of the generated samples. During generation, we first generate higher-level latents ${\mathcal{Z}}_{3}$, which determine the main structures of the 3D models. Then we use ${\mathcal{Z}}_{3}$ as a condition to generate ${\mathcal{Z}}_{2}$, which adds major details to the models. In the end, we generate ${\mathcal{Z}}_{1}$ conditioned on both ${\mathcal{Z}}_{3}$ and ${\mathcal{Z}}_{2}$. This final step adds some minor details to the samples. A visual illustration can be found in Fig. 13.

Figure: Figure 11: Category-Conditioned Generative Results on ShapeNet.
Refer to caption: x3.jpg

Figure: Figure 12: Unconditional Generative Results on Objaverse-10k.
Refer to caption: x4.jpg

Figure: Figure 13: Latent Levels. Each small $4\times 4$ block shares the same level 3 latents ${\mathcal{Z}}_{3}$. 3D models in the same block have similar structures. In each block, every $1\times 4$ line shares the same level 2 latents ${\mathcal{Z}}_{2}$. In each line of a block, 3D models look almost the same except for some minor details. Thus, we argue that ${\mathcal{Z}}_{3}$ controls the *structure*, ${\mathcal{Z}}_{2}$ affects the *major details* and ${\mathcal{Z}}_{1}$ is responsible for *minor details*.

## 5 Conclusion

We proposed LaGeM (Large Geometry Model), an architecture for encoding 3D geometry. Different from previous approaches, the latent space is modeled as a hierarchical latent VecSets. To make this work, our model employs a U-Net-style design and a new regularization technique for the bottleneck. We showed that this model can be trained much faster with much lower GPU memory costs, especially for larger networks and datasets. This enables scaling of the network for large-scale datasets. We release our model trained on a 600k geometry dataset. Additionally, we proposed a cascaded diffusion model to show some preliminary generative results with the hierarchical latent space.

#### Limitation.

Since the latent space is divided into multiple levels, training a diffusion model on all levels still takes a lot of time. Our method does not solve the high training cost problem of diffusion itself.

## Appendix A Data preprocessing

The data preprocessing is based on (Zhang et al., ).

### A.1 Volume points sampling.

We sample volume points uniformly in the bounding sphere.

### A.2 Near points sampling

The near-surface points are obtained by sampling Gaussian-jittered surface points.

## Appendix B Data augmentations

#### Random axis scaling.

The augmentation is from (Zhang et al., ). We randomly sample a scaling factor for each axis which ranges from [0.75, 1.25].

#### Unit sphere normalization.

We normalize each mesh to a unit sphere, i.e., the max point norm of the point clouds is 1.

#### Random rotations.

We apply random rotations during the training of the autoencoder,

$$ ${\mathbf{R}}(\alpha,\beta,\gamma)=\begin{bmatrix}\cos\alpha&-\sin\alpha&0\\ \sin\alpha&\cos\alpha&0\\ 0&0&1\\ \end{bmatrix}\begin{bmatrix}\cos\beta&0&\sin\beta\\ 0&1&0\\ -\sin\beta&0&\cos\beta\\ \end{bmatrix}\begin{bmatrix}1&0&0\\ 0&\cos\gamma&-\sin\gamma\\ 0&\sin\gamma&\cos\gamma\\ \end{bmatrix},$ (13) $$

where $\alpha$, $\beta$, and $\gamma$ are yaw, pitch, and roll, respectively. Our meshes are firstly normalized into a unit sphere. Thus after the random rotations, the models will still be inside of a unit sphere.

## Appendix C Regularization

The proposed regularization (see Table 2) is implemented with layer normalization (PyTorch code).

## Appendix D Training time query points sampling

In the previous work (Zhang et al., ), the sampling strategy is uniformly sampling 1024 points in the bounding volume during training. We found this is not working on Objaverse. Since lots of meshes have very thin structures, this strategy will cause no inside points to be sampled during training. This heavily imbalenced data classficiation severely affects the occupancy loss.

We propose the following solution. In each iteration, we make sure half of the points have positive labels and the other half have negative labels.

## Appendix E Training loss

The loss is binary cross entropy as in previous work (Zhang et al., ). Formally, we have

$$ ${\mathcal{L}}=\mathbb{E}_{\mathbf{p}\in\mathbb{R}^{3}}\left[\mathrm{BCE}\left( \mathcal{\hat{O}}(\mathbf{p}),\mathcal{O}(\mathbf{p})\right)\right].$ (14) $$

In practice, we use the empircal loss

$$ $\mathbb{E}_{\mathbf{p}\in{\mathcal{Q}}^{\text{vol}}}\left[\mathrm{BCE}\left( \mathcal{\hat{O}}(\mathbf{p}),\mathcal{O}(\mathbf{p})\right)\right]+0.1\cdot \mathbb{E}_{\mathbf{p}\in{\mathcal{Q}}^{\text{near}}}\left[\mathrm{BCE}\left( \mathcal{\hat{O}}(\mathbf{p}),\mathcal{O}(\mathbf{p})\right)\right].$ (15) $$

Here, ${\mathcal{Q}}^{\text{vol}}$ is the set of volume query points, and ${\mathcal{Q}}^{\text{near}}$ is the set of near-surface query points.

## Appendix F Diffusion

We use the formulation EDM (Karras et al., ) for the diffusion models. The inference/sampling algorithm is also taken from the paper.

## Appendix G Latents analysis

We analyze how latents are affecting the final reconstruction. The latents are partially replaced by standard Gaussian noise (this is because our latents are also zero mean and unit variance). We show the visual results in Fig. 14.

Figure: Figure 14: Latent with red color ${\mathcal{Z}}$ means it is replaced by Gaussian noise. Latent with blue color ${\mathcal{Z}}$ means it is generated with the diffusion models.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: For the rebuttal, please refer to the Weaknesses section. Additionally, I have a few questions and suggestions:

- For the diffusion transformer, conducting cross-attention only once may not be sufficient.

- In Figure 1, there is no explanation for 'FtoL'.
soundness: 3
strengths: - The proposed method extends prior work VecSet to a hierarchical architecture, which improves generalization ability.

- The hierarchical autoencoder encodes the 3D shape into different levels of latent representations, with each level controlling different geometric details. This feature is highly beneficial for 3D generation.

- The writing in this paper is clean and easy to follow. The comparison with previous work (Table 1) is a valuable addition.
summary: This paper introduces a novel hierarchical autoencoder that compresses 3D models into a highly compact latent space, designed to handle large-scale datasets and support generative modeling using diffusion. Unlike previous approaches, the hierarchical autoencoder works on unordered sets of vectors, with each level controlling different geometric details. The model effectively represents a wide range of 3D models while preserving high-resolution details, and it reduces training time and memory usage compared to the baseline. Additionally, the authors propose a cascaded diffusion framework for generative modeling in the hierarchical latent space, allowing control over the level of detail in generated 3D models.
weaknesses: - The improvement in training time (by 0.7×) and memory consumption (by 0.58×) does not seem significant. In this case, having three levels of latent representations might be too heavy.

- The experiments presented in the paper involve up to 2K latent representations, which is not a substantial sequence length for Transformers with Flash Attention. Recent work [a] has scaled up to 64K latents.

- When using three levels of latent representations, we need three levels of diffusion models, which may lead to additional error accumulation. It would be worthwhile to mention how this issue is addressed for the proposed diffusion model.

- For the proposed regularization, it seems to force the datasets to share the same mean and standard deviation (Eq. 7), which could negatively impact model performance. In contrast, for KLD, very small values are typically used to avoid harming reconstruction. One suggestion is to try the method from [b] without applying any regularization.

[a] Meshtron: High-Fidelity, Artist-Like 3D Mesh Generation at Scale. https://openreview.net/forum?id=mhzDv7UAMu

[b] AutoDecoding Latent 3D Diffusion Models. NeurIPS 2023.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: For the experiments presented in Table 4, Table 5 and Figure 8, I wonder if both LaGeM and VecSet are using the same latent size and type of latent regularization? If not, the comparison could be unfair, as these differences could dominate the performance difference rather than the hierarchical architecture itself.
soundness: 2
strengths: * The proposed autoencoder significantly outperforms VecSet, the previous work, in terms of reconstruction quality and generalization capability.
* The paper demonstrates the controllability of individual level-of-detail, which is not possible with previous works due to having only a single level of latent vectors.
summary: The paper combines the idea of hierarchical VAE (Vahdat & Kautz, 2020) with VecSet (Zhang et al., 2023), leading to a 3D shape auto-encoder with hierarchical latent space. The proposed mode takes point clouds as input, and produces an occupancy field, which follows VecSet. It is more efficient than a single-level VecSet autoencoder, while being scaleable to large 3D datasets. The proposed auto-encoder outperforms VecSet in terms of reconstruction quality, while being more efficient. When coupled with a cascaded generative model (e.g. a cascaded diffusion model), this also enables control over the individual level of detail of the generated shapes.
weaknesses: * It is well-known that having KL divergence on the latent trades reconstruction quality for a smoother and more compact latent space. Such a well-behaved latent space could be helpful for the performance of the upstream generative model. I thus won't consider removing the KL loss "new regularization technique" (L477-478).
* The idea of controllability over different level-of-detail is interesting. However, according to Figure 13, its effect is very subtle and not particularly useful. This does increase the complexity and training cost of the diffusion model however, as stated in the Limitation section.
* The effect of using different latent regularizations (Table 2) is heavily discussed in Sec. 3.2 but not ablated. 
* Table 3: It would be clearer if each column had a header explaining their differences and giving the setting a name.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 3
questions: Please refer to the weaknesses above.
soundness: 3
strengths: 1. The paper introduces a novel hierarchical 3D autoencoder with faster training time and low memory consumption, addressing the expressiveness of SoTA method 3DShape2VecSet which is unable to scale to larger datasets.
2. Extensive experiments demonstrate the method's efficacy, outperforming previous state-of-the-art methods on key datasets.
3. The paper is generally well-written and easy to follow. The figures are helpful in illustrating hierarchical architecture.
summary: This paper introduces a novel 3D autoencoder called LAGEM, which maps 3D models into a highly compressed latent space. The key contribution of this paper is the hierarchical autoencoder architecture, which takes 0.70x the time and 0.58x the memory compared to the baseline. Experiments on Objaverse and ShapeNet demonstrate promising results.
weaknesses: 1. My main concern is whether having a better 3D autoencoder will lead to a better 3D generative model. On one hand, a better 3D autoencoder implies a higher upper bounder on the quality of the generated results. On the other hand, it requires that the latent space be smoother and easier to learn. Therefore, it would be even better if quantitative metrics for the 3D generated results could be provided.
2. Training a diffusion model on multi-levels takes a lot of training time. So is it possible to only train a single diffusion model to generate latent codes. And train a feed-forward network which take the latent code from the previous level as input to predict the latent code for the next level.
3. Is it possible to generalize the LaGeM to scene-level datasets like Matterport3D or the Replica dataset?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: The authors claim that it is difficult to scale VecSet to Objaverse training (#L328) and why?
soundness: 2
strengths: 1. The paper is clearly written and effectively expresses the motivation and role of hierarchical learning. 

2. The reporting of experimental hyperparameters is very detailed, enhancing the technical solidity of the paper. 

3. The model has been tested on a large number of datasets, validating the robustness of the proposed method.

4. Additionally, I enjoy Fig. 13.
summary: This paper introduces a novel hierarchical autoencoder that maps 3D models into a highly compressed latent space. The hierarchical autoencoder is specifically designed to tackle the challenges arising from large-scale datasets and generative modeling using diffusion. The paper also proposes a cascaded diffusion framework where each stage is conditioned on the previous stage based on the above hierarchical autoencoder.
weaknesses: The main weakness of this paper lies in the lack of comparison experiments. 

1. Regarding reconstruction accuracy, this work is only compared with one baseline (VecSet). In reality, there are more comparative options for reconstruction, including autoSDF[1], LION[2], and 3DQD[3]. 

2. Additionally, this work does not compare the quality of generated results with other baselines. The authors have listed several potential comparison baselines in Table 1. From the perspective of visualizing generated quality, the quality of the generated results does not show a significant advantage over many baselines listed in Table 1. 

3. There is also a lack of ablation experiments, such as the rationale behind using 3 levels of latent features.

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

[2] LION: Latent Point Diffusion Models for 3D Shape Generation

[3] 3DQD: Generalized Deep 3D Shape Prior via Part-Discretized Diffusion Process

Meta-review: The paper introduces a hierarchical autoencoder for 3D shapes that allows to train diffusion models in for generation in latent space.
The paper was well-received by all reviewers, converging to positive scores, recommending acceptance. The reviewers highlighted the good writing and extensive experiments, clearly showing benefits over previous methods. After rebuttal, remaining concerns regarding ablations and missing metrics have been resolved.

I agree with the reviewers and follow their suggestion with an accept recommendation.

Decision: Accept (Poster)