On Quantizing Implicit Neural Representations

TL;DR

This work shows that a non-uniform quantization of neural weights can lead to significant improvements and demonstrates that a clustered quantization enables improved reconstruction of signals using binary neural networks.

Abstract

The role of quantization within implicit/coordinate neural networks is still not fully understood. We note that using a canonical fixed quantization scheme during training produces poor performance at low-rates due to the network weight distributions changing over the course of training. In this work, we show that a non-uniform quantization of neural weights can lead to significant improvements. Specifically, we demonstrate that a clustered quantization enables improved reconstruction. Finally, by characterising a trade-off between quantization and network capacity, we demonstrate that it is possible (while memory inefficient) to reconstruct signals using binary neural networks. We demonstrate our findings experimentally on 2D image reconstruction and 3D radiance fields; and show that simple quantization methods and architecture search can achieve compression of NeRF to less than 16kb with minimal loss in performance (323x smaller than the original NeRF).

Figures (10)

• Figure 1: Comparison of cluster and uniform quantization on a small NeRF model. Top: 4-layers, 64 neurons per layer. Bottom: 4-layers, 128 neurons per layer. Quantization to 3-bits-per-weight.
• Figure 2: Left: Different uniform quantization ranges. Right: Comparison between decision boundaries for the examined uniform and cluster quantizations. A uniform quantization scheme equally divides the quantization range. A cluster quantization divides the quantization range such that each partition contains equal mass of the distribution.
• Figure 3: QAT Per epoch training cycle: The current full-precision weight matrix $W$ is quantized $q(W)=W_k$. The loss function $L(W_k)$ is calculated for the input data, and the error derivatives $\pdv{L(W_k)}{W_k}$ calculated. The full-precision weight matrix $W$ is then updated using backpropagation.
• Figure 4: Effect of repartitioning on Total Layerwise Quantization Error (TLQE) [log scale]. Recalculating partitions reduces quantization error at the cost of increased computation. Architecture: 1 hidden layer, 18 neurons, 5-bit weights, K-means quantization.
• Figure 5: Perceptual evaluation (CIFAR10)