Magnituder Layers for Implicit Neural Representations in 3D

Zip-NeRF is known for its large memory consumption (referring to NeRF-baselines). Is it possible to reduce the massive overheads by applying MAG layers or the improvement can only be time-wise?

Thank you for this question. The large memory consumption in Zip-NeRF primarily stems from its explicit grid structure, which MAG layers alone cannot reduce. This is precisely where fully implicit models, such as vanilla NeRF, become advantageous under strict memory constraints. In such scenarios, MAG layers serve as an effective and easy-to-implement method for accelerating inference without adding to memory overhead.

Quantitative results are not satisfactory enough

We would like to clarify that we get gains for not only NERF but for iSDF, MIP-NERF as well when we train them from scratch. For NERF, we get 24 % speedups while suffering only 1% drop in quality. Similarly, for iSDF we get 23% gain (see the trade-off plot in Fig 7 (right)). To get more quality gains, we borrowed ideas from recent works on random feature mechanisms.

[1,2] have shown improved performance when the projection matrix $\mathbf{G}$ can be trained instead of sampled from a fixed distribution. We show that it is indeed true in our case as well. Note that this incurs no additional cost in inference. In all cases, including NERF, D-NERF, MipNERF, ZipNERF, iSDF, ViT we show that using trainable $\mathbf{G}$, our performance drop is minimal and we have updated all the tables in the paper with the new updated results. For convenience, we summarize them in the table below :

Thus our method becomes quite powerful in INRs and in models where there are a lot of MLP layers. Furthermore these gains are higher than some existing published work (for ex: see [3] where the inference gain is around 18% with a quality drop of 3%).

Our main focus in this work is in INRs where one has a lot of MLPs. For other shallower architectures, our methods provide gains but the gains are small compared to INRs as we have a limited room to inject our MAG layers.

Finally, our distillation method is actually a real world test case of Q1 where the distributions are coming from pretrained models. Our MAG-layers can approximate complex outputs well enough that they can simply replace FFLs, thus justifying the design choice of $f$. Thus this simple layer wise distillation method is supposed to be a test-time inference speed up method where one does not have access to a lot of resources. This distillation takes less than a second and very limited GPU resources to train. Our new results on layer-wise distillation can actually maintain quality while reducing inference speeds (see Fig 8 and 9). This method is designed to complement other methods like pruning and quantization to speed up inference.

Finally since the number of random features is smaller than the dimension of the models, the memory footprint for training these models is less than the base models. Furthermore, the design of these layers allow for combining the MAG-layer with the subsequent linear linear effectively creating a single linear layer (please see the details in Appendix A) which further reduces the memory footprint during inference which is beneficial for deployment in low resource scenarios like edge devices.

[1] On Learning the Transformer Kernel. Chowdhury et al. 2022

[2] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. Zhang et al. ICLR 2024

[3] CoordX: Accelerating Implicit Neural Representation with a Split MLP Architecture. Liang et al. ICLR 2022

Converge of MAG-MLP vs regular MLP

Thank you for the question. We do not see any issues in convergence for the MAG-MLP. We simply use ADAM as an optimizer and the auto-differentiation in PyTorch to train these models.

Pseudo-code for our algorithm

We have added a PyTorch style pseudo-code for our algorithm in Appendix F.

Dear Reviewer,

Thank you for your response. We will add a discussion around low rank decompositions with focus around dictionary learning and separable convolutions. Please do not hesitate to reach out if you have any questions.

If our rebuttal has successfully answered your concerns, we would kindly request the Reviewer to adjust their scores.

Sincerely,

The Authors.

Given the current trend towards neural 3D representations that use very few and tiny MLP layers or even no MLP layers (e.g., InstantNGP, 3DGS), the benefits of applying MAG-layers to such architectures would likely be very insignificant.

While methods like Instant-NGP, and 3DGS achieve high speed-ups, they rely on hybrid or explicit representations that increase memory usage. Thus fully implicit methods are useful in scenarios with limited memory resources like edge devices. Both SDFs and NERFs find applications in various robotics tasks. We present two examples below :

• Truncated Signed Distance Fields were recently successfully applied to create 3D-fields for Robotic manipulation (re-arrangement problem) (applying large VLMs in zero-shot fashion + MPC on learned dynamics) [1] .
• In [2], the authors used a NERF to create reconstructions of an office environment.

Thus our main focus in this work is on INRs. For models with a large number of MLPs, like NERF, iSDF, ViT and MIP-NERF, we consistently get almost $20$% gains in inference speeds with minimal drop in quality (about 2%).

[1] $D^3$ Fields: Dynamic 3D Descriptor Fields for Zero-Shot Generalizable Rearrangement. Wang et al. CoRL 2024

[2] Mobility VLA: Multimodal Instruction Navigation with Long-Context VLMs and Topological Graphs. Chiang et al. CoRL 2024

MAG Layers for other applications

Thank you for this question.

We apply MAG-layers for the task of image classification using Vision Transformer (ViT). In this setting, we replace a part of the Feed-Forward Network (FFN) block in ViT with the MAG layer, namely the expansion layer with the GeLU activation. Note that : the MAG layer can then be combined with the following linear layer to create a single linear layer of size $[r, 768]$, where $r$ is the number of random features. This allows for a significant compression of the ViT model. To summarize our findings : by replacing top-6 layer’s MLP blocks with MAG reduces the size of the ViT model from 346 Mb to 196.85 Mb. This speeds up inference by almost $\mathbf{35}$% and has minimal impact on accuracy (less than 2%) . We present more details in Appendix E Image Classification Experiments. For convenience, we also show the results comparing inference parameters (in millions) vs accuracy in the table below :

Speed up during training

Thank you for this question. One of the benefits of using a smaller number of random features is that it shrinks the memory footprint during training. We compare the training time for iSDF where our method achieves 10% improvement in training time.

Reducing training parameters of baseline NERF

Thank you for proposing this experiment. We added these results in Fig 5. We show that our models offer similar or higher reconstruction quality than the models with the same number of training parameters while being consistently faster. The speed-ups are achieved due to the compression trick as explained in Appendix A which allows us to combine any subsequent FFL. This results in effectively converting a 2-layer neural network into a single layer. For NERF, we observe that reducing the size of the model to match our parameter count results in 2 % degradation in quality, while still being almost 5 % slower.

Furthermore, we also show that reducing the layer size in iSDF to match the number of training parameters in our proposed MAG-iSDF results in a 4.5% performance decrease compared to MAG-iSDF. The smaller iSDF is also 2% slower than our MAG-iSDF due to our novel compression trick. (see Fig 7).

Thus simple speed-ups techniques do not work quite as well as our MAG-layers which can retain quality as well as achieve speedups.

I went through the author's responses. I appreciate the additional efforts made by the authors during the rebuttal. Below are my responses:

1. The authors made additional experiments to show that MAG layer indeed performs better than simply the number of parameters in the layers. The results on NeRF & D-NeRF, the dense MLPs, show the superiority of MAG layers. However, on Zip-NeRF, MAG-layer has a slight quality drop compared to the reduced parameter version. These results indicate that the MAG layer can be a better way to reduce parameter size and improve the inference speed, for models with dense MLP layers. I acknowledge this contribution claimed in the paper.

2. I pretty much agree with Reviewer TtzD's comments. The MAG-layer should not only positioned as some acceleration technique for 3D representation. The proposed method is more generalizable and can be applied to models beyond INR, for example, like ViT, LLM, etc.

Given the current status of this submission, I will keep my rating score around the borderline of acceptance.

Other methods in 3D representation, such as 3D Gaussian Splatting, Instant-NGP, and Neus2, that significantly increase inference speed (up to 10x) without compromising rendering quality.

While methods like 3D-Gaussian Splatting (3D-GS), Instant-NGP, and NeuS2 achieve high speed-ups, they rely on hybrid or explicit representations that increase memory usage. This high memory usage is not well-suited for scenarios with limited resources like on edge devices, where purely explicit methods may not be feasible. Both SDFs and NERFs find applications in various robotics tasks. We present two examples below :

• Truncated Signed Distance Fields were recently successfully applied to create 3D-fields for Robotic manipulation (re-arrangement problem) (applying large VLMs in zero-shot fashion + MPC on learned dynamics) [1] .
• In [2], the authors used a NERF to create reconstructions of an office environment.

Thus our main focus in this work is on INRs. Our MAG layer is designed to reduce both inference time and memory requirements with only a minimal drop in quality and thus replacing MLPs by our layers can speed up INRs. For models with a large number of MLPs, like NERF, iSDF, ViT and MIP-NERF, we consistently get at least $20$% gains inference speeds with minimal drop in quality.

To summarize : For NeRF, we achieve a 24% improvement. Similarly, iSDF shows a 23% gain. Moreover, for ViT, we achieve nearly a 35% boost in inference speed. For MIP-NERF, we show an improvement of 18%. Overall, the quality drop across all cases is consistently minimal (about 2%).

Furthermore these gains for NERF are higher than some existing published work (for ex: see [3] where the inference gain is around 18% with a quality drop of 3%).

[1] $D^3$ Fields: Dynamic 3D Descriptor Fields for Zero-Shot Generalizable Rearrangement. Wang et al. CoRL 2024

[2] Mobility VLA: Multimodal Instruction Navigation with Long-Context VLMs and Topological Graphs. Chiang et al. CoRL 2024

[3] CoordX: Accelerating Implicit Neural Representation with a Split MLP Architecture. Liang et al. ICLR 2022

A simple baseline by reducing the number of layers or matching the training parameters

Thank you for proposing this experiment. We added these results in Fig 5. We show that our models offer similar or higher reconstruction quality than the models with the same number of training parameters while being consistently faster. The speed-ups are achieved due to the compression trick as explained in Appendix which allows us to combine any subsequent FFL. This results in effectively converting a 2-layer neural network into a single layer. For NERF, we observe that reducing the size of the model to match our parameter count results in 2 % degradation in quality, while still being almost 5 % slower.

Furthermore, we also show that reducing the layer size in iSDF to match the number of training parameters in our proposed MAG-iSDF results in a 4.5% performance decrease compared to MAG-iSDF. The smaller iSDF is also 2% slower than our MAG-iSDF due to our novel compression trick. (see Fig 7).

Thus simple speed-ups techniques do not work quite as well as our MAG-layers which can retain quality as well as achieve speedups.

Other tasks like classification using Transformers.

Thank you for proposing this experiment. In this setting, we replace a part of the Feed-Forward Network (FFN) block in Vision Transformer (ViT) with the MAG layer, namely the expansion layer with the GeLU activation. Note that : the MAG layer can then be combined with the following linear layer to create a single linear layer of size $[r, 768]$, where $r$ is the number of random features. This allows for a significant compression of the ViT model. To summarize our findings : by replacing top-6 layer’s MLP blocks with MAG reduces the size of the ViT model from 346 Mb to 196.85 Mb. This speeds up inference by almost $\mathbf{35}$% and has minimal impact on accuracy (less than 2%). We present more details in Appendix E Image Classification Experiments. For convenience, we also show the results comparing inference parameters (in millions) vs accuracy in the table below :

I appreciate the authors' detailed response to my comments. However, I still have a few remaining concerns.

The claim that "Our method is 1% behind the baseline in terms of PSNR but 24% faster in inference speed" is problematic. PSNR is defined in the logarithmic domain, making it inappropriate to directly compare percentage differences between PSNR and inference speed. Moreover, the newly provided Figure 5 indicates that, when the number of parameters is the same, the performance of the original MLP and the proposed MAG layer is very similar, with differences so minor that they are unlikely to have any meaningful impact.

As a result, I find the significance of this work within the context of 3D reconstruction, which is the primary focus of this paper as suggested by its title, to be quite limited. That said, this work might have greater relevance in the area of neural network design. However, to make such a claim, the authors would need to rethink the narrative of their paper and include comparisons to other neural networks such as KAN. This is outside my area of expertise, but I personally believe this work might be better positioned in the field of neural network design, where 3D reconstruction serves as just one application of the proposed method.