11institutetext: School of Computation, Information and Technology, Technical
University of Munich, Germany 22institutetext: Munich Center for Machine
Learning, Germany 33institutetext: Institute of Machine Learning in Biomedical
Imaging, Helmholtz Munich, Germany 44institutetext: School of Biomedical
Engineering and Imaging Sciences, King’s College London, UK 55institutetext:
School of Medicine & Health, Klinkum Rechts der Isar, Technical University of
Munich, Germany 66institutetext: German Cancer Consortium (DKTK), Partner Site
Munich, Germany
66email<EMAIL_ADDRESS>
# Data-Driven Tissue- and Subject-Specific Elastic Regularization for Medical
Image Registration

Anna Reithmeir 112233 Lina Felsner 1133 Rickmer Braren 5566 Julia A. Schnabel
11223344 Veronika A. Zimmer 113355

###### Abstract

Physics-inspired regularization is desired for intra-patient image
registration since it can effectively capture the biomechanical
characteristics of anatomical structures. However, a major challenge lies in
the reliance on physical parameters: Parameter estimations vary widely across
the literature, and the physical properties themselves are inherently subject-
specific.

In this work, we introduce a novel data-driven method that leverages
hypernetworks to learn the tissue-dependent elasticity parameters of an
elastic regularizer. Notably, our approach facilitates the estimation of
patient-specific parameters without the need to retrain the network.

We evaluate our method on three publicly available 2D and 3D lung CT and
cardiac MR datasets. We find that with our proposed subject-specific tissue-
dependent regularization, a higher registration quality is achieved across all
datasets compared to using a global regularizer. The code is available at
https://github.com/compai-lab/2024-miccai-reithmeir.

###### Keywords:

Spatially Adaptive Regularization Hypernetworks.

##  1 Introduction

Accurate and reliable intra-patient registration of medical images is
ubiquitous for image-based diagnostics and treatment, for example, in
radiotherapy and disease monitoring [19]. While most state-of-the-art deep
learning-based registration algorithms resort to generic and global
regularization methods (regularizers) [9, 2, 14], physics-inspired
regularizers, such as the linear elastic regularizer [4], are particularly
well-suited for intra-patient registration. They model biomechanical
principles and thus ensure the physical plausibility of the registration.
First, the selection of physical parameter values is ambiguous. Prior works
often resort to estimations found across the literature [15, 13] even though
they can differ significantly not only in their absolute values but also in
their ratio [8, 20]. Second, different anatomical structures have different
physical properties. For instance, bones are nearly rigid, while soft tissue
is more elastic. A global regularizer can not accommodate such varying
properties, and thus, a spatially adaptive regularizer that can vary within
the image is required. However, spatial variability further complicates the
choice of physical parameter values, and precise knowledge about anatomical
structure locations is necessary. Third, physical properties can vary among
individuals. On the one hand, age is a factor as, e.g., with increasing age,
the elasticity of bones and arteries tends to decrease [23, 5]. On the other
hand, the presence of pathologies can impact tissue properties. For example,
increased stiffness is observed in fibrotic liver and lung tissue [20, 7] as
well as in cardiac tissue of patients with obesity and atherosclerosis [5].

Figure 1: Overview of our method. We train a globally and spatially adaptive
network. After training, the optimal tissue-specific elasticity parameters are
estimated with the global network. The spatially adaptive network then
predicts the deformation field for the registration of a new image pair that
follows the physical properties specified by the parameter values.

These challenges hinder the use of a one-suits-all set of physical parameter
values for the physics-inspired regularization of image registration.
Recently, works have been introduced that learn the elastic regularization in
the context of image registration [1, 17]. Arratia López et. al. [1] adopt a
hyperelastic regularizer to promote incompressibility in the registration
process. While the regularizer varies between the tissue in focus and the
background, it does not distinguish between different tissue types. Since this
method is based on a physics-informed neural network (PINN), retraining is
necessary for every image pair. Reithmeir et. al. [17] use a hypernetwork to
learn the elasticity parameters of the linear elastic regularizer for image
registration. A global regularizer is applied and tissue-specific properties
are not considered. Neither of the two methods above learns patient-specific
regularization parameters. The recent work of Ragoza et. al. [16] focuses on
the elasticity parameter estimation from medical images. They propose a PINN
for spatially varying tissue elasticity reconstruction from images [16].
However, it requires a specific medical imaging modality (MR elastography) and
cannot be applied to standard MR and CT images. Here, image registration is
not considered.

In this work, we build upon our previous work in [17] and introduce a novel
framework that not only learns tissue- but also subject-specific parameters
for the elastic regularization of image registration. To this end, we employ
two hypernetworks: One for the parameter estimation and one for the
registration with the spatially varying elastic regularization (Fig. 1). Our
main contributions are twofold:

  * •

We propose a data-driven spatially adaptive elastic regularization that varies
the elasticity properties for different anatomical structures.

  * •

To the best of our knowledge, we are the first to introduce a registration
algorithm that learns the subject-specific linear elastic regularization where
the elasticity parameters of the linear elastic regularizer are estimated for
each image pair individually without the need to retrain the network.

##  2 Methods

We propose a novel method that learns the spatially adaptive linear elastic
regularization from the data. It comprises two independent neural networks
that differ in their regularization scheme and which we will refer to as
global and spatially adaptive network in the following (Fig. 1 (i)). The
global network employs a global linear elastic regularizer and is used to
extract the tissue-dependent elasticity parameters. The spatially adaptive
network uses a spatially varying regularizer and is used for the regularized
registration.

For the estimation of optimal tissue- and patient-specific elasticity
parameters, we exploit the ability of hypernetworks [6] to adapt to
hyperparameter values at test time. Hypernetworks are small networks that can
learn the effect of hyperparameters on the output of a larger network. In this
work, we apply two hypernetworks that learn the effect of the elasticity
parameters of the linear elastic regularizer [4] on the registration. This
allows us to explore the regularization parameter space for new data at test
time without the need to retrain the networks.

Our method consists of three stages (see Fig. 1), which are described in
detail below: (i) Training the two networks, (ii) estimation of the optimal
subject- and tissue-specific elasticity parameters, and (iii) registration of
new data with the identified parameters.

###  2.1 Training of Global and Spatially Adaptive Networks

We employ two networks to learn the globally and spatially adaptive
regularized registration, respectively. Each consists of a U-Net [18]
registration network and a hypernetwork [6] and predicts the deformation
$\phi:\mathbbm{R}^{D}\rightarrow\mathbbm{R}^{D}$ between a moving image
$M:\Omega\subset\mathbbm{R}^{D}\rightarrow\mathbbm{R}$ and a fixed image
$F:\Omega\subset\mathbbm{R}^{D}\rightarrow\mathbbm{R}$. We parameterize $\phi$
as the dense displacement field $\phi(x)=x+\mathbf{u}(x)$ for each pixel
location $x$.

The inputs to both models are the image pair $(M,F)$ and the Lamé parameters
of the linear elastic regularizer [4] that control the elasticity properties
that are modeled.

####  2.1.1 Global Model:

This network is similar to [17] and based on [11]. It applies a global linear
elastic regularizer, i.e., it enforces the same elasticity properties
everywhere in the image domain $\Omega$. The two global Lamé parameters
$\lambda,\mu\in\mathbbm{R}$ are given as inputs to the hypernetwork. The
network is trained unsupervised, using the loss function in [17], with the
change of constraining $\lambda,\mu\in[0,1]$.

####  2.1.2 Spatially adaptive Model:

This network is an extension of [17] with a spatially varying linear elastic
regularizer. It adapts to different elasticity properties based on the
underlying tissue type at each $x\in\Omega$. This is achieved by tissue-
specific Lamé parameters $\\{(\lambda_{i},\mu_{i})|\ \forall i\in 1,...,C\\}$
with $C$ as the number of considered tissue classes.

Two elasticity parameter maps
$\mathbf{\Lambda},\mathbf{\Gamma}:\Omega\in\mathbbm{R}^{D}\rightarrow\mathbbm{R}$
are formed with the help of the segmentation map
$\mathbf{S}:\Omega\subset\mathbbm{R}^{D}\rightarrow\mathbbm{Z}$ of the fixed
image $F$. They hold the respective tissue-specific parameters at each
$x\in\Omega$, i.e., $\mathbf{\Lambda}(x)=\lambda_{c}$ and
$\mathbf{\Gamma}(x)=\mu_{c}$ for $\mathbf{S}(x)=c$. The training loss function
of the network is

$\begin{split}\mathcal{L}(M,F,\phi,\mathbf{\Lambda},\mathbf{\Gamma})=&(2*\mathbb{1}-\mathbf{\Lambda}-\mathbf{\Gamma})\
NCC(F,\phi\circ M)\\\
&+\int_{\Omega}\frac{1}{4}\mathbf{\Gamma}\sum_{i,j=1}^{D}(\partial_{x_{i}}u_{j}+\partial_{x_{j}}u_{i})^{2}+\frac{1}{2}\mathbf{\Lambda}(\mathrm{div\
}\mathbf{u})^{2}\ d\mathbf{x}\enspace,\end{split}$ (1)

where $\mathbb{1}:\Omega\subset\mathbbm{R}^{D}\rightarrow\mathbbm{R}$ is a
constant weighting map of ones, NCC is the local normalized cross-correlation
similarity metric, $\partial_{x}u$ are the partial derivatives of $\mathbf{u}$
and div is the divergence. To balance the regularization and the similarity
terms, we require that $0\leq\lambda(x),\mu(x)\leq 1\forall x\in\Omega$.

For simplicity, we omit an additional regularization weight which controls the
amount of regularization and absorb it in the elasticity parameters as in
[17]. During training, the elasticity parameters are individually sampled from
the uniform distribution $\mathcal{U}(0,1)$ in each training step. The two
networks are trained separately on the same training data.

###  2.2 Elasticity Parameter Identification

Hypernetworks allow the efficient optimization of hyperparameters at test
time. We exploit this advantage and use the trained global model to estimate
the optimal elasticity parameters of the regularizer. Notably, we do this not
only for each tissue type but also for each image pair individually. Thus, we
obtain tissue- and subject-specific elasticity parameters, which, to the best
of our knowledge, is not performed in prior works. Similar to [17], we
identify the optimal parameters by a grid search over the parameter space at
test time. Without increasing the computational complexity: (i) we perform a
sample-based grid search to identify the optimal parameters for each subject
individually, and (ii) we base the grid search on a class-wise heuristic to
find the optimal parameters for each tissue type. In this work, we use the
class-wise Dice score.

Overall, for each test sample $n$, the grid search results in the parameter
set $\\{(\lambda_{i}^{n},\mu_{i}^{n})|\ \forall i\in 1,...,C\\}$. For $N$ test
samples, in total $2CN$ parameters are estimated.

###  2.3 Registration with Spatially Adaptive Elastic Regularization

For the registration of an unseen test sample, the optimal Lamé parameters
that have been identified for this specific sample with the global model are
given as input to the hypernetwork of the spatially adaptive model. A single
forward pass through the trained spatially adaptive network then registers the
given image pair. Since this network is trained with a spatially varying
regularizer, the predicted deformation field is regularized in the different
anatomical regions according to the specified local elasticity properties.

##  3 Experiments and Results

###  3.1 Implementation Details and Datasets

We implement our method with TensorFlow 2.12.1 and use the default network
architecture of HyperMorph [11]. We use the Adam optimizer [12] and a learning
rate of 1e-4. All networks are trained until convergence using an NVIDIA
RTXA6000 GPU. The linear elastic regularizer is implemented with forward
finite differences. The code will be made publicly available upon acceptance.

We evaluate our method on three intra-patient datasets in 2D and 3D; (I) NLST:
The lung CT data of the National Lung Screening Trial dataset provided by the
Learn2Reg challenge [21, 9]. We use the middle coronal slice of each of the
209 annotated images and a train/val/test split of 169/10/30. We clip the
intensities to (-980, 1518) HU. (II) ACDC: The 2D cardiac MRI dataset is taken
from the Automatic Cardiac Diagnosis Challenge (ACDC) [3]. We use the middle
slices of the end-systolic and end-diastolic frames, a split of 50/20/30 and
we crop the images to size (128,128). All images are resampled to isotropic
pixel size, and their intensities are normalized. (III) L2R-Lung: The 3D lung
CT dataset is taken from the Learn2Reg challenge [10] and contains 30
inhale/exhale image pairs. The fixed images are only partly visible. We use a
split of 20/4/4, clip the intensities to (-980, 600) HU and padded the images
to size (192,128,192).

In addition to the already provided lung segmentation masks in the NLST and
L2R-Lung datasets, we generate liver and bone (vertebrae, ribs, and sternum)
masks with the TotalSegmentator tool [22]. The ACDC dataset provides
segmentations for the right ventricle, myocardium, and left ventricle blood
pool.

###  3.2 Analysis of Subject-Specific Elasticity Parameters

In the first experiment, we analyze the sample-specific elasticity parameters
that are estimated with the trained global network as described in Subsection
2.2. We use the class-wise Dice score as the grid search optimization
heuristic and extract the optimal Lamé parameters for each tissue type and
subject individually. The results for the NLST dataset are shown in Fig. 2. It
can be seen that the values for both parameters ($\lambda,\mu$) vary between
the subjects. Within each tissue class, a tendency towards low values for
$\mu$ and $\lambda$ can be observed. However, for each tissue type, multiple
samples can be found for which high values are optimal. Furthermore, the value
distribution differs between the tissue types. For example, the bone-specific
$\mu$ generally favors values that are smaller than $0.5$, while the liver-
specific values stretch out over the full parameter space. Similar results are
obtained for the other two datasets (see Appendix A).

Figure 2: The subject-specific tissue-wise regularization parameters
$\lambda,\mu$ that have been estimated with the global network for the NLST
dataset. Each blue dot corresponds to one test sample. The results show that
the optimal parameters vary between tissue types and subjects.

###  3.3 Comparison of Regularization Methods

In the second experiment, we evaluate our proposed method. To this end, we
compare our subject-specific spatially varying regularization (spatially
adaptive models) to the global regularization (global models). The global
models are evaluated with the estimated optimal global regularization
parameters according to the mean Dice score. The spatially adaptive models are
evaluated with the estimated optimal tissue-specific regularization parameters
according to the class-wise Dice scores. Furthermore, for better comparison,
we estimate tissue- and subject-specific regularization weights in the same
manner for the diffusion regularizer of the HyperMorph model [11]. Here, the
regularization weights specify the amount of smoothness per tissue type. As
evaluation metrics, we use the mean and class-wise Dice scores (DSC) and the
95th percentile Hausdorff distance (HD95) between the fixed and predicted
warped segmentations, the target registration error (TRE, Euclidean distance)
between the fixed and predicted warped keypoints and the fraction of negative
Jacobian determinants of the displacement field ($\%negJ$).

The quantitative results are shown in Tab. 1. The results for the diffusion
regularizer, in general, outperform the results for the linear elastic
regularizer. Similar observations are made in [17]. In terms of the class-wise
DSC, all spatially adaptive models outperform the global models in at least
two out of three tissue classes. For the ACDC dataset and the elastic
regularizer, all three class-wise DSCs show higher performance with the
spatially adaptive regularization, even though this is the only case where a
lower HD95 is obtained. Furthermore, the mean DSC is consistently higher for
the spatially adaptive models. Interestingly, for the 2D datasets, the amount
of folding, i.e., the fraction of negative Jacobian determinants, increases
when spatially adaptive regularization is applied, while it decreases in the
3D case. Visual inspection of the results reveals that for some test samples,
a high amount of folding is observed inside the lung and background (compare
second-to-top row of Fig.3). Note, however, that nearly no folding is observed
here in the liver and bone structures. The qualitative results in Fig. 3
furthermore confirm the superior registration quality for the proposed method:
The breathing motion of the lower lung lobes (NLST dataset, second-to-top row)
and the omnidirectional expansion of the two heart ventricles (ACDC dataset,
bottom row) is better reflected in the predicted deformations as can be seen
in the deformation plots and difference images. Further qualitative examples
can be found in Appendix B.

Table 1: Quantitative results for the linear elastic (elas) and diffusion (diff) regularized models. The models with global (g) and spatially adaptive and subject-specific (sas) regularization parameters are compared. The results are shown as mean(standard deviation). For the class DSC the classes correspond to [lung, bone, liver] for the NLST/L2R-Lung and [right ventricle, myocardium, left ventricle blood pool] for the ACDC dataset. The results show that with the tissue- and subject-specific regularization, higher registration quality is achieved. For metric descriptions see Sec. 3.3 2D | Model | Class DSC$\uparrow$ | Mean DSC$\uparrow$ | HD95$\downarrow$ | %negJ$\downarrow$ |
---|---|---|---|---|---|---
NLST | initial | [0.914, 0.531, 0.775] | 0.749(0.064) | 11.755(10.328) | - |
diff-g [11] |  [0.972, 0.799, 0.855] | 0.885(0.067) | 8.227(10.293) |  0.022(0.007) |
diff-sas (ours) |  [0.972, 0.801, 0.855] |  0.886(0.065) |  8.157(9.998) | 0.038(0.028) |
elas-g [17] |  [0.969, 0.784, 0.855] | 0.879(0.068) | 9.211(10.636) |  0.009(0.004) |
elas-sas (ours) |  [0.970, 0.796, 0.852] |  0.883(0.067) |  8.598(10.621) | 0.040(0.025) |
ACDC | initial | [0.748, 0.520, 0.667] | 0.645(0.119) | 7.714(3.162) | - |
diff-g [11] |  [0.897, 0.757, 0.888] | 0.847(0.052) | 2.696(1.776) |  0.006(0.005) |
diff-sas (ours) |  [0.903, 0.771, 0.888] |  0.854(0.056) |  2.659(1.967) | 0.011(0.008) |
elas-g [17] | [0.886, 0.735, 0.867] | 0.829(0.066) |  3.139(1.682) |  0.002(0.002) |
elas-sas (ours) |  [0.888, 0.743, 0.869] |  0.833(0.068) | 3.175(2.010) | 0.009(0.004) |
3D | Model | Class DSC$\uparrow$ | Mean DSC$\uparrow$ | TRE$\downarrow$/HD95$\downarrow$ | %negJ$\downarrow$ |
L2R-Lung | initial | [0.911, 0.494, 0.758] | 0.721(0.061) | 3.227(0.734)/16.092(7.328) | - |
diff-g [11] | [0.959, 0.693, 0.829] | 0.827(0.023) | 3.041(0.551)/12.865(7.452) | 0.030(0.014) |
diff-sas (ours) |  [0.960, 0.695, 0.839] |  0.832(0.022) |  2.568(0.532)/12.481(7.579) |  0.027(0.010) |
elas-g [17] |  [0.961, 0.678, 0.822] | 0.820(0.022) | 2.645(0.484)/13.330(7.942) | 0.013(0.005) |
elas-sas (ours) |  [0.959,0.684,0.842] |  0.828(0.024) |  2.638(0.512)/12.499(6.786) |  0.006(0.002) |
Figure 3: Qualitative results for the NLST (top) and ACDC (bottom) datasets.
The left column shows the moving (M) and fixed (F) images. From left to right:
Warped moving image, predicted deformation, Jacobian determinant (blue:
positive values, red: negative values), difference image between fixed and
warped image, segmentation map boundaries. The results show that the subject-
specific spatially adaptive regularization (elas-sas) leads to more plausible
deformation than the global regularization (elas-g).

##  4 Discussion and Conclusion

In this work, we proposed a novel method that estimates tissue- and subject-
specific parameters for the physics-inspired regularization of data-driven
image registration. We focused on the parameters of the linear elastic
regularizer and evaluated the method on three diverse datasets.

Our analysis of the subject-specific elasticity parameters that have been
estimated with the global model in Sec.3.2 confirms that the parameters can
differ not only between tissue types but also between individual subjects.
This highlights the importance of the adaptation of physics-inspired
regularizers to the local anatomical structures and to the individual
patients. The comparison of the different regularization methods in Sec.3.3
shows that physiologically more meaningful deformations are obtained with
spatially adaptive, subject-specific linear elastic regularization. While our
method results in more folding, we observe consistently increased class Dice
scores for the bones in Tab.1, which indicates that small anatomical
structures are better registered with our method than with a globally
regularized model.

A current limitation of this work is that the estimated elasticity parameters
cannot be interpreted in a physically meaningful manner. The parameter $\mu$
is the shear modulus and serves as a measure of stiffness: Rigid materials
have higher values for $\mu$ than more elastic materials. This, however, is
not reflected in our estimations since, e.g., the bone-specific $\mu$ do not
favor larger values compared to the liver- or lung-specific $\mu$. A possible
reason for this could be the uniform sampling space of the parameters in our
method, and we plan to investigate this further in the future. Also, we aim to
extend the evaluation of our method to additional datasets and more tissue
types.

In conclusion, we proposed a novel image registration method that learns
tissue- and sample-specific elasticity parameters for physics-inspired
regularization. We believe that the patient-specific adaptation of the
regularization is fundamental for successful intra-patient registration and
that it can further improve data-driven registration methods in the future.

## References

  * [1] Arratia López, P., Mella, H., Uribe, S., Hurtado, D.E., Sahli Costabal, F.: Warppinn: Cine-mr image registration with physics-informed neural networks. Medical Image Analysis 89, 102925 (2023)
  * [2] Balakrishnan, G., Zhao, A., Sabuncu, M., Guttag, J., Dalca, A.: Voxelmorph: A learning framework for deformable medical image registration. IEEE Transactions on Medical Imaging 38, 1788–1800 (2019)
  * [3] Bernard, O., Lalande, A., Zotti, C., Cervenansky, F., Yang, X., Heng, P.A., Cetin, I., Lekadir, K., Camara, O., Ballester, M.A.G., et al.: Deep learning techniques for automatic mri cardiac multi-structures segmentation and diagnosis: is the problem solved? IEEE transactions on medical imaging 37(11), 2514–2525 (2018)
  * [4] Broit, C.: Optimal Registration of Deformed Images. Graduate School of Arts and Sciences, University of Pennsylvania (1981)
  * [5] Cocciolone, A.J., Hawes, J.Z., Staiculescu, M.C., Johnson, E.O., Murshed, M., Wagenseil, J.E.: Elastin, arterial mechanics, and cardiovascular disease. American Journal of Physiology-Heart and Circulatory Physiology 315(2), H189–H205 (2018)
  * [6] Ha, D., Dai, A.M., Le, Q.V.: Hypernetworks. arXiv preprint arXiv:1609.09106 (2017)
  * [7] Haak, A.J., Tan, Q., Tschumperlin, D.J.: Matrix biomechanics and dynamics in pulmonary fibrosis. Matrix Biology 73, 64–76 (2018)
  * [8] Hagemann, A., Rohr, K., Stiehl, H., Spetzger, U., Gilsbach, J.: Biomechanical modeling of the human head for physically based, nonrigid image registration. IEEE Transactions on Medical Imaging 18(10), 875–884 (1999)
  * [9] Hering, A., Hansen, L., Mok, T.C.W., et al.: Learn2reg: Comprehensive multi-task medical image registration challenge, dataset and evaluation in the era of deep learning. IEEE Transactions on Medical Imaging 42(3), 697–712 (2023)
  * [10] Hering, A., Murphy, K., van Ginneken, B.: Learn2reg challenge: Ct lung registration - training data [data set]. zenodo. http://doi.org/10.5281/zenodo.3835682 (2020)
  * [11] Hoopes, A., Hoffmann, M., Fischl, B., Guttag, J., Dalca, A.V.: Hypermorph: Amortized hyperparameter learning for image registration. Information Processing in Medical Imaging (IPMI), vol. 12729, p. 3–17 (2021)
  * [12] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 (2014)
  * [13] Min, Z., Baum, Z.M.C., Saeed, S.U., Emberton, M., Barratt, D.C., Taylor, Z.A., Hu, Y.: Non-rigid medical image registration using physics-informed neural networks. In: Information Processing in Medical Imaging (IPMI) (2023)
  * [14] Mok, T.C.W., Chung, A.C.S.: Large deformation diffeomorphic image registration with laplacian pyramid networks. Medical Image Computing and Computer Assisted Intervention (MICCAI), vol. 12263, pp. 211–221 (2020)
  * [15] Qin, C., Wang, S., Chen, C., Bai, W., Rueckert, D.: Generative myocardial motion tracking via latent space exploration with biomechanics-informed prior. Medical Image Analysis 83, 102682 (2023)
  * [16] Ragoza, M., Batmanghelich, K.: Physics-informed neural networks for tissue elasticity reconstruction in magnetic resonance elastography. In: Medical Image Computing and Computer Assisted Intervention (MICCAI). pp. 333–343 (2023)
  * [17] Reithmeir, A., Schnabel, J.A., Zimmer, V.A.: Learning physics-inspired regularization for medical image registration with hypernetworks. In: Medical Imaging 2024: Image Processing. International Society for Optics and Photonics, SPIE (2024 (accepted))
  * [18] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention (MICCAI). pp. 234–241 (2015)
  * [19] Rueckert, D., Schnabel, J.A.: Medical Image Registration, pp. 131–154. Springer, Berlin, Heidelberg (2011)
  * [20] Talwalkar, J.A.: Elastography for detecting hepatic fibrosis: options and considerations. Gastroenterology 135(1), 299–302 (2008)
  * [21] Team, N.L.S.T.R.: Data from the national lung screening trial (nlst) [data set]. the cancer imaging archive. (2013)
  * [22] Wasserthal, J., Breit, H.C., Meyer, M.T., Pradella, M., Hinck, D., Sauter, A.W., Heye, T., Boll, D.T., Cyriac, J., Yang, S., Bach, M., Segeroth, M.: Totalsegmentator: Robust segmentation of 104 anatomic structures in ct images. Radiology: Artificial Intelligence 5(5) (2023)
  * [23] Zioupos, P.: Ageing human bone: Factors affecting its biomechanical properties and the role of collagen. Journal of Biomaterials Applications 15(3), 187–229 (2001)

## Supplementary Material

### Appendix A

(a) ACDC dataset

(b) L2R-Lung dataset

Figure 4: The subject-specific tissue-wise regularization parameters
$\lambda,\mu$ for the ACDC and L2R-Lung datasets that are estimated with the
trained global model. In Fig. 4(a), the classes correspond to the right
ventricle (RV), myocardium (LV-Myo), and left ventricle blood pool (LV-BP).

### Appendix B

Figure 5: Further visual registration examples of the NLST (top) and ACDC
(bottom) test dataset. The left column shows the moving (M) and fixed (F)
images. From left to right: Warped moving image, predicted deformation,
Jacobian determinant (blue: positive values, red: negative values), difference
image between fixed and warped image, segmentation map boundaries.


*captionUnknown document class (or package)

# Learned Point Cloud Compression for Classification

Mateen Ulhaq and Ivan V. Bajić  School of Engineering Science
Simon Fraser University
Burnaby, BC, Canada
<EMAIL_ADDRESS>School of Engineering Science
Simon Fraser University
Burnaby, BC, Canada
<EMAIL_ADDRESS>
###### Abstract

Deep learning is increasingly being used to perform machine vision tasks such
as classification, object detection, and segmentation on 3D point cloud data.
However, deep learning inference is computationally expensive. The limited
computational capabilities of end devices thus necessitate a codec for
transmitting point cloud data over the network for server-side processing.
Such a codec must be lightweight and capable of achieving high compression
ratios without sacrificing accuracy. Motivated by this, we present a novel
point cloud codec that is highly specialized for the machine task of
classification. Our codec, based on PointNet, achieves a significantly better
rate-accuracy trade-off in comparison to alternative methods. In particular,
it achieves a 94% reduction in BD-bitrate over non-specialized codecs on the
ModelNet40 dataset. For low-resource end devices, we also propose two
lightweight configurations of our encoder that achieve similar BD-bitrate
reductions of 93% and 92% with 3% and 5% drops in top-1 accuracy, while
consuming only 0.470 and 0.048 encoder-side kMACs/point, respectively. Our
codec demonstrates the potential of specialized codecs for machine analysis of
point clouds, and provides a basis for extension to more complex tasks and
datasets in the future.

###### Index Terms:

Point cloud compression, coding for machines

##  I Introduction

Point clouds are used to represent 3D visual data in many applications,
including autonomous driving, robotics, and augmented reality. Recent advances
in deep learning have led to the development of deep learning-based methods
for machine vision tasks on point cloud data. Common tasks include
classification, segmentation, object detection, and object tracking. However,
current deep learning-based methods often require significant computational
resources, which impose hardware requirements. Such significant requirements
may not be physically or economically feasible for end devices.

One approach to address the issue of insufficient end-device computational
resources is to transmit the point cloud and other sensor data to a server for
processing. However, this introduces its own challenges, including the effects
of network availability, latency, and bandwidth. In order to reduce network
requirements, the end device may compress the point cloud data before
transmission. However, network capabilities vary depending on various factors,
including end-device location and network congestion. This means that
sufficient bandwidth may still not be available to transmit the point cloud
data. A hybrid strategy is to perform part of the machine task on the end
device itself. This can reduce the amount of data that needs to be transmitted
to the server for further processing, without exceeding the computational
budget of the end device [1]. This enhances the robustness of the end device
to varying network conditions, while potentially improving overall system
latency and adaptability [2].

We propose a novel learned point cloud codec for classification. Our learned
codec takes a point cloud as input, and outputs a highly compressed
representation that is intended solely for machine analysis. To our knowledge,
this is the first point cloud codec specialized for machine analysis. Existing
codecs for point clouds are designed to reconstruct point clouds intended for
human viewing. This means that a significant amount of bits is wasted on
encoding information that is not strictly relevant to machine analysis. By
partially processing the point cloud before compression, our codec is able to
achieve significantly better compression performance, without compromising
task accuracy.

In this paper, we present our task-specialized codec architecture in full and
lightweight configurations. We evaluate its rate-accuracy (RA) performance on
the ModelNet40 dataset [3]. We also investigate how the number of input points
(and thus reduced computation) affects the RA performance of our proposed
codec. Furthermore, we compare our proposed codec’s performance against
alternative non-specialized methods. Our code for training and evaluation is
available online111 https://github.com/multimedialabsfu/learned-point-cloud-
compression-for-classification.

##  II Related work

Point cloud classification models can be organized into groups based on the
type of input data they accept. Models such as VoxNet [4] take as input point
clouds that have been preprocessed into a voxel grid. Unfortunately, these
methods often use 3D convolutions, which require a significant amount of
computational resources. Additionally, since most voxels are usually empty,
these methods arguably waste a significant amount of computation on empty
space. Furthermore, the voxel grid representation is not very compact, and
thus requires a significant amount of memory for higher spatial resolutions
(e.g., a 32-bit tensor of shape $1024\times 1024\times 1024$ occupies 32 GB).
Models such as OctNet [5] take octrees as input. Octrees offer a more compact
representation of the voxelized point cloud by encoding the node occupancy in
bitstrings. Large unoccupied regions of space may be represented via a single
“0” node in an octree. Point-based models such as PointNet [6] and PointNet++
[7] directly accept raw point lists $(x_{1},x_{2},\ldots,x_{P})$, where
$x_{i}\in\mathbb{R}^{3}$ represents a point in a 3D space and $P$ is the
number of points. Some challenges faced with this input format include
designing order-invariant models (due to the lack of a worthwhile canonical
ordering of points), as well as in devising operations capable of using the
metric structure induced by point locality. Despite the challenges, point-
based models are able to achieve surprisingly competitive accuracy, and offer
the most promise in terms of minimizing computational requirements.

PointNet [6], which our proposed architecture is based on, can be represented
as an input permutation-invariant function:

$f(x_{1},\ldots,x_{n})=(\gamma\circ\pi)(h(x_{1}),\ldots,h(x_{n})),$

where $h$ is applied to each point $x_{i}$ individually, $\pi$ is a simple
permutation-invariant function, and $\gamma$ may be any function. In the
original PointNet architecture, $h$ is a weight-shared MLP, $\pi$ is a max
pooling function applied across the point dimension, and $\gamma$ is an MLP.

In the related field of learned image compression, Ballé _et al._ [8] proposed
a variational autoencoder (VAE) architecture for image compression. Here, the
model transforms the input ${\boldsymbol{x}}$ into a latent representation
${\boldsymbol{y}}$, which is then quantized into ${\boldsymbol{\hat{y}}}$ and
losslessly compressed using a learned entropy model. The codec is trained end-
to-end using the loss function

$\mathcal{L}=R+\lambda\cdot D({\boldsymbol{x}},{\boldsymbol{\hat{x}}}),$ (1)

where $D({\boldsymbol{x}},{\boldsymbol{\hat{x}}})$ is the distortion measure
between input ${\boldsymbol{x}}$ and decoded ${\boldsymbol{\hat{x}}}$, and $R$
is the estimate of the entropy of ${\boldsymbol{\hat{y}}}$. One simple entropy
model, known in the literature as an _entropy bottleneck_ , makes a
“factorized” prior assumption — that each element within a latent channel is
independently and identically distributed. It models a monotonically
increasing non-parametric cumulative distribution function using a
differentiable MLP. This mechanism has also shown effectiveness in learned
codecs in other fields, including learned point cloud compression (PCC), and
has been incorporated by a variety of works including [9, 10, 11, 12, 13].

We have also based our work on ideas introduced for machine tasks on images.
Early works demonstrated the use of standard codecs in compressing the latent
features [14]. More recently, approaches such as Video Coding for Machines
(VCM) [15] and Coding for Machines (CfM) have gained traction in the research
community. For instance, works such as [16, 17] demonstrate the potential
bitrate savings of scalable image compression in a multi-task scenario for a
machine vision task (e.g., facial landmark detection, or object detection) and
human vision (e.g., image reconstruction). In this work, we focus solely on a
single machine vision task applied to point cloud data.

##  III Proposed codec

###  III-A Input compression

In fig. 1(a), we show an abstract representation of an input codec, similar to
the “chain” configuration explored by [18] for end-to-end image compression
for machines. In this codec, the input point cloud ${\boldsymbol{x}}$ is
encoded directly, without any intermediate feature extraction. On the decoder
side, the point cloud is then reconstructed as ${\boldsymbol{\hat{x}}}$. Any
point cloud compression codec can be used for this purpose, including standard
non-learned codecs such as G-PCC [19]. Finally, the reconstructed point cloud
${\boldsymbol{\hat{x}}}$ is fed into a classification model (e.g., PointNet)
in order to obtain the class prediction ${\boldsymbol{\hat{t}}}$. This
approach provides a baseline for comparison with our proposed codec.

(a) input compression

(b) proposed

Figure 1: High-level comparison of codec architectures.

###  III-B Motivation for the proposed codec

An efficient task-specific codec can be developed using the concept of
Information Bottleneck (IB) [20]:

$\min_{p({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})}\quad
I({\boldsymbol{x}};{\boldsymbol{\hat{y}}})-\beta\cdot
I({\boldsymbol{\hat{y}}};{\boldsymbol{\hat{t}}}),$ (2)

where $I(\cdot;\cdot)$ is the mutual information [21],
$p({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})$ is the mapping from the input
point cloud ${\boldsymbol{x}}$ to the latent representation
${\boldsymbol{\hat{y}}}$, and $\beta>0$ is the IB Lagrange multiplier [20]. We
can think of $p({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})$ as feature
extraction followed by quantization. Hence, ${\boldsymbol{\hat{y}}}$ is fully
determined whenever ${\boldsymbol{x}}$ is given, so
$H({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})=0$, where $H(\cdot\mid\cdot)$
is the conditional entropy [21]. Therefore,
$I({\boldsymbol{x}};{\boldsymbol{\hat{y}}})=H({\boldsymbol{\hat{y}}})-H({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})=H({\boldsymbol{\hat{y}}})$.

Figure 2: Proposed codec architecture.

Furthermore, since decreasing $-\beta\cdot
I({\boldsymbol{\hat{y}}};{\boldsymbol{\hat{t}}})$ would improve the accuracy
of the task, we can use $\lambda\cdot
D({\boldsymbol{t}},{\boldsymbol{\hat{t}}})$ as a proxy for $-\beta\cdot
I({\boldsymbol{\hat{y}}};{\boldsymbol{\hat{t}}})$, where ${\boldsymbol{t}}$ is
the ground-truth label, ${\boldsymbol{\hat{t}}}$ is the label produced using
the compressed latent representation, and
$D({\boldsymbol{t}},{\boldsymbol{\hat{t}}})$ is a distortion measure.
Therefore, in our case, the IB (2) becomes:

$\min_{p({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})}\quad
H({\boldsymbol{\hat{y}}})+\lambda\cdot
D({\boldsymbol{t}},{\boldsymbol{\hat{t}}}).$ (3)

It is clear that the form of IB in (3) is analogous to the loss function (1).
We make use of this analogy to develop the proposed codec, which is described
next.

###  III-C Proposed architecture

In fig. 1(b), we show a high-level representation of our proposed codec
architecture. Following the terminology of [8], we refer to $g_{a}$ as the
_analysis transform_ , and $g_{s}$ as the _synthesis transform_. In this
architecture, the input point cloud ${\boldsymbol{x}}$ is first encoded into a
latent representation ${\boldsymbol{y}}=g_{a}({\boldsymbol{x}})$, which is
then quantized as ${\boldsymbol{\hat{y}}}=Q({\boldsymbol{y}})$, and then
losslessly compressed using a learned entropy model. Therefore,
$p({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})$ from the IB is $Q\circ g_{a}$.
Then, the reconstructed latent representation ${\boldsymbol{\hat{y}}}$ is used
to predict the classes ${\boldsymbol{\hat{t}}}=g_{s}({\boldsymbol{\hat{y}}})$.

Our proposed architecture, visualized in fig. 2, is based on the PointNet [6]
classification model. The input 3D point cloud containing $P$ points is
represented as a matrix ${\boldsymbol{x}}\in\mathbb{R}^{3\times P}$. The input
${\boldsymbol{x}}$ is fed into a sequence of encoder blocks. Each encoder
block consists of a convolutional layer, a batch normalization layer, and a
ReLU activation. As described in [6], the early encoder blocks must be applied
to each input point independently and identically, in order to avoid learning
input permutation-specific dependencies. Therefore, we use pointwise
convolutional layers with kernel size 1, which are exactly equivalent to the
“shared MLP” described in [6].

Following the sequence of encoder blocks, a max pooling operation is applied
along the point dimension to generate an input permutation-invariant feature
vector of length $N$. The resulting feature vector is then multiplied element-
wise by a trainable gain vector $v\in\mathbb{R}^{N}$, which is initialized to
$[1,1,\ldots,1]$ before training. Due to the batch normalization layers, the
resulting feature vector has small values. Thus, in order to improve training
stability and the rate of convergence, the feature vector is multiplied by a
constant scalar value of $10$. The resulting vector is the output of the
encoder-side analysis transform, which we label as ${\boldsymbol{y}}$.

Then, we quantize ${\boldsymbol{y}}$ via uniform quantization (specifically,
integer rounding) to obtain ${\boldsymbol{\hat{y}}}$. During training, uniform
quantization is simulated using additive uniform noise
$\mathcal{U}(-0.5,0.5)$. The quantized vector ${\boldsymbol{\hat{y}}}$ is then
losslessly encoded using a fully-factorized learned entropy model introduced
in [8].

On the decoder side, the decoded vector ${\boldsymbol{\hat{y}}}$ is fed into
an MLP consisting of fully-connected layers interleaved with ReLU activations
and batch normalizations. Before the last fully-connected layer, we use a
dropout layer that randomly sets 30% of its inputs to zero. The output of the
MLP is a vector of logits ${\boldsymbol{\hat{t}}}\in\mathbb{R}^{T}$, where $T$
is the number of classes.

We provide “full”, “lite”, and “micro” configurations of the proposed codec.
For each configuration, table I lists the number of layer output channel sizes
along with the estimated MAC (multiply-accumulate) counts222One MAC operation
may be considered equivalent to a FLOP (floating-point operation) on most
hardware.. Group convolutional layers are specified in the format
“output_size/group”. In contrast to [6], we do not use any input or feature
transformations in order to simplify the architectures for clearer analysis,
as well as to reduce the computational requirements.

TABLE I: Layer sizes and MAC counts for various proposed codec types Proposed | Encoder | Decoder | Encoder | Decoder
---|---|---|---|---
codec | layer sizes | layer sizes | MAC/pt | MAC
full | 64 64 64 128 1024 | 512 256 40 | 150k | 670k
lite | 8 8 16 16/2 32/4 | 512 256 40 | 0.47k | 160k
micro | 16 | 512 256 40 | 0.048k | 150k

###  III-D Lightweight and micro architectures

In addition to our “full” proposed codec, we also provide a lightweight
configuration, which we denote as “lite”. In this architecture, the encoder-
side layers contain fewer output channels. To further reduce encoder-side
computational costs, they also use group convolutions with channel shuffles in
between, as is done in ShuffleNet [22]. After training, the gain and batch
normalization layers may be fused into the preceding convolutional layer. The
“lite” architecture strikes a balance between RA performance and encoder
complexity. In fact, the encoder-side transform requires only 0.47k
MACs/point, which is significantly less than the 150k MACs/point required by
the “full” architecture encoder. For input point clouds consisting of $P=256$
points, the total MAC count for the “lite” codec is 120k, which is below the
corresponding decoder-side MAC count of 160k.

Additionally, we examine a “micro” architecture, whose encoder-side transform
consists of only a single encoder block with 16 output channels, and a max
pooling operation. This codec is useful for analysis and comparison — and yet,
it is also capable of surprisingly competitive RA performance.

##  IV Experiments

Our models were trained on the ModelNet40 [3] dataset, which consists of 12311
different 3D object models organized into 40 classes. We used an Adam
optimizer with a learning rate of 0.001. Our code was written using the
PyTorch, CompressAI [23], and CompressAI Trainer [24] libraries.

The loss function that is minimized during training is:

$\mathcal{L}=R+\lambda\cdot D({\boldsymbol{t}},{\boldsymbol{\hat{t}}}),$

where the rate $R=-\log p_{{\boldsymbol{\hat{y}}}}({\boldsymbol{\hat{y}}})$ is
the log of the likelihoods outputted by the entropy model, and the distortion
$D({\boldsymbol{t}},{\boldsymbol{\hat{t}}})$ is the cross-entropy between the
one-hot encoded labels ${\boldsymbol{t}}$ and the softmax of the model’s
prediction ${\boldsymbol{\hat{t}}}$. We trained different models to operate at
different rate points by varying the hyperparameter $\lambda\in[10,16000]$.

###  IV-A Proposed codec

For each of the proposed codec architectures, we trained a variation of each
codec to accept an input point cloud containing $P\in\mathcal{P}$ points. We
trained eight such variations for each of the values in the set
$\mathcal{P}=\\{8,16,32,64,128,256,512,1024\\}$. Although each codec is
capable of handling a variable number of points, training a separate model for
each $P$ guarantees that each codec is well-optimized for each rate-accuracy
trade-off.

###  IV-B Input compression codec

We compare our proposed codec against an “input compression” codec
architecture. For this codec, the encoder may be taken from any point cloud
codec. We have tested multiple codecs, including TMC13 [25] v14 (an
implementation of the G-PCC v2 [19] standard), OctAttention [12], and Draco
[26]. On the decoder side is the corresponding point cloud decoder, followed
by a PointNet classification model. We trained a PointNet model (without the
input and feature transforms) for each $P\in\mathcal{P}$.

We generated eight separate datasets of $P$-point point clouds, where each
point cloud was uniformly subsampled from the test dataset. Then, we
compressed and decompressed each point cloud from each $P$-point dataset at
various compression ratios. The compression ratio can be effectively
controlled by varying the amount of input scaling, which we denote by $S$.
(The input scaling parameter is directly proportional to the number of bins
used during uniform quantization of the input points.) We varied $S$ over the
set $\mathcal{S}=\\{1,2,4,\ldots,256\\}$ and $P$ over $\mathcal{P}$ to produce
$|\mathcal{P}|\cdot|\mathcal{S}|$ distinct datasets. We evaluated each dataset
associated with the pair $(P,S)\in\mathcal{P}\times\mathcal{S}$ on the
correspondingly trained PointNet models to obtain a set of rate-accuracy
points. Finally, we took the Pareto front of this set to obtain the best rate-
accuracy curve achieved by the tested input compression codec.

###  IV-C Reconstruction

In order to visually assess the contents of the machine task-specialized
bitstream, we trained a point cloud reconstruction network on top of our
trained models. This auxiliary network was trained to minimize the loss
function $\mathcal{L}=D({\boldsymbol{x}},{\boldsymbol{\hat{x}}})$, where we
used Chamfer distance for $D$.

We also identify a critical point set for a fixed point cloud. A critical
point set is a minimal set of points which generate the exact same latent
${\boldsymbol{y}}$, and correspondingly, the same bitstream. Formally, for any
given point cloud ${\boldsymbol{x}}$, let
${\boldsymbol{x}}_{C}\subseteq{\boldsymbol{x}}$ denote a (not necessarily
unique) critical point set. Then,
$g_{a}({\boldsymbol{x}}_{C})=g_{a}({\boldsymbol{x}})={\boldsymbol{y}}$, and
there is uniquely one valid critical point set $({\boldsymbol{x}}_{C})_{C}$
for ${\boldsymbol{x}}_{C}$, and it is itself. Since $g_{a}$ contains a max
pooling operation, the critical point set is not theoretically unique;
however, in practice, it is rare for there to be more than one critical point
set. A critical point set may be computed by
${\boldsymbol{x}}_{C}=\bigcup_{1\leq j\leq
N}\operatorname*{arg\,max}_{{\boldsymbol{x}}_{i}\in{\boldsymbol{x}}}\,(h({\boldsymbol{x}}_{i}))_{j}$,
where $\\{h({\boldsymbol{x}}_{i}):1\leq i\leq P\\}$ represents the entire set
of generated latent vectors immediately preceding max pooling.

##  V Results

fig. 3 shows the rate-accuracy (RA) curves for the proposed “full”, “lite”,
and “micro” codecs in comparison with the input compression codec. Also
included are two baseline accuracies taken from the official PointNet paper
[6], for the model with (89.2%) and without (87.1%) the input/feature
transforms. Since our compression models were all trained _without_ the
input/feature transforms, the lower baseline offers a more direct comparison.
In table II, we list the peak accuracies attained by each codec, as well as
the Bjøntegaard-Delta (BD) [27] improvements in rates and accuracies relative
to the reference input compression codec.

(a) “full” codec

(b) “lite” codec

(c) “micro” codec

(d) input compression codecs

Figure 3: Rate-accuracy curves evaluated on the ModelNet40 test set. TABLE II:
BD metrics and max attainable accuracies per codec

Codec | Max acc (%) | BD rate (rel %) | BD acc (%)
---|---|---|---
Input compression |  |  |
TMC13 [25] | 88.5 | 0.0 | 0.0
OctAttention [12] | 88.4 | -13.2 | +2.1
IPDAE [13] | 87.0 | -23.0 | +3.6
Draco [26] | 88.3 | +780.7 | -4.2
Proposed (full) |  |  |
$P=1024$ | 88.5 | -93.8 | +16.4
$P=512$ | 88.0 | -93.7 | +15.9
$P=256$ | 87.6 | -93.3 | +15.4
$P=128$ | 87.1 | -92.7 | +14.9
$P=64$ | 86.1 | -91.1 | +13.2
$P=32$ | 81.8 | -90.6 | +9.3
$P=16$ | 70.4 | -86.8 | -2.3
$P=8$ | 46.8 | -88.5 | -25.3
Proposed (lite) |  |  |
$P=1024$ | 85.0 | -93.0 | +13.5
$P=512$ | 85.5 | -92.8 | +14.2
$P=256$ | 84.4 | -92.4 | +12.8
$P=128$ | 84.0 | -91.6 | +12.5
$P=64$ | 81.3 | -88.5 | +9.8
$P=32$ | 76.3 | -88.7 | +4.9
$P=16$ | 66.2 | -86.1 | -4.1
$P=8$ | 43.6 | -90.2 | -28.0
Proposed (micro) |  |  |
$P=1024$ | 83.6 | -91.8 | +12.7
$P=512$ | 82.5 | -91.6 | +11.6
$P=256$ | 81.6 | -91.1 | +11.0
$P=128$ | 80.1 | -90.9 | +9.9
$P=64$ | 76.6 | -89.9 | +6.5
$P=32$ | 70.3 | -89.0 | +0.1
$P=16$ | 59.4 | -87.6 | -10.8
$P=8$ | 41.9 | -88.3 | -28.8

  * •

$P$ is the number of points in the input ${\boldsymbol{x}}$. The BD metrics
were computed using the TMC13 input compression codec as the reference anchor.

Our “full” proposed codec, which is an extension of PointNet (no transforms),
achieves the lower baseline accuracy at $120$ bits, and an 80% accuracy at
$30$ bits. Our “lite” proposed codec saturates in RA performance at around
$P=512$ input points. At around $P=256$, the total MAC count of the proposed
“lite” encoder is roughly equal to the decoder. As shown by the rate-accuracy
curves, the $P=256$ model does not suffer too significant a drop in rate-
accuracy performance. This suggests that our method is capable of achieving a
good trade-off between rate, accuracy, and runtime performance. Similarly, our
“micro” codec suffers a further slight drop in rate-accuracy performance, but
achieves another significant improvement in runtime performance. The input
compression codec is the worst performing codec, and attains the lower
baseline accuracy at roughly 2000 bits.

In fig. 5, we show various point clouds that have been reconstructed from the
bitstreams generated by our proposed codecs. For each codec, we include
samples reconstructed from bitstreams compressed at different rates. Above
each reconstructed point cloud, we show the corresponding reference point
cloud, with critical points marked in red.

i 284 bits

ii 43 bits

iii 18 bits

iv 12 bits

(a) “full” codec

i 124 bits

ii 36 bits

iii 19 bits

iv 13 bits

(b) “lite” codec

i 76 bits

ii 27 bits

iii 23 bits

iv 14 bits

(c) “micro” codec

Figure 5:  Reconstructions of a sample airplane 3D model from the ModelNet40
test set for various codecs and bitrates. For each reconstruction, its
corresponding reference point cloud is marked with _critical points_ in red.

##  VI Discussion

For input point clouds containing $P=1024$ points, our “full”, “lite”, and
“micro” codec configurations achieve an accuracy of 80% with as few as 30, 40,
and 50 bits. For comparison, $\log_{2}(40)\approx 5.3$ bits are required to
losslessly encode uniformly distributed class labels of the 40 classes from
ModelNet40. Our codec comes surprisingly close to this theoretical lower
bound, despite the fact that our architecture design omits the traditional MLP
“classifier” within the encoder. The same pointwise function is applied to all
points, and the only operation that “mixes” information between the points is
a pooling operation. This suggests that our encoder should possess limited
classification abilities, and yet it consumes only a few times more bits than
the theoretical lower bound.

To put this into perspective, consider coding for image classification, which
is one of the best developed areas in the field of coding for machines.
Current state-of-the-art (SOTA) approaches [28, 29, 30] for coding for image
classification on ImageNet [31] require upwards of $0.01$ bits per pixel (bpp)
to maintain a reasonable top-1 accuracy. With the typical input image
resolution of $224\times 224$, this works out to be around $500$ bits.
However, the maximum classifier output entropy with $1000$ classes is only
$\log_{2}(1000)\approx 10$ bits, which is several orders of magnitude lower.
Hence, the gap between the current SOTA and theoretical limits on coding for
image classification is much higher than what is achieved by our proposed
codec for point cloud classification.

To explore why our codec comes so close to the theoretical lower bound, we
propose the following arguments. Let ${\boldsymbol{x}}$ represent a possible
input point cloud, and let ${\boldsymbol{\hat{y}}}=(Q\circ
g_{a})({\boldsymbol{x}})$ be its quantized transformed latent representation.
Applying the data processing inequality to the Markov chain
${\boldsymbol{x}}\to{\boldsymbol{x}}\to{\boldsymbol{\hat{y}}}$, we determine
that $I({\boldsymbol{x}};{\boldsymbol{x}})\geq
I({\boldsymbol{x}};{\boldsymbol{\hat{y}}})$. Furthermore, since $Q\circ g_{a}$
is deterministic, $H({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})=0$, and so

$\displaystyle H({\boldsymbol{x}})=I({\boldsymbol{x}};{\boldsymbol{x}})\geq
I({\boldsymbol{x}};{\boldsymbol{\hat{y}}})=H({\boldsymbol{\hat{y}}})-H({\boldsymbol{\hat{y}}}\mid{\boldsymbol{x}})=H({\boldsymbol{\hat{y}}}).$

This indicates theoretically that the quantized latent representation
${\boldsymbol{\hat{y}}}$ must on average be at least as compressible as the
input point cloud ${\boldsymbol{x}}$ that it was derived from.

Since the critical point set ${\boldsymbol{x}}_{C}\subseteq{\boldsymbol{x}}$
produces the exact same ${\boldsymbol{\hat{y}}}$ as the original input point
cloud ${\boldsymbol{x}}$, we may use the same arguments as above to argue that

$H({\boldsymbol{x}})\geq H({\boldsymbol{x}}_{C})\geq
H({\boldsymbol{\hat{y}}}).$

This provides us with a potentially tighter bound. In fact, as shown in fig.
5iii, much of the general shape of the shown sample point cloud can be
reconstructed from only 23 bits of information. Furthermore, since
$|{\boldsymbol{x}}_{C}|\leq N$, there are only at most $32$ and $16$ distinct
critical points for the “lite” and “micro” codecs, respectively. This suggests
part of the reason for why our proposed codec achieves such big gains in
comparison to input compression.

##  VII Conclusion

In this paper, we proposed a new codec for point cloud classification. Our
experiments demonstrated that the “full” configuration of the codec achieves
stellar rate-accuracy performance, far exceeding the performance of
alternative methods. We also presented “lite” and “micro” configurations of
the codec whose encoders consume minimal computational resources, and yet
achieve comparable gains in rate-accuracy performance.

Our work may be extended to other point cloud tasks, such as segmentation and
object detection, or to more complex tasks involving larger models and larger
point clouds from real-world datasets. Our work also sets a good starting
point for further research into approaches for scalable and multi-task point
cloud compression. We hope that our work will help towards achieving the end
goal of more capable end devices.

## References

  * [1] Y. Kang, J. Hauswald, C. Gao, A. Rovinski, T. Mudge, J. Mars, and L. Tang, “Neurosurgeon: Collaborative intelligence between the cloud and mobile edge,” _SIGARCH Comput. Archit. News_ , vol. 45, no. 1, pp. 615–629, 2017.
  * [2] N. Shlezinger and I. V. Bajić, “Collaborative inference for AI-empowered IoT devices,” _IEEE Internet of Things Magazine_ , vol. 5, no. 4, pp. 92–98, Dec. 2022.
  * [3] Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao, “3D ShapeNets: A deep representation for volumetric shapes,” in _Proc. IEEE/CVF CVPR_ , 2015, pp. 1912–1920.
  * [4] D. Maturana and S. A. Scherer, “VoxNet: A 3D convolutional neural network for real-time object recognition,” in _Proc. IEEE/RSJ Int. Conf. Intelligent Robots and Systems (IROS)_ , 2015, pp. 922–928.
  * [5] G. Riegler, A. O. Ulusoy, and A. Geiger, “OctNet: Learning deep 3D representations at high resolutions,” in _Proc. IEEE/CVF CVPR_ , 2017, pp. 6620–6629.
  * [6] C. Qi, H. Su, K. Mo, and L. J. Guibas, “PointNet: Deep learning on point sets for 3D classification and segmentation,” in _Proc. IEEE/CVF CVPR_ , 2017, pp. 77–85.
  * [7] C. Qi, L. Yi, H. Su, and L. J. Guibas, “PointNet++: Deep hierarchical feature learning on point sets in a metric space,” in _NIPS_ , 2017.
  * [8] J. Ballé, D. C. Minnen, S. Singh, S. J. Hwang, and N. Johnston, “Variational image compression with a scale hyperprior,” _arXiv preprint arXiv:1802.01436_ , 2018.
  * [9] W. Yan, Y. shao, S. Liu, T. H. Li, Z. Li, and G. Li, “Deep autoencoder-based lossy geometry compression for point clouds,” _arXiv preprint arXiv:1905.03691_ , 2019.
  * [10] Y. He, X. Ren, D. Tang, Y. Zhang, X. Xue, and Y. Fu, “Density-preserving deep point cloud compression,” _arXiv preprint arXiv:2204.12684v1_ , 2022.
  * [11] J. Pang, M. A. Lodhi, and D. Tian, “GRASP-Net: Geometric residual analysis and synthesis for point cloud compression,” in _Proc. 1st Int. Workshop on Advances in Point Cloud Compression, Processing and Analysis_ , 2022.
  * [12] C. Fu, G. Li, R. Song, W. Gao, and S. Liu, “OctAttention: Octree-based large-scale contexts model for point cloud compression,” in _Proc. AAAI_ , vol. 36, no. 1, Jun. 2022, pp. 625–633.
  * [13] K.-S. You, P. Gao, and Q. T. Li, “IPDAE: Improved patch-based deep autoencoder for lossy point cloud geometry compression,” in _Proc. 1st Int. Workshop on Advances in Point Cloud Compression, Processing and Analysis_ , 2022.
  * [14] H. Choi and I. V. Bajić, “Near-lossless deep feature compression for collaborative intelligence,” in _Proc. IEEE MMSP_ , 2018, pp. 1–6.
  * [15] L.-Y. Duan, J. Liu, W. Yang, T. Huang, and W. Gao, “Video Coding for Machines: A paradigm of collaborative compression and intelligent analytics,” _IEEE Trans. Image Process._ , vol. 29, pp. 8680–8695, 2020.
  * [16] Y. Hu, S. Yang, W. Yang, L.-Y. Duan, and J. Liu, “Towards coding for human and machine vision: A scalable image coding approach,” in _Proc. IEEE ICME_ , 2020, pp. 1–6.
  * [17] H. Choi and I. V. Bajić, “Scalable image coding for humans and machines,” _IEEE Trans. Image Process._ , vol. 31, pp. 2739–2754, Mar. 2022.
  * [18] L. D. Chamain, F. Racapé, J. Bégaint, A. Pushparaja, and S. Feltman, “End-to-end optimized image compression for machines, a study,” in _Proc. DCC_ , 2020, pp. 163–172.
  * [19] K. Mammou, P. A. Chou, D. Flynn, M. Krivokuća, O. Nakagami, and T. Sugio, “G-PCC codec description v2,” 2019, ISO/IEC JTC1/SC29/WG11 N18189.
  * [20] N. Tishby, F. C. Pereira, and W. Bialek, “The information bottleneck method,” in _Proc. 37th annual Allerton Conf. on Communication, Control, and Computing_ , Sept. 1999, pp. 368–377.
  * [21] T. M. Cover and J. A. Thomas, _Elements of Information Theory_ , 2nd ed. Wiley, 2006.
  * [22] X. Zhang, X. Zhou, M. Lin, and J. Sun, “ShuffleNet: An extremely efficient convolutional neural network for mobile devices,” in _Proc. IEEE/CVF CVPR_ , 2018, pp. 6848–6856.
  * [23] J. Bégaint, F. Racapé, S. Feltman, and A. Pushparaja, “CompressAI: a PyTorch library and evaluation platform for end-to-end compression research,” _arXiv preprint arXiv:2011.03029_ , 2020.
  * [24] M. Ulhaq and F. Racapé, “CompressAI Trainer,” 2022. [Online]. Available: https://github.com/InterDigitalInc/CompressAI-Trainer
  * [25] D. Flynn and K. Mammou, “MPEG-PCC-TMC13,” 2021. [Online]. Available: https://github.com/MPEGGroup/mpeg-pcc-tmc13
  * [26] Google, “Draco: 3D data compression,” 2017. [Online]. Available: https://google.github.io/draco/
  * [27] G. Bjøntegaard, “Calculation of average PSNR differences between RD-curves,” in _ITU-T SC16/Q6 VCEG-M33_ , Apr. 2001, pp. 2–4.
  * [28] Y. Matsubara, R. Yang, M. Levorato, and S. Mandt, “Supervised compression for resource-constrained edge computing systems,” in _Proc. IEEE/CVF Winter Conf. on Applications of Computer Vision (WACV)_ , 2022, pp. 923–933.
  * [29] Z. Duan and F. Zhu, “Efficient feature compression for edge-cloud systems,” in _Proc. Picture Coding Symposium (PCS)_ , Dec. 2022, pp. 187–191.
  * [30] N. Ahuja, P. Datta, B. Kanzariya, V. S. Somayazulu, and O. Tickoo, “Neural rate estimator and unsupervised learning for efficient distributed image analytics in split-DNN models,” in _Proc. IEEE/CVF CVPR_ , Jun. 2023, pp. 2022–2030.
  * [31] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and F.-F. Li, “ImageNet: A large-scale hierarchical image database,” in _Proc. IEEE/CVF CVPR_ , 2009, pp. 248–255.