Masked Wavelet NeRF: Compact Neural Rendering
- The paper introduces a novel method that integrates multi-level DWT with trainable binary masks to achieve over 95% sparsity while preserving high PSNR.
- It employs rank-1 factorization of 3D grids and differentiable masking to compress neural radiance fields under 2 MB without sacrificing rendering quality.
- Experimental results demonstrate that Masked Wavelet NeRF outperforms prior methods by maintaining superior image quality and fast convergence at ultra-compact model sizes.
Masked Wavelet NeRF is a method for highly compact neural radiance field (NeRF) representations, integrating discrete wavelet transforms (DWT) and trainable masking for sparsity, optimized for fast convergence and efficiency in neural rendering tasks. The approach is designed for practitioners seeking to minimize memory footprint without sacrificing rendering quality, operating at state-of-the-art performance with budgets below 2 MB while employing differentiable components compatible with existing grid-based NeRF frameworks (Rho et al., 2022).
1. Foundation: Grid-Based Neural Radiance Fields
Baseline “fast NeRF” variants—such as DVGO, Plenoxels, and TensoRF—replace an MLP-dominant scene representation with explicit 3D grids and a small MLP for color regression. Specifically, the dense radiance field is described by:
- A 3D density grid of dimensions containing scalar opacity values.
- A 3D feature grid of identical size, holding F-dimensional feature vectors per voxel.
Rendering a ray proceeds via:
- Sampling a 3D point along the ray.
- Trilinear interpolation of both grids at to yield (opacity) and (feature vector).
- Passing through an MLP (with the viewing direction) to predict color:
Standard volumetric rendering integrates and 0 for the final output. This baseline method converges rapidly but typically incurs tens to hundreds of megabytes in grid storage (Rho et al., 2022).
2. Discrete Wavelet Transform for Parameter Efficiency
Masked Wavelet NeRF introduces a multi-level DWT to compactly encode the 3D scene representation. Direct application of a full 3D DWT is computationally prohibitive (cubic complexity), so a rank-1 factorization (as in TensoRF) decomposes each 3D grid into three 2D planes and a 1D vector per rank. To each 2D plane:
- A single-level DWT produces four subbands: 1, notated as:
2
- The transform recursively applies to the 3 band for 4 levels:
5
Resulting coefficients 6 are the representation stored and optimized.
The inverse DWT reconstructs 7 by upsampling and convolving with appropriate scaling and wavelet filters, summing over all subbands. Level-wise, high-pass bands are scaled by 8 during mask learning to stabilize training—without this, high-pass bands dominate gradients and degrade reconstruction fidelity.
3. Trainable Binary Mask: Sparsity via Differentiable Masking
To further enhance compactness, Masked Wavelet NeRF learns a binary mask over each wavelet coefficient:
- Each coefficient 9 possesses an associated real-valued mask parameter 0.
- The binary mask is 1, employing the Heaviside step function (2).
- During forward passes, masked coefficients are 3.
- Backpropagation utilizes the straight-through estimator (STE) with a sigmoid surrogate 4 for gradient flow:
5
Joint training optimizes all subbands, mask parameters, and MLP weights by minimizing a sum of standard rendering loss (e.g., MSE or RGB 6) and a sparsity-promoting regularizer,
7
with 8 setting the balance between accuracy and sparsity. In practice, 9 yields 0 sparsity with negligible PSNR loss (1 dB).
4. Compression and Inference Pipeline
Following training, the retained representation consists of:
- Nonzero wavelet coefficients 2.
- The compact binary mask 3.
Compression is achieved through:
- Coefficient thresholding and quantization: Discard zeroed coefficients and apply uniform 4-bit quantization to the remainder. Store quantized values 5, scale 6, and offset 7.
- Mask compression, totaling three steps:
- Packing mask bits into 8-bit bytes.
- Run-length encoding to exploit zero runs.
- Huffman coding on RLE symbols.
Subband grouping and level-wise blockwise encoding, using separate RLE+Huffman codes for additional space reduction (∼30%).
This process yields an end-to-end compression pipeline that is fully differentiable (except Huffman coding) and inserts seamlessly into grid-based NeRF engines.
5. Experimental Outcomes and Quantitative Benchmarks
Masked Wavelet NeRF was evaluated on multiple datasets (NeRF Synthetic, NSVF Synthetic, Tanks&Temples, LLFF), using TensoRF VM-192 as the backbone (17.9 MB raw storage):
- Under a 2 MB budget, sparsity reaches up to 8 zeroed wavelet coefficients.
- 8-bit quantization is used for kept coefficients; MLP weights use 32 bits.
- On NeRF Synthetic:
- Baseline VM-192 (17.9 MB): 32.91 dB
- Masked Wavelet NeRF @0.83 MB: 31.95 dB
- CP-384 [25] @0.72 MB: 31.18 dB
- cNeRF [14] @0.70 MB: 30.49 dB
- CCNeRF [13] @4.4 MB: 30.55 dB
- For NSVF Synthetic, at 0.87 MB, Masked Wavelet NeRF achieves 35.18 dB versus CP-384’s 33.92 dB at 0.72 MB.
The approach consistently outperforms prior compression methods at comparable or lower sizes, preserving visual quality and PSNR.
6. Ablation Studies and Methodological Insights
A series of ablation studies highlight method robustness and optimal settings:
- DWT vs. spatial vs. DCT: DWT-based grids lose only 0.2 dB at >90% sparsity, compared to 2–3 dB in spatial grids; DWT outperforms DCT by 0.6 dB at 95% sparsity.
- Multi-level DWT: Increasing decomposition levels (9 to 0) incrementally raises sparsity (80%→97%) and improves PSNR.
- Level-wise scaling: Absence of 1 scaling reduces PSNR from 31.95 dB to 31.15 dB.
- Mask learning vs. fixed thresholding: Learned binary masks plus sparsity regularization yield a 0.5–1 dB higher PSNR than magnitude thresholding at the same sparsity.
- Optimizer parameters: Adam, with 0.02 base LR for grid/mask, 0.001 for MLP, decayed by 0.1 every 30k iterations; mask binarization is staged to stabilize early training.
- Overhead and efficiency: The approach is ≈1.5× slower than spatial grids (A100) but remains 5–10× faster end-to-end than classical NeRF-MLP models.
7. Significance and Integration with Existing Frameworks
Transitioning from a dense spatial grid to a wavelet basis with a jointly-trained binary mask achieves radiance field models under 1 MB without loss of rendering fidelity compared to 17 MB grid baselines. All major operations—including DWT transforms, mask learning, and entropy coding—are differentiable, enabling seamless adoption in contemporary grid-based NeRF infrastructures. Such compactness, achieved without accuracy compromise, is salient for resource-constrained deployment and efficient neural scene representation (Rho et al., 2022).