NAISR: A 3D Neural Additive Model for Interpretable Shape Representation

TL;DR

NAISR tackles the lack of covariate-aware, interpretable 3D shape representations by introducing a Neural Additive Interpretable Shape Representation that deforms a learned atlas via covariate-driven displacement fields. It fuses deep implicit shape representations with an additive, disentangled deformation model to enable shape reconstruction, evolution, disentanglement, and transfer, while preserving interpretability. Evaluations on Starman, ADNI hippocampus, and pediatric airway demonstrate competitive reconstruction and superior real-data transfer and covariate analysis, outperforming baselines in key real-world scenarios. This framework supports scientific discovery and personalized predictions by making geometric changes attributable to individual covariates and enabling covariate-guided shape synthesis across longitudinal data.

Abstract

Deep implicit functions (DIFs) have emerged as a powerful paradigm for many computer vision tasks such as 3D shape reconstruction, generation, registration, completion, editing, and understanding. However, given a set of 3D shapes with associated covariates there is at present no shape representation method which allows to precisely represent the shapes while capturing the individual dependencies on each covariate. Such a method would be of high utility to researchers to discover knowledge hidden in a population of shapes. For scientific shape discovery, we propose a 3D Neural Additive Model for Interpretable Shape Representation ($\texttt{NAISR}$) which describes individual shapes by deforming a shape atlas in accordance to the effect of disentangled covariates. Our approach captures shape population trends and allows for patient-specific predictions through shape transfer. $\texttt{NAISR}$ is the first approach to combine the benefits of deep implicit shape representations with an atlas deforming according to specified covariates. We evaluate $\texttt{NAISR}$ with respect to shape reconstruction, shape disentanglement, shape evolution, and shape transfer on three datasets: 1) $\textit{Starman}$, a simulated 2D shape dataset; 2) the ADNI hippocampus 3D shape dataset; and 3) a pediatric airway 3D shape dataset. Our experiments demonstrate that $\textit{Starman}$ achieves excellent shape reconstruction performance while retaining interpretability. Our code is available at $\href{https://github.com/uncbiag/NAISR}{https://github.com/uncbiag/NAISR}$.

Figures (3)

• Figure 1: Neural Additive Implicit Shape Representation. During training we learn the template $\mathcal{T}$ and the multi-layer perceptrons (MLPs) $\{g_i\}$ predicting the covariate-wise displacement fields $\{\mathbf{d}_i\}$. The displacement fields are added to obtain the overall displacement field $\mathbf{d}$ defined in the target space; $\mathbf{d}$ provides the displacement between the deformed template shape $\mathcal{T}$ and the target shape. Specifically the template shape is queried not at its original coordinates $\mathbf{p}$, but at $\Tilde{\mathbf{p}}=\mathbf{p}+\mathbf{d}$ effectively spatially deforming the template. At test time we evaluate the trained MLPs for shape reconstruction, evolution, disentanglement, and shape transfer.
• Figure 2: Visualizations of airway and hippocampus reconstruction with different methods. The red and blue circles show the structure in the black circle from two different views. Hippocampus shapes are plotted with two $180^{\circ}$-flipped views. NAISR can produce detailed and accurate reconstructions as well as impute missing airway parts. More visualizations are available in Section \ref{['subsec.supp_shape_reconstruction']} of the supplementary material.
• Figure 3: Template shape extrapolation in covariate space using A-SDF and NAISR on three datasets. For the Starman shape extrapolations, the blue shapes are the groundtruth shapes and the red shapes are the reconstructions. The shapes in the middle white circles are the template shapes. The template shape is generated with zero latent code and is used to create a template covariate space. The shapes in the green and yellow boxes are plotted with $\{\Phi_i\}$, representing the disentangled shape evolutions along each covariate respectively. The purple shadows over the space indicate the covariate range that the dataset covers. Cyan points represent male and purple points female patients in the dataset. The points represent the covariates of all patients in the dataset. The colored shades at the boundary represent the covariate distributions stratified by sex. Example 3D shapes in the covariate space are visualized with their volumes ($cm^3$) below. NAISR is able to extrapolate the shapes in the covariate space given either an individualized latent code $\mathbf{z}$ or template latent code $\mathbf{0}$, whereas A-SDF struggles. The supplementary material provides more visualizations of individualized covariate spaces in Section \ref{['subsec:disentangled_shape_evolution']}. (Best viewed zoomed.)