S WITCH -N E RF: L EARNING S CENE D ECOMPOSITION WITH M IXTURE OF E XPERTS FOR L ARGE - SCALE N EURAL R ADIANCE F IELDS
Anonymous authors
Paper under double-blind review
A BSTRACT
The Neural Radiance Fields (NeRF) have been recently applied to reconstruct building-scale and even city-scale scenes.
To model a large-scale scene efficiently, a dominant strategy is to employ a divide-and-conquer paradigm via performing scene decomposition, which decomposes a complex scene into parts that are further processed by different sub-networks.
Existing large-scale NeRFs mainly use heuristic hand-crafted scene decomposition, with regular 3D-distance-based or physical-street-block-based schemes.
Although achieving promising results, the hand-crafted schemes severely limit the capabilities of NeRF in large-scale scene modeling.
First , it is extremely challenging to manually design a universal scene decomposition rule for different complex scenes, leading to adaptation issues while applying the model to different scenarios.
Second , the decomposition procedure is not learnable, hindering the network from jointly optimizing the scene decomposition and the radiance fields in an end-to-end manner, to better model the scene with different sub-networks.
Third , the different sub-networks are typically optimized independently, and thus the inconsistency among them cannot be effectively handled during the optimization.
To tackle these issues, in this paper, we propose Switch-NeRF, a novel end-to-end large-scale NeRF with learning-based scene decomposition.
We design a gating network to dispatch 3D points to different NeRF sub-networks.
The gating network can be optimized together with the
NeRF sub-networks for different scene partitions, by a design with the Sparsely
Gated Mixture of Experts (MoE).
The outputs from different sub-networks can also be fused in a learnable way in the unified framework to effectively guarantee the consistency of the whole scene.
Furthermore, the proposed MoE-based
Switch-NeRF model is carefully implemented and optimized to achieve both highfidelity scene reconstruction and efficient computation.
Our method establishes clear state-of-the-art performances on several large-scale datasets.
To the best of our knowledge, we are the first to propose an applicable end-to-end sparse NeRF network with learning-based decomposition for large-scale scenes.
I NTRODUCTION
The Neural Radiance Fields (NeRF) method (Mildenhall et al., 2020) has gathered wide popularity in novel-view synthesis and 3D reconstruction due to its high quality and simplicity.
It encodes a 3D scene from multiple 2D posed images.
The original NeRF typically targets small scenes or objects, while in real-world applications such as autonomous driving and augmented reality (AR) / virtual reality (VR), it is critically important to build NeRF models to effectively handle large-scale scenes.
The problem of a large-scale NeRF is that more data typically requires a higher network capacity (number of network parameters).
A na ¨ ıve solution is to densely increase the network width and depth.
However, this will also greatly increase the computation for each sample and is harder to optimize.
A more applicable network should have a large capacity while maintaining almost constant computational cost for each sample.
Therefore, building an applicable large-scale NeRF can be considered as building a sparse neural network.
The core of the design is to select different network parameters ( i.e. sub-networks) for different inputs.
This procedure can be formulated as a scene decomposition problem in the NeRF task.
Each sub-network handles a different part of the scene.
[Figure]
The Block-NeRF (Tancik et al., 2022) clusters images by dividing the whole scene according to street blocks.
The sub-networks in both methods are trained separately.
Our Switch-NeRF learns to decompose the 3D points by a trainable gating network and the whole network is trained end-to-end.
Along the scene decomposition and learning a sparse neural network, recent Mega-NeRF (Turki et al., 2022) and Block-NeRF (Tancik et al., 2022) have extended NeRF to building-scale and even city-scale scenes based on heuristic hand-crafted scene decomposition.
As depicted in Fig.
1, the
Mega-NeRF and Block-NeRF simply use 3D sampling distances or street blocks to decompose the scene and train different NeRF models separately.
With promising results on large-scale scenes, their hand-crafted scene decomposition methods still have several limitations: (i) large-scale scenes are essentially complex and irregular.
It is extremely challenging to design a universal scene decomposition rule for different scenes in a hand-crafted way.
This accordingly brings adaptation issues when applying them to model distinct scenarios in the real world.
(ii) It requires rich priors of the target scene, such as the bounds of the scene, to deploy the partition centroids as in Mega-NeRF and the physical distribution of the scene images as in Block-NeRF.
These priors may not be available in practical applications.
(iii)
Their decomposition is not learnable, hindering the network from jointly optimizing the scene decomposition and the radiance fields in an end-to-end manner.
The gaps between the decomposition, composition and NeRF optimization may lead to sub-optimal results.
Besides, the different sub-networks are typically trained separately, and thus the inconsistency among them cannot be effectively handled during the optimization.
They usually consider a large overlapping between different partitions to align different sub-networks (Tancik et al., 2022).
To address above-mentioned issues, in this paper, we make the following contributions.
An end-to-end framework for joint learning of scene decomposition and NeRF.
We present
Switch-NeRF, an end-to-end sparse neural network framework, which jointly learns the scene decomposition and NeRF.
As shown in Fig.
1c, we propose a learnable gating network for scene decomposition.
It dynamically selects and sparsely activates a sub-network for each 3D point.
The overall network is trained end-to-end without any heuristic intervention.
We do not require any priors of the 3D scene shape or the distribution of scene images, leading to a generic framework for large-scale scenes.
Since the selection of sub-networks in training is a discrete operation, a critical problem is how to back-propagate gradients into the gating network.
We use the strategy from the Sparsely-Gated Mixture-of-Experts (MoE) (Shazeer et al., 2017) to deal with this problem.
We structure our sub-networks as NeRF experts for different scene partitions.
3D points are dispatched into different NeRF experts based on the gating network.
Besides the gating network, we also design a head to unify the predictions of multiple NeRF experts, which aligns the high-level implicit features from different NeRF experts to effectively address the inconsistency problem.
Efficient network design and implementation.
With the framework design of Switch-NeRF, however, it is not trivial to optimize and implement it efficiently and stably.
In NeRF, the number of 3D points in a forward pass is orders of magnitude larger than that of the input tokens in other NLP and vision tasks within an MoE framework.
Dispatching samples to different NeRF experts inevitably introduces large computation and memory usage.
Therefore, we consider dispatching 3D points only once with an effective gating network, and we design a deeper gating network to guarantee enough parameters to boost the accuracy of the scene rendering.
Another problem is that common MoE implementations (Hwang et al., 2022) define a capacity factor to limit the number of tokens dispatched to each expert.
This dynamically drops overflow 3D points for our NeRF experts.
It works well when training the network but brings a large influence on the testing accuracy.
To address this issue, we implement a full dispatch operation by CUDA based on Tutel (Hwang et al., 2022) to significantly improve the testing performance, while avoiding unnecessary memory allocation.
High-quality results.
Extensive experiments are conducted on large-scale challenging benchmarks.
Qualitative results show that our network can learn reasonable decomposition of large-scale complex scenes.
It establishes state-of-the-art performances and shows clearly superior results and scalability with much less network parameters compared to hand-crafted decomposition methods.
R ELATED W ORK
Neural Radiance Field .
The Neural Radiance Field (NeRF) is proposed by Mildenhall et al.
(2020) to use volumetric rendering for novel view synthesis from posed images.
It encodes a 3D scene into a multilayer perceptron (MLP), which is simple and requires very limited priors of the scene.
Due to the success of NeRF in high-quality rendering and 3D reasoning, many works have been proposed to improve its efficiency (Reiser et al., 2021;
Yu et al., 2021; M ¨ uller et al., 2022), accuracy (Barron et al., 2021; Verbin et al., 2022) and apply it to challenging scenes (Zhang et al., 2020; MartinBrualla et al., 2021; Xiangli et al., 2022) and 3D reconstruction tasks (Wang et al., 2021).
We pay more attention to the closely related works, i.e. NeRF methods for large-scale scenes.
As shown in Fig.
1a, Mega-NeRF (Turki et al., 2022) proposes to use a simple 3D distance-based method to cluster training pixels into parts that can be trained separately by different NeRF models.
It samples centroids uniformly in the 3D scene and groups 3D points in testing.
Block-NeRF (Tancik et al., 2022) proposes to scale NeRF to city-level scenes.
As shown in Fig.
1b, it divides the whole scene based on physical distribution of the scene images, i.e. partitioning through street blocks.
The scene images in different street blocks are trained separately by sub-networks.
In the testing, both methods consider offline fusion of prediction results from different sub-networks.
In contrast to
Mega-NeRF and Block-NeRF which use hand-craft scene decomposition, our Switch-NeRF jointly learns the scene decomposition and a large-scale NeRF in an end-to-end manner.
Our method is not dependent on any priors of the 3D shape and the physical image distribution of a target scene.
Therefore, our method is more generic for arbitrary large-scale scenes.
Mixture of Experts (MoE) .
Modern MoE methods mainly follow the work of Shazeer et al.
It proposes a Sparsely-Gated-Mixture-of-Experts layer in place of the feed-forward network ( i.e. FFN or MLP) in a language model.
It designs a vanilla Top-
k gating network to dispatch samples into k experts, and proposes an auxiliary loss for balancing the training of different experts.
The MoE has been widely used in Natural Language Processing (NLP)
(Lepikhin et al., 2021;
Fedus et al., 2022) and Vision (Riquelme et al., 2021).
Switch Transformer (Fedus et al., 2022) suggests that Top- 2 gating is not necessary.
It trains a network of high quality with Top- 1 gating to largely reduce the dispatch computation and communication.
Besides, different gating mechanisms are also proposed, such as Hash Routing (Roller et al., 2021) and BASE (Lewis et al., 2021).
There have been popular implementations of MoE in Mesh-TensorFlow (Shazeer et al., 2018), Deepspeed (Rajbhandari et al., 2022) and Tutel (Hwang et al., 2022), which focus on improving large-scale distributed training.
To the best of our knowledge, none of the existing works considers developing a Mixture-of-NeRF-Experts (MoNE) for large-scale scenes.
We design an effective structure of MoNE for jointly learning scene decomposition and NeRF, and also improve the efficiency of
MoNE by handling the issue of dispatching large-scale 3D points to different NeRF experts.
S WITCH -N E RF
Our Switch-NeRF is a Sparse Neural Radiance Field network targeting large-scale scenes.
A framework overview of Switch-NeRF is depicted in Fig.
We represent a large-scale 3D scene as a sparse 5D radiance function F Θ : ( x , d ) → ( c , σ ) , where ( x , d ) are a 3D point and its view direction, ( c , σ )
[Table]
z )
[Table]
+jIKV3UU9qVtGiq/p5IsTBmJCLXKbAdmHlvIv7ntRLbuwhTJuPEUklmi3oJR1ahyfeoyzQllo8cwUQzdysiA6wxsS6jggshmH95kdRPysFZ+fTutFS5zeLIwwEcwjEEcA4VuIEq1ICAgGd4hTdPey/eu/cxa8152cw+/IH3+QPSpZB3</latexit> AE
[Figure]
base
[Figure] [Figure]
/v2J2ThCP/wsZCEVv/jZ3/xk1yhSY
[Figure]
GHi8N8PMPD
[Figure]
WwqDrfju5tfWNza38dmFnd2//
[Figure]
/FmtPQl7zpj25mfvOJayMi9YCTmHdDOlAiEIyilR47yMfoB
[Figure]
l42iuW3LI7B1klXkZKkKHWK
[Figure]
hFLQq6QSWpM23Nj7KZUo2CSTwudxPCYshEd8LaliobcdNP5xVNyZpU
[Figure]
the corresponding gate value and sent to a head for density σ and color c prediction with direction d and appearance embedding.
The rendering loss is used for supervision.
The images on the left of each expert are the visualization of 3D radiance fields handled by different experts.
are predicted color and density, Θ represents network parameters.
F Θ sparsely activates only a part of its parameters for an input x each time.
The overall structure of our network mainly consists of a gating network G , a set of n experts { E i } ni =1 , and a shared prediction head H for generating σ and c .
The actual inputs of our network are the positional encoding ( PE( x ) , PE( d ) ) for ( x , d ) as in the vanilla NeRF (Mildenhall et al., 2020).
We omit PE ( · ) in our following equations for simplicity.
Given an input ( x , d ) , we first send x into the gating network and obtain the gate values G ( x ) .
Then, we apply a Top- 1 operation on G ( x ) to determine which expert should be activated.
As shown in
Fig. 2, only one expert ( i.e. Expert 2) is activated.
The point x will be dispatched to this selected expert.
Other experts do not participate in the processing of x .
The output feature of the expert E ( x ) will be multiplied by the gate value corresponding to this expert.
This makes the gating network be trained jointly with the expert networks.
After that, the feature is used to predict the density σ and color c together with d and the appearance embedding AE.
In the next, we first introduce details of the trainable gating and our gating network architecture in Sec. 3.1.
Then we introduce our expert and head network architectures in Sec. 3.2.
We discuss the capacity factor and full dispatch in Sec.
We finally formulate the rendering procedure and our loss functions in Sec. 3.4.
3. S PARSE G ATING IN S WITCH -N E RF
The sparse gating network plays an important role in our Switch-NeRF because it determines the optimization routes of different NeRF experts.
In Switch-NeRF, we only consider one gating network because there is a very large number of 3D points that require gating and dispatching in NeRF optimization.
Multiple gating operations will largely decrease the training and testing speed.
Trainable gating in Switch-NeRF.
The trainable gating in our network follows the mechanism in the MoE method Switch Transformer (Shazeer et al., 2017).
The gate values G ( x ) is a vector of n -dimensions normalized via Softmax, in which G ( x ) i represents the probability of selecting the i -th NeRF expert.
We apply a Top- 1 function on G ( x ) to sparsely select only 1 expert E s from a set of NeRF experts, i.e. { E i } ni =1 , for each 3D point.
The input x will be dispatched into the selected expert E s and obtain an output E s ( x ) .
The final output ˜ E ( x ) the output of E s multiplied by the corresponding gate value:
[Equation]
As the predicted gate values are multiplied to the corresponding outputs of NeRF experts, the gating network can be optimized together with the NeRF experts in the backward pass.
This makes the network able to directly learn scene decomposition during network training.
Our network structure is highly sparse because we use Top- 1 to select only one experts for each sample.
[Table]
> PE( x ) Linear
Linear & ReLU Linear & ReLU
Linear & ReLU Linear & ReLU
Linear & ReLU Linear & ReLU
Linear & ReLU Linear & ReLU
Linear & LayerNorm
Linear
Softmax
Mul
Gate
Expert
Linear & Softplus <latexit sha1_base64="bRgyqLH9q6jwKhziyjrSlljThc4=">AAAB7XicbVDLSgNBEOyNrxhfUY9eBoPgKexKUI8BLx48RDAPSJYwO5lNxsxjmZkVwpJ/8OJBEa/+jzf/xkmyB00saCiquunuihLOjPX9b6+wtr6xuVXcLu3s7u0flA+PWkalmtAmUVzpToQN5UzSpmWW006iKRYRp+1ofDPz209UG6bkg50kNBR4KFnMCLZOavUMGwrcL1f8qj8HWiVBTiqQo9Evf/UGiqSCSks4NqYb+IkNM6wtI5xOS73U0ASTMR7SrqMSC2rCbH7tFJ05ZYBipV1Ji+bq74kMC2MmInKdAtuRWfZm4n9eN7XxdZgxmaSWSrJYFKccWYVmr6MB05RYPnEEE83crYiMsMbEuoBKLoRg+eVV0rqoBpfV2n2tUr/L4yjCCZzCOQRwBXW4hQY0gcAjPMMrvHnKe/HevY9Fa8HLZ47hD7zPH6JPjzY=</latexit>  Linear
Linear & ReLU
Linear & Sigmoid <latexit sha1_base64="r6UGcqpeAxNPrDn8rjg9Nzh0jok=">AAAB8XicbVA9SwNBEJ2LXzF+RS1tFoNgFe5E1DJgY2ERwXxgcoS9zV6yZG/v2J0Tw5F/YWOhiK3/xs5/4ya5QhMfDDzem2FmXpBIYdB1v53Cyura+kZxs7S1vbO7V94/aJo41Yw3WCxj3Q6o4VIo3kCBkrcTzWkUSN4KRtdTv/XItRGxusdxwv2IDpQIBaNopYcu8icMwoxNeuWKW3VnIMvEy0kFctR75a9uP2ZpxBUySY3peG6CfkY1Cib5pNRNDU8oG9EB71iqaMSNn80unpATq/RJGGtbCslM/T2R0ciYcRTYzoji0Cx6U/E/r5NieOVnQiUpcsXmi8JUEozJ9H3SF5ozlGNLKNPC3krYkGrK0IZUsiF4iy8vk+ZZ1buont+dV2q3eRxFOIJjOAUPLqEGN1CHBjBQ8Ayv8OYY58V5dz7mrQUnnzmEP3A+fwAM45E1</latexit
[Figure]
Gating network architecture.
In previous MoE methods (Shazeer et al., 2017; Lepikhin et al., 2021), the gating network is typically a simple linear mapping.
It is reasonable because they usually have multiple MoE layers.
The input of the gating networks in deep MoE layers can be high-level features from previous layers.
Therefore, their gating networks share more information from the main network.
Let the original input be the 3D point x .
Typically x will first go through a sub-network
S with several layers and learn an internal feature S ( x ) .
The real gating operation of the deep MoE layers can be written as:
[Equation]
Therefore, the gating networks in these methods actually share an S network from the main network.
Considering that we only use one gate network in our SwitchNeRF for efficiency, it is necessary to allocate more parameters for our gating network to learn powerful gating.
We can put the gating network closer to the prediction head.
It will share more parameters with the main network.
However, this will also make each sample point share more layers before the gating network.
In this case, the capacity of the whole network shrinks, as the network capacity is controlled by the number of layers of each unshared expert.
Therefore, as shown in the Fig. 3, we put the gating network at the beginning of Switch-NeRF to maximize the layer numbers of the expert networks.
In order to allocate more parameters to the gating network, we use a shallow MLP instead of a linear mapping.
Our gating network G ( x ) consists of 4 Linear layers and 1 LayerNorm as shown in Fig.
Our design balances the efficiency, network sparsity and the number of parameters in the gating network.
3. E XPERT AND H EAD N ETWORKS
NeRF Expert Network.
The set of n NeRF experts { E i } ni =1 in Switch-NeRF contains most of the network parameters.
As shown in Fig.
3, each expert in Switch-NeRF is a deep MLP with a skip connection.
The structure and depth of the experts are aligned with the main structure of vanilla
NeRF.
Each expert only processes a part of the 3D points, determined by the gating network.
The output of the selected expert is multiplied by the corresponding gate value to obtain a feature vector ˜ E ( x ) .
This feature vector is then used to predict density σ and the direction-dependent color c .
By increasing the number of NeRF experts n , we can easily scale the network’s capacity.
Unified NeRF Head.
The head network H is designed for the final predictions of each input sample.
It is shared for all the samples, as shown in Fig.
After obtaining the expert output ˜ E ( x ) , we use a
Linear layer with a Softplus activation (Zheng et al., 2015) to predict σ .
The use of Softplus follows
Mip-NeRF (Barron et al., 2021) for stable prediction of σ .
Then, ˜ E ( x ) goes through a Linear layer and is concatenated with PE( d ) and an appearance embedding AE.
The color c is predicted from the concatenated feature by an MLP.
The appearance embedding AE is a trainable vector to capture image-level photometric and environmental variations (Martin-Brualla et al., 2021).
3. C APACITY F ACTOR AND F ULL D ISPATCH
Capacity factor for training.
In our Switch-NeRF, it is important to efficiently dispatch 3D points to different experts.
Einops-based dispatch (Lepikhin et al., 2021) causes memory overflow due to the large number of 3D points.
In our training, we use the CUDA-based fast dispatch in Tutel (Hwang et al., 2022).
Following previous MoE methods (Lepikhin et al., 2021), we set a capacity factor for each NeRF expert in the training.
It caps the number of sample points dispatched to each
NeRF expert, leading to uniform tensor shapes, balanced computation and communication.
Let B be the batch size of the whole network, C f be a capacity factor and n be the number of NeRF experts.
Then the maximum number of sample points going into each NeRF expert is B e = ceil ( kBC f n ) .
The dispatch with a capacity factor can be called a uniform dispatch.
Fig.
4 shows the uniform dispatch with C f = 1 . 0 .
Overflow points are dropped.
If the capacity is not fully used, it will be zero-padded.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
images rendered by uniform dispatch have apparent artifacts.
A larger capacity factor decreases the dropping ratio but increases the memory and computation.
In our network, we set the capacity factor to 1.0 without requiring extra memory.
We use the Batch Prioritized Routing (Riquelme et al., 2021) to improve the training with lower expert capacity.
Full dispatch for testing.
Previous MoE methods usually use the uniform dispatch for both the training and testing (Lepikhin et al., 2021; Fedus et al., 2022), which inevitably drops sample points.
In our Switch-NeRF, although the uniform dispatch works well in the training, we observe that dropping sample points can significantly decrease the test accuracy.
A possible reason is that we do not use stacked MoE layers and skip connections between them as previous MoE methods usually do, to maintain the network efficiency.
To tackle this problem, we implement an efficient full dispatch strategy based on Pytorch, CUDA, and Tutel.
It can dispatch all the points to their corresponding expert with only slight memory increase.
A description of the strategy is depicted in Fig.
With this efficient full dispatch, we can largely improve the test accuracy.
3. V OLUME R ENDERING AND L OSSES
Our volume rendering procedure follows the vanilla NeRF (Mildenhall et al., 2020).
The training data of Switch-NeRF consists of multiple posed images.
For each pixel in the training images, we back-project a camera ray r into a 3D space.
A set of N samples are sampled along the ray.
For each sample ( x , d ) , our network F Θ predicts a volume density σ and a color c =
( r, g, b ) .
Let δ i be the distance between adjacent points.
The expected color ˆ C ( r ) of this pixel is synthesized along the ray by the volume rendering function in NeRF (Mildenhall et al., 2020):
[Equation]
T
[Equation]
exp
[Equation]
Rendering loss.
The main loss of Switch-NeRF is the rendering loss L r .
After rendering the color ˆ C ( r ) of a ray from our network through Equation 3, we compute L r for supervision:
[Equation]
where R is the set of sampled rays.
C ( r ) is the ground truth color of ray r in the training images.
Auxiliary loss tackling imbalanced optimization.
One problem of training MoE-based networks is that the gating network can favor only a few experts (Shazeer et al., 2017).
The optimization and utilization of NeRF experts will thus be imbalanced.
This could even cause several experts not trained, and then the whole network converges to a sub-optimal solution and cannot scale well through the training since some of the network capacities are not fully utilized.
Following Shazeer et al.
; Lepikhin et al.
(2021), we use an auxiliary loss L a to regularize the gating network and balance the utilization of NeRF experts.
Specifically, following the differentiable load balancing loss proposed in GShard (Lepikhin et al., 2021), we define our auxiliary loss as follows.
Given n
NeRF
experts and a batch B with N sample points, let c i be the number of points dispatched to each expert
E i by Top- 1 .
We first compute the gating values m i distributed to E i with m i
= (cid:80) x ∈B G ( x ) i .
The auxiliary loss L a can be computed as L a = nN 2 (cid:80) ni c
i m i .
This auxiliary loss encourages balanced gating because it is minimized when the dispatching is ideally balanced.
Under the balanced gating, c i and m i are both expected to be N n .
Then (cid:80) n i c
i m i will be N 2 /n .
The loss L a will be 1 .
The total loss L is the weighted sum of the above-mentioned two losses:
[Equation]
where λ is the weight for our auxiliary loss.
We set λ = 5 × 10 − 4 for all our main results and it is sufficient to balance the utilization of experts.
M e g a - N e R F
S w it
c
hN e R F
Building
Sci-Art
Residence
Rubble
Campus
G r ound
T r u t h
[Table]
E XPERIMENTS 4. D ATASETS , M ETRICS AND V ISUALIZATION
Datasets.
We use the Building, Rubble datasets from Mill 19 (Turki et al., 2022) and Residence,
Sci-Art, Campus datasets from UrbanScene3D (Liu et al., 2021) to evaluate our Switch-NeRF.
Each scene contains thousands of high-resolution images.
The camera parameters for all the images are the same as Mega-NeRF (Turki et al., 2022).
These datasets cover large enough areas while they can still be handled by consumer-level workstations with commonly used GPUs.
Metrics.
We use the PSNR, SSIM (Wang et al., 2004)
(both higher is better), and VGG implementation of LPIPS (Zhang et al., 2018) (lower is better) to quantitatively evaluate our results on the novel view synthesis.
The PSNR is to measure the mean squared error between two images in logarithmic space.
The SSIM focuses more on structural similarity.
The LPIPS measures perceptual similarity.
Visualization.
Besides the rendered images, we visualize the 3D radiance fields.
We sample 3D points along rays and use the α = 1 − exp ( − σ i δ i ) as the opacity of each 3D point.
We use Point
Cloud Library (Rusu & Cousins, 2011) to show the color and opacity of each 3D point
. 4. S ETTING
Similar to NeRF++ (Zhang et al., 2020) and Mega-NeRF, the 3D scene space is spilt into a foreground and a background.
We use 8 experts with Top- 1 gating and C f = 1 . 0 for the foreground to end-to-end learn scene decomposition, and one original NeRF in the background.
This differs from
Mega-NeRF using a foreground and a background NeRF for each of its 8 sub-networks, requiring more network parameters.
Each of our NeRF experts contains 7 layers with each layer 256 channels.
We sample 256 coarse and 512 fine points per ray in the main network and 128/256 samples in the background network.
We use 8 NVIDIA RTX 3090 GPUs for distributed data-parallel training and sample 1024 rays for each GPU.
We use Adam optimizer (Kingma & Ba, 2015) and a learning rate decaying exponentially from 5 × 10 − 4 to 5 × 10 − 5 , and use bfloat16 in training and float16 in testing to reduce memory and time.
We train 500k iterations for each dataset and test on validation images.
S w it
c
h - N e R
F
All
Expert Expert
Expert Expert
M e g aN e R F
Centroid
Centroid Centroid
Centroid 5 All
4. B ENCHMARK P ERFORMANCE
The main quantitative results are reported in Table 1 and qualitative results in Fig.
The statistics of NeRF (Mildenhall et al., 2020), NeRF++ (Zhang et al., 2020), SVS (Riegler & Koltun), DeepView (Flynn et al., 2019), and Mega-NeRF are quoted from Mega-NeRF.
The MLP width of NeRF and NeRF++ is both set to 2048 to obtain a similar capacity between Mega-NeRF and Switch-NeRF.
As shown in Table 1, our Switch-NeRF achieves state-of-the-art accuracy on all the datasets in terms of PSNR, SSIM, and LPIPS.
It produces much better accuracy than NeRF and NeRF++.
This confirms the effectiveness of scaling NeRF with sparse Mixture-of-Experts instead of densely expanding the network.
Our method also outperforms the sparse Mega-NeRF, showing the advantage of learning decomposition instead of the hand-crafted scene decomposition.
Fig.
5 presents the rendered images of Mega-NeRF and Switch-NeRF.
Our method can render more tiny structures and details compared to Mega-NeRF.
These main results demonstrate the overall performance of our method.
4. M ODEL A NALYSIS
[Table]
Besides the main results, we perform ablation experiments to analyze different aspects of our model.
The Mill 19 Building scene is used by default.
Decomposition.
We analyze our learning-based decomposition compared to other decomposition methods.
The random decomposition randomly dispatches 3D points into sub-networks.
The distance decomposition dispatches 3D points based on the 3D distances in training and testing.
The network design and the training setup for these variants are the same as Switch-NeRF.
We also include the results of
Mega-NeRF, which uses distances to cluster pixels in training and 3D points in testing.
The results are shown in Table 2.
The random decomposition cannot fully utilize more parameters in the network.
The distance decomposition performs better and is similar to Mega-NeRF.
Our Switch-NeRF with learned decomposition achieves the best accuracy.
NeRF expert specialization.
We visualize the 3D radiance fields of the UrbanScene3D-Sci-Art for
Mega-NeRF and Switch-NeRF in Fig.
The Mega-NeRF partitions the scene into regular parts.
In Switch-NeRF, the experts have been roughly specialized to different semantic parts of the scene.
Experts 1 and 2 focus on different buildings.
Expert 3 focuses on the green grasses and trees on the ground.
Expert 5 specializes on other parts of the ground.
It is clear that our Switch-NeRF can achieve more reasonable and semantically meaningful scene decomposition in a learning way.
[Figure]
[Table] [Table]
with memory, our network uses less testing memory and has much less
[Table] [Table]
Gating network.
We analyze different designs of our gating network.
As shown in Table 2, the
Linear gating just feeds the PE( x ) into a trainable linear layer.
The w/o Norm version does not have a LayerNorm.
We also evaluate the gating without the auxiliary loss L a .
The Linear gating with fewer parameters does not generate satisfactory results.
The MLP+Norm can boost the performance of our whole network.
The network without auxiliary loss
L a does not even converge.
This shows
L a is vital for the success of MoE methods, consistent with the observations in Shazeer et al.
Unified head.
We remove the shared unified head and add a head separately for each NeRF expert.
Table 2 shows that without the unified head the accuracy significantly decreases.
We also observe that in this case, the network is not robust enough to the mixed-precision training and testing.
A possible reason is that the gating value is directly multiplied to the predictions rather than the highlevel features, which makes the gating and prediction unstable in training and testing.
This verifies our motivation to design a unified head for multiple NeRF experts.
Scalability.
Table 3 shows the accuracy and efficiency of Switch-NeRF with different numbers of experts.
With only 4 experts, we already achieve similar accuracy to Mega-NeRF using 8 subnetworks.
Compared to the model with 8 experts, the model with 16 experts increases the performance remarkably, while without a large increase of the memory and time in testing.
Compared to Mega-NeRF with 16 sub-networks, our Switch-NeRF with 16 experts scales much better with a higher accuracy, fewer number of the parameter and less memory footprint in testing.
All of these prove the scalability of our Switch-NeRF.
It can obtain much better results by increasing network capacity while maintaining almost constant computational and memory costs.
Top- k and capacity factor.
Table 4 studies different Top- k and C f when training Switch-NeRF.
Increasing C f from 1.0 to 1.5 or using k = 2 in training can improve the accuracy while increasing the training time and memory, especially for k = 2 .
With our design of Switch-NeRF, Top- 1 with
C f = 1 . 0 already obtains good results with acceptable efficiency, suggesting that it is better to increase C f instead of k because C f in training does not affect testing with the full dispatch.
Efficiency.
In Table 5, we compare the efficiency of Switch-NeRF with Mega-NeRF.
Our model achieves better accuracy using only half of the parameters of Mega-NeRF.
Mega-NeRF uses a background NeRF and different versions of appearance embeddings (AE) for each sub-network, leading to much more parameters.
As Switch-NeRF is trained end-to-end, it only uses one background NeRF and one version of AE.
Our network is more parameter efficient and with better accuracy.
Since it has a gating network and is trained end-to-end, it reasonably requires more floating point operations (FLOPs) for each point and costs slightly more memory and time for training.
Notably, compared to Mega-NeRF, it uses around 20% less testing memory, and only a very minor increase in testing time.
Our Switch-NeRF achieves a good efficiency together with a better accuracy.
Full dispatch.
We test our trained model with a uniform dispatch and our full dispatch discussed in Sec. 3.3.
A capacity factor C f of 2 . 0 and 4 . 0 with Batch Prioritized Routing (Riquelme et al., 2021) is used in the uniform dispatch.
Table 6 shows that the results of the uniform dispatch with
C f = 2 . 0 largely decreases.
Although the results with C f = 4 . 0 improve, they are still far lower than the full dispatch.
Moreover, it costs more time than the full dispatch because of a large C f with zero padding.
Fig.
4 shows that images rendered with the uniform dispatch contain much more artifacts than the full dispatch, proving the effectiveness of the full dispatch in Switch-NeRF.
C ONCLUSION
In this paper, we present Switch-NeRF.
To the best of our knowledge, it is the first sparse large-scale
NeRF with learnable scene decomposition.
We propose an end-to-end mixture-of-NeRF-experts framework to learn scene decomposition jointly with NeRF.
We further design and implement an efficient network architecture and a full dispatch strategy to boost accuracy.
Extensive experiments demonstrate that our network can learn more reasonable scene decomposition and shows state-ofthe-art accuracy on the scene synthesis compared to hand-crafted decomposition methods. A PPENDIX A B UNGEE N E RF D ATASET
We compare Switch-NeRF with BungeeNeRF (Xiangli et al., 2022) on the 56Leonard and
Transamerica datasets provided from BungeeNeRF.
We first analyze the different motivations and focuses between the BungeeNeRF and our Switch-NeRF.
The BungeeNeRF focuses on modeling scenes with multi-scales and different levels of details, while Switch-NeRF targets large-scale scene modeling with large data volumes.
BungeeNeRF has several special designs for training.
It decomposes the training data by predefined scales and uses multi-level supervision for its middle blocks.
It designs special data feeding scheme and model growing training.
We conducted experiments on the BungeeNeRF dataset, and show the performance comparison in Table 7.
It is worth noting that the data volume of BungeeNeRF (around 450 images with 640x360 resolution for each scene) is similar to the vanilla NeRF but is significantly smaller compared to that of Mega-NeRF (the Building dataset with 1940 images of 4608x3456 resolution).
Thus, we use less experts for BungeeNeRF dataset.
As seen in Table 7, our Switch-NeRF with 4 experts can outperform the BungeeNeRF in all scales.
[Table]
B B LOCK -N E RF D ATASET
We train a Switch-NeRF on the San Francisco Mission Bay Dataset of Block-NeRF (Tancik et al., 2022).
It consists of 12,000 images with an image resolution of around 1200x900.
Although the data volume is similar to the datasets used by Mega-NeRF and Switch-NeRF, the training setting of
Block-NeRF is quite different from ours.
The Block-NeRF uses much larger computing and memory resources in its training setting.
It trains each scene block with 32 TPU v3 cores combining offer
GB memory, which is not easy to be followed by the common computer workstations.
The batchsize of Block-NeRF for each block is 16384, significantly larger than the averaging batch-sizefor each sub-network in Mega-NeRF and Switch-NeRF.
The networks of BlockNeRF are trained with full precision while ours with bfloat16 precision.
Besides, an important part of the training data, i.e., the masks for dynamic objects in the scenes used by Block-NeRF, are also not released in the available training data.
The Block-NeRF does not open-source its code, which makes us difficult to align its hyper-parameters on their training data.
Due to our lack of comparable computing resources and the unavailability of the same training data used in Block-NeRF, it is not feasible for us to train a Switch-NeRF based on the same setting as
Block-NeRF, to allow a direct and fair comparison.
We thus train our network on the Block-NeRF dataset with a similar setting used in our paper.
We use 8 NVIDIA RTX 3090 GPUs to train a SwitchNeRF model with 8 experts, and sample 1664 rays for each GPU, resulting in a batch-size of 1664 on average for each expert.
This is far less than the batch-size of 16384 for each Block in Block-NeRF.
The width of each layer is set as 512 as in Block-NeRF.
We also use the positional encoding of
Mip-NeRF as used by Block-NeRF.
We sample 256 points each ray for the coarse and fine networks in training, and 512 points each ray for testing.
We use the validation dataset of Block-NeRF to evaluate the performance with the PSNR, SSIM and LPIPS metrics.
The experimental results are shown in Table 8.
[Table]
C T WO G ATING O PERATIONS
We add an additional ablation study on two gating operations.
We train a Switch-NeRF with 2 gating networks, each acting as a gating operation.
The first gating network is added after the 1 -st linear layer of the main network, and the second gating network is added after the 4 -th linear layer of the main network.
The total number of expert layers is the same as that of the Switch-NeRF usinggating network to make them have similar network capacity.
The results in terms of both accuracy and efficiency are shown in Table 9.
From the results, we can clearly see that two gating operations cost more training time, training memory
[Table]
and testing time, while achieving a similar accuracy to the model with one gating operation.
This suggests that the number of gating operations is not critical for the model performance, and one gating operation already shows sufficient capabilities in the dispatch of points.
[Table]
152s
D M ORE E XPERTS
The expert number directly controls the network capacity of Switch-NeRF.
With large-scale scenes we may need to increase the expert number to get better results.
In this section we add more experiments on the expert number.
In Table 10, the results show that more experts can consistently produce better results for larger-scale dataset.
We additionally visualize the 3D radiance fields of a Switch-NeRF trained with 16 experts on the
UrbanScene3D-Sci-Art dataset in Figure 7.
Compared with the Figure 6 in the main paper, we can see that our network can still roughly specialized to different fine-grained semantic parts of the scene.
[Table]
[Figure]
E F AST R ENDERING
We explore the capability of Switch-NeRF to be integrated into fast rendering techniques based on octrees similar to the interactive exploration used in Mega-NeRF-Dynamic Turki et al.
As the detailed evaluation protocol for the faster rendering is not provided by Mega-NeRF-Dynamic, we thus define a protocol to evaluate for both methods, on the validation images of each dataset.
Specifically, we first convert the trained model into a coarse octree.
Given a view point, we render a coarse image directly from the coarse octree.
Then, we follow the dynamic octree refinement strategy proposed by Mega-NeRF-Dynamic, to refine the octree for the current view point for several rounds by querying the trained model.
Here, we consider 16 rounds.
We finally render a refined image from the refined octree.
We evaluate both Mega-NeRF model and Switch-NeRF model on the same fast-rendering protocol we just described, and report the PSNR and average refinement time of octree in Table 11, and images in Figure 8.
As shown in Table 11, since the octrees before refinement are very coarse, the coarse images for both Mega-NeRF and Switch-NeRF are of low quality.
After the dynamic octree refinement, the quality of images improves.
Since the resolution of refined octrees is still limited, the quality is not comparable to the original model.
However, our Switch-NeRF can obtain similar or better results on the validation datasets compared to MegaNeRF.
It should be noted that the octrees do not handle the background, and thus the rendered images may have black regions in both Mega-NeRF and Switch-NeRF, which decreases the PSNR values.
However, as shown in Figure 8, our method can render good quality of refined images in the foreground.
This shows our Switch-NeRF can be flexibly and effectively integrated into existing faster rendering techniques.
[Table]
M e g a - C o a r s e
S w it
c
h - C o
a r s e
Building
Sci-Art
Residence
Rubble
Campus
G r o u nd T r u t h
M e g a - R e f i n e d
S w it
c
h - R
e f i n e d