Paper: $\texttt{NAISR}$: A 3D Neural Additive Model for Interpretable Shape Representation (ICLR-2024)

Reviewer comment (Theory):
The paper assumes the given population to be unimodal and representable fully by a learned atlas. However, the methodology's ability to handle multimodal distributions and topological variations remains unclear.

Author response:
The current approach is not targeted at problems with changing topology as it is based on deforming an estimated template image. We will clarify this in the paper. This property might be a restriction in certain application contexts and could, for example, be overcome by allowing subject-specific modifications to the template level set function. However, note that in medical applications preserving topology is often desirable. Such applications are currently our main target. Our model currently also only considers one template shape. However, extensions to multiple templates, which would then be selected based on the best fit for a specific subject, are conceivable. Note though that while supporting multiple templates would be a more flexible approach for shape fitting it would at the same time complicate the construction of spatial correspondences between all shapes which is often a primary requirement in medical shape analysis. We will add related discussions to the paper.

Reviewer comment (Experiments):
The experiments do not adequately demonstrate the model's performance under limited training sample sizes, a typical scenario in medical shape analysis.

Author response:
We conducted an experiment to train $\texttt{NAISR}$ with only 50 airway shapes and found $\texttt{NAISR}$ fails to work on such a limited sample size with a median Chanfer distance of 0.7440 (0.0390 using the full training set).

Very limited sample sizes are likely not a good match for our approach as our goal is to not only capture a shape space for a population but also to capture variability with respect to given covariates. Hence, a sufficient number of observations across the desired covariate range is desirable. The smallest dataset we used (the pediatric airway dataset) contains 357 shapes, which is a medium-sized dataset. However, we argue that this dataset is challenging and representative for medical scenarios because 1) most subjects (n=229) only have one time observation; and 2) large shape and size variations can be observed as the dataset spans a large age range from infant to adult.

Our experiments show that the required sample sizes can be much smaller than the methods designed for general computer vision datasets, such as A-SDF, which we believe is a significant improvement. However, they cannot be too small based on the results using only 50 samples discussed above.

Reviewer comment (Presentation):
Authors provide a discussion of limitations in the supplementary material, however this should be part of the main paper.

Author response:
Thanks for the advice. We will include a short discussion on limitations (due to page constraints) in the main manuscript.

Reviewer comment (Evaluation):
The paper does not discuss the computational complexity and runtime for various processes, including training, inference, shape transfer, evolution, and disentanglement. Given that neural implicit representation works at the point coordinate level, understanding these metrics is crucial.

Author response:
Thank you for the suggestion. We now measured runtimes and have added these measures to our paper in S.3.1 (see the paragraph on computational runtime) in the supplementary material.

Reviewer comment (Reproducibility):
What are the computational and time requirements to reconstruct a full shape using the proposed implicit representation?

Author response:
The model is trained on an Nvidia GeForce RTX 3090 GPU approximately for 12 hours for 3000 epochs for the pediatric airway dataset, and for around 3.5 hours for 300 epochs for the ADNI hippocampus dataset. It takes around 130 seconds for 800 iterations and 30 seconds for 200 iterations to infer the latent code $\mathbf{z}$ and the covariates $\mathbf{c}$ respectively. However, we found 200 iterations should be sufficient to produce a reasonable reconstruction in most cases. Then, it takes approximately 30 seconds to sample the SDF field and apply Marching Cubes to extract the shape mesh. Regarding shape transfer, evolution, and disentanglement for a specific case, we first need to optimize the latent code $\mathbf{z}$ as in shape reconstruction. Once the latent code $\mathbf{z}$ is optimized or assigned (e.g., to $\mathbf{0}$ as a template), it will take around 30 seconds to produce a new mesh controlled by the covariate $\mathbf{c}$.

Reviewer comment (Theory):
The paper does not adequately justify the use of a non-amortized latent code, particularly the latent code encapsulating shape parameters not observed or captured by the covariate. The latent code is being optimized during training and inference as opposed to using a neural network to estimate such latent codes.

Author response:
Our method derives from DeepSDF, the pioneering approach in using a coordinate-based neural network to represent shapes, with an auto-decoder to infer the latent code $\mathbf{z}$ during the testing stage. The reason for using an auto-decoder instead of an encoder is that it is difficult to make sure the high-dimensional latent space produced by the encoder is condensed yet sufficient to represent realistic unseen shapes (as discussed in the supplementary material of the DeepSDF paper). A powerful encoder requires the model to have generative ability. However, in our case, we may only obtain several hundreds of airway shapes in the medical dataset, which makes it challenging to obtain a powerful encoder. Therefore, we choose to perform test time inference to obtain $\mathbf{z}$, which can prevent us from the risk of overfitting an auto-encoder while still retaining all of our desired capabilities of the model.

Reviewer comment (Novelty):
It is not clear how SD based techniques will be beneficial for medical images. Majority of medical imaging techniques use voxel-based representations. The proposed implicit shape representation (as other baselines) smoothed out details (compare ‘’gold standard’’ to all methods in Figure 2, for instance) and it is not clear why a method with such artifacts would be beneficial in practice. Authors should provide more clear motivation of why the proposed method will be useful for medical imaging.

Author response:
This is a good question. We are excited to share the background and motivation of our research problem which motivates $\texttt{NAISR}$.

* While medical images are volumetric, shape analysis based on shapes extracted from volumetric segmentations is an important task in medical image analysis to understand shape differences. These shapes are, of course, derived from the volumetric images. In our case, the objective is to capture normal airway shape development, w.r.t, age, and weight, as little is known about it, and understanding it better will be able to assess airway abnormalities. Think about it as the equivalent of a growth chart for children. Here we simply do not limit ourselves to height but want to capture shape change with age.

* In many shape analysis problems, acounting for covariates is important, because the shape population contains large variations and the covariates help with indicating which part of the shape population we should look at when analyzing new cases. Take our pediatric airway shape dataset as an example: if we want to characterize airway abnormalities for a 10-year-old child, airways of newborn babies and adults will be quite different. In our case, we use $\texttt{NAISR}$  to learn how the airway shape changes based on the covariates (e.g., age and weight).

We agree that there is a smoothing effect. However, we believe that this smoothing effect can be beneficial (and should therefore not be considered an artifact) based on two reasons:

* The jagged appearance of the gold standard shapes in Figure 2 is due to the voxel representation of the underlying segmentations. Hence, these are not representations of the true geometry as, for example, real hippocampi or airways are smooth. In fact, in classical medical shape analysis, shapes are commonly smoothed before analyses or represented by smooth basis function (for example, based on spherical harmonics). Our approach naturally provides a reasonable level of smoothing preserving key details (as can be seen in Figure 2) while reducing voxelization effects. Note also that in real clinical data voxel resolutions may differ between image acquisitions: this is for example the case for the airway shapes. Hence, removing these voxelization artifacts is desirable. It is true that strong geometric shape distortions might be problematic, but this is not what we observe. Ultimately, the success of these shape approximations should be judged in the context of a downstream task as part of follow-up work.

* The core question we are trying to answer with $\texttt{NAISR}$ is, how an average shape (atlas shape) for particular covariates looks like. This population average is expected to be smooth.

Reviewer comment (Theory):
As described in Section 3.1, it is assumed that “the overall displacement field is the sum of displacement fields that are controlled by individual parameters”, plus there is a contribution of the unknown covariate (z) which cannot be controlled. However, it isn’t clear what happens if a key covariate is unknown (e.g., not provided in metadata) or a not useful covariate is present. The authors should clearly explain implications of such assumptions.

Author response:
We thank the reviewer for posing this excellent question. This may happen in common clinical scenarios. There are two scenarios to be discussed here. **The first scenario is that the covariate is known at the training stage but missing at the test stage.** We discussed this in Sec.3.4, mainly in the Shape Reconstruction and Shape transfer part, but provide more details below which we will use to improve the discussion in the paper.

Suppose we already have a trained $\texttt{NAISR}$ model $\Phi$ and $\Phi$ has been trained well to represent population-level shape variations with respect to the covariates.

For a new shape $s_{tgt}$ with unknown covariates, we first need to recover its corresponding latent code $\mathbf{z}$ and the covariates $\mathbf{c}$. To estimate these quantities, the network parameters stay fixed and we optimize over the covariates $\mathbf{c}$ and the latent code $\mathbf{z}$ which are both randomly initialized. Specifically, we solve the optimization problem,

$$\hat{\mathbf{c}},\hat{\mathbf{z}}=\underset{\mathbf{c},\mathbf{z}}{\arg \min }~\mathcal{L}(\Phi, \mathbf{c}, \mathbf{z})$$

If the covariates  (e.g., age or weight at imaging time) are known we only infer the latent code $\mathbf{z}$ by optimizing 

$$\hat{\mathbf{z}}=\underset{\mathbf{z}}{\arg \min }~\mathcal{L}(\Phi,\mathbf{c}, \mathbf{z})\$$

A new patient shape with different covariates can then be generated by extracting the zero level set of $\Phi(\mathbf{p}, \mathbf{c}_{\mathrm{new}}, \hat{\mathbf{z}})$.

**The second scenario is where a key covariate is never recognized and thus not modeled as a covariate with $\texttt{NAISR}$.** In this case the deformation will purely be driven by the latent code $\mathbf{z}$. If we never capture an important covariate, $\texttt{NAISR}$ will gracefully  revert to a simple deformation-based shape representation model, from which we can still obtain shape reconstructions but which cannot be used to observe the missing covariates' effects of shapes.

Reviewer comment (Novelty):
Relationship to other methods (Mu 2021) is not clearly described. What is the methodological advantage that the proposed technique brings?

Author response:
We thank the reviewer for this question. We will update our manuscript to better clarify the motivation of our method. We apologize for not being able to include the full version of our research survey in the main text due to the 9-page limit. A more complete version of the research survey is provided in supplementary material S.1., including discussions of **deep implicit functions**, **point correspondences**, **disentangled representation learning**, **articulated shapes**, and **explainable artificial intelligence**, as well as how these methods relate to yet differ from our method.

Regarding the methodological advantage of $\texttt{NAISR}$ over other methods, we provide comparisons in Table 1 in the main text and discuss our experimental observations in Section 4.2 and in the supplementary material (S.3.3). Among all the comparison methods, only A-SDF (Mu 2021) is suitable for similar tasks (shape transfer and disentangled evolution) as our method.

However, A-SDF 1) does not consider deformations; 2) uses covariates as different inputs into one MLP which are therefore entangled together inside the network; 3) assumes that each articulation affects a separate object part while in medical scenarios, covariates often affect shapes in a more entangled and complex way; 4) is trained on a large longitudinal dataset, where each shape has at least 31 different articulations (e.g., a single door has 31 observations for different door opening angles.) On the other hand, $\texttt{NAISR}$ is a deformable shape representation that can work on smaller shape datasets and disentangle complex covariate interactions based on the neural additive structure of the deformation model.

The goal of the work by Wei et al. is semantic editing of geometric shapes. For example one might have an input chair represented as a point cloud for which we want to adjust the height of the back. To be able to do such semantic editing operations requires a set of defined shape attributes that correspond to our covariates. An encoder is then trained to predict shape attributes from a given point cloud. The decoder to go from an attribute vector to a shape is not learned however, but instead the geometric deformation model is assumed known and applied to a known template shape, termed an analytical decoder. Hence, the approach easily allows synthesizing synthetic shapes for different attribute vectors. As it is assumed known how a given shape is deformed by a set of attributes it is possible to create deformed shapes for two different attribute vectors and then to directly infer a deformation field between the two shapes (as both shapes come from the same template and we therefore know the point-to-point correspondences).

Transferring such an attribute change to a real, non-synthetic input point cloud then corresponds to predicting the attributes for this shape. Two synthetic shapes and their corresponding deformation are then created based on the predicted attribute vector and the edited predicted attribute vector (for example, by adjusting the attribute for chair back height). The corresponding deformation is then applied to the input shape. While our approach has conceptual similarities it is ultimately significantly different in the following ways: 1) it estimates a template shape and does not assume one; 2) it estimates the deformation models, i.e., how the covariates deform the template shape, rather than assuming it is known how the covariates affect the deformations; and 3) it is not based on simulated deformations.

The problem settings for the works by Mu et al. and Wei et al. do not carry over to typical scientific shape discovery scenarios because 1) they in general do not provide densely sampled longitudinal data for every patient; and 2) sample sizes are usually relatively small. For example, in our airway dataset, most kids only have one observation (Table.S.7). Our experiments show that A-SDF (Mu 2021) works exceptionally well at reconstructing shapes in the training set but does not work well for unseen airways in the testing set, indicating that A-SDF (Mu 2021) overfits the training set by memorizing shapes with their covariates. In contrast, $\texttt{NAISR}$ is designed for realistic scientific shape discovery problems and works for small and medium-sized datasets.

Reviewer comment (Novelty):
Related work is limited to representations focusing on implicit functions. There is other work in allowing parametrized shape editing, e.g., see for instance: Learning to Infer Semantic Parameters for 3D Shape Editing (Wei, 3DV 2020).

Author response:
We thank the reviewer for this question. We will update our manuscript to better clarify the motivation of our method. We apologize for not being able to include the full version of our research survey in the main text due to the 9-page limit. A more complete version of the research survey is provided in supplementary material S.1., including discussions of **deep implicit functions**, **point correspondences**, **disentangled representation learning**, **articulated shapes**, and **explainable artificial intelligence**, as well as how these methods relate to yet differ from our method.

Regarding the methodological advantage of $\texttt{NAISR}$ over other methods, we provide comparisons in Table 1 in the main text and discuss our experimental observations in Section 4.2 and in the supplementary material (S.3.3). Among all the comparison methods, only A-SDF (Mu 2021) is suitable for similar tasks (shape transfer and disentangled evolution) as our method.

However, A-SDF 1) does not consider deformations; 2) uses covariates as different inputs into one MLP which are therefore entangled together inside the network; 3) assumes that each articulation affects a separate object part while in medical scenarios, covariates often affect shapes in a more entangled and complex way; 4) is trained on a large longitudinal dataset, where each shape has at least 31 different articulations (e.g., a single door has 31 observations for different door opening angles.) On the other hand, $\texttt{NAISR}$ is a deformable shape representation that can work on smaller shape datasets and disentangle complex covariate interactions based on the neural additive structure of the deformation model.

The goal of the work by Wei et al. is semantic editing of geometric shapes. For example one might have an input chair represented as a point cloud for which we want to adjust the height of the back. To be able to do such semantic editing operations requires a set of defined shape attributes that correspond to our covariates. An encoder is then trained to predict shape attributes from a given point cloud. The decoder to go from an attribute vector to a shape is not learned however, but instead the geometric deformation model is assumed known and applied to a known template shape, termed an analytical decoder. Hence, the approach easily allows synthesizing synthetic shapes for different attribute vectors. As it is assumed known how a given shape is deformed by a set of attributes it is possible to create deformed shapes for two different attribute vectors and then to directly infer a deformation field between the two shapes (as both shapes come from the same template and we therefore know the point-to-point correspondences).

Transferring such an attribute change to a real, non-synthetic input point cloud then corresponds to predicting the attributes for this shape. Two synthetic shapes and their corresponding deformation are then created based on the predicted attribute vector and the edited predicted attribute vector (for example, by adjusting the attribute for chair back height). The corresponding deformation is then applied to the input shape. While our approach has conceptual similarities it is ultimately significantly different in the following ways: 1) it estimates a template shape and does not assume one; 2) it estimates the deformation models, i.e., how the covariates deform the template shape, rather than assuming it is known how the covariates affect the deformations; and 3) it is not based on simulated deformations.

The problem settings for the works by Mu et al. and Wei et al. do not carry over to typical scientific shape discovery scenarios because 1) they in general do not provide densely sampled longitudinal data for every patient; and 2) sample sizes are usually relatively small. For example, in our airway dataset, most kids only have one observation (Table.S.7). Our experiments show that A-SDF (Mu 2021) works exceptionally well at reconstructing shapes in the training set but does not work well for unseen airways in the testing set, indicating that A-SDF (Mu 2021) overfits the training set by memorizing shapes with their covariates. In contrast, $\texttt{NAISR}$ is designed for realistic scientific shape discovery problems and works for small and medium-sized datasets.

Reviewer comment (Experiments):
In section 4.2, authors describe experiments with shape reconstruction, and indicate that the proposed method and other SDF techniques can help complete shapes. It isn’t clear whether the technique is trained for shape completion, and if so, where do the partially complete shapes come from? Does the fact that a shape is partial negatively influence method training?

Author response:
As proposed in the DeepSDF paper, one advantage of implicit shape representations is that inference can be performed from an arbitrary number of SDF samples. Shape completion under such implicit frameworks amounts to solving for the shape code $\mathbf{z}$ that best explains a partial shape observation via Eq. (9) or Eq. (10). Given the shape code a complete shape can be rendered as the decoder implicitly captures the full shapes.

In our pediatric airway dataset, partial shapes exist during data collection due to differing fields of views during image acquisition. For example, some images contain the carina (the first branch of the airway) and others do not. Analyzing partial shapes is therefore inevitable in our research problem.

Our method is not trained specifically for shape completion. Instead, shape completion is an additional benefit we obtain from the implicit shape representation design. Different from voxel-based representations, an individual point is the input to the network, instead of a volume, meaning that it can be evaluated anywhere and extrapolation (within reason) is possible.

Reviewer comment (Presentation):
Some of the limitations (invertibility problem) should be mentioned in the main manuscript instead of the supplementary. Ending the conclusion of the paper with a reference to the supplementary does not seem like a good approach to me.

Author response:
Thanks for the suggestions. We will try our best to create some space and update our manuscript accordingly.

Reviewer comment (Experiments):
All the shapes in the dataset seem to be pre-aligned and centered. Is it a requirement that the shapes need to be pre-aligned and centered? In such a case, it should be clearly mentioned in Problem Description (Section 3.1)

Author response:
This is a good question. We mentioned in Supplementary S.2 that the hippocampus shapes are rigidly aligned using the ICP algorithm. The airway shapes are rigidly aligned using the anatomical landmarks. We now mention this in the Problem Description (Section 3.1). Thanks for the suggestion.

Reviewer comment (Reproducibility):
The construction details of the algorithm need to be improved.

Author response:
We are sorry about not placing enough construction details in the main text. In the supplementary material (S.3.1), we include the implementation details, such as parameter settings and data preprocessing steps. We now added more details about model construction. We now also added Figure.S.5 to the supplementary material for a clearer illustration of our approach. We will also make our implementation publicly available on *GitHub* so others can reproduce and build upon our results. See also our responses to the two specific construction questions below.

Reviewer comment (Reproducibility):
It would be better if the authors give more details of how to train the additive model $g_i$?

Author response:
Each subnetwork, including the template network $\mathcal{T}$ and the displacement networks $\{f_i\}$, are all parameterized with an $N_l$-layer MLP using $sine$ activations. We use $N_l=8$ for Starman and the ADNI hippocampus dataset; we use $N_l=6$ for the pediatric airway dataset. There are 256 hidden units in each layer. The architecture of the $\{f_i\}$ follows DeepSDF, in which a skip connection is used to concatenate the input of $(\mathbf{p}, c_i)$ to the input of the middle layer, as shown in the added Figure.S.5 in the supplementary material. We use a latent code $\mathbf{z}$ of dimension 256 ($L=256$). We follow SIREN for the network setting and parameter initialization.

For each training iteration, the number of points sampled from each shape is 750 ($N=750$), of which 500 are on-surface points ($N_{on} = 500$) and the others are off-surface points ($N_{off} = 250$). We train $\texttt{NAISR}$ for 3000 epochs for the airway dataset and for 300 epochs for the ADNI hippocampus and Starman datasets using Adam with a learning rate $5e-5$ and batch size of 64.

Also, we jointly optimize the latent code $\mathbf{z}$ with $\texttt{NAISR}$ using Adam with a learning rate of $1e-3$. We use a latent code $\mathbf{z}$ of dimension 256 ($L=256$). During training, $\lambda_1 =\lambda_5 = 1 \cdot 10^1$; $\lambda_2 = 3 \cdot 10^1$; $\lambda_3 = 1 \cdot 10^1$, $\lambda_4 = 1 \cdot 10^2$. For $\mathcal{L}_{lat}$, $\lambda_6 = \frac{2}{L}$; $\sigma=0.01$ (following DeepSDF). During inference, the latent codes are optimized for $N_t$ iterations with a learning rate of $5e-3$. $N_t$ is set to 800 for the pediatric airway dataset; $N_t$ is set to 200 for the Starman and ADNI Hippocampus datasets.

These details are included in S.3.1 in the supplementary material.

Reviewer comment (Experiments):
Does the inclusion of group structure contribute to the improvement of model performance? Please refer to the Group Sparse Additive Model (Yin et al 2012; Chen et al 2017).

Author response:
Thanks for the excellent suggestion of discussing group structure in covariates. We will include the suggested references and add a brief forward-looking discussion for future work. Note that our experiments are currently only based on very few covariates: two for the Starman dataset, three for the airway dataset, and four for the hippocampus dataset. In contrast, Yin et al.'s work targets 8,141 genes, and Chen et al. uses on the order of 20 covariates in their synthetic and real experiments. Exploring how group sparsity could be included in our NAISR model for high-dimensional covariates (for example, to explore the interaction between shape and gene expression on datasets that provide both datatypes) and how it would improve results would be extremely interesting future work.