Editing Conditional Radiance Fields

In plain words

A neural radiance field stores a 3D scene inside a neural network, so you can render it from any camera. Normally one network holds one scene. This paper trains one network on many chairs, or many cars, so it learns what a chair usually looks like. A user then draws rough strokes on one picture of a chair and says "make this seat red" or "remove this arm". The network carries that change through the whole 3D part, and the change still shows from other angles.

Editing the whole network each time is slow. The authors found that the late layers hold most of what a colour or a part shape depends on. So each edit changes only those layers, which keeps the edit quick and leaves the rest of the object alone.

TL;DR

Edits to a conditional radiance field are made by optimising a few late layers against the user's scribbles, with a penalty that keeps the weights close to the original. The paper reports better colour and shape edits than finetuning all weights, a single-instance NeRF, or a GAN editing method (Model Rewriting), on three shape datasets. Skip it if you need edits in under a second: the shape edit takes roughly 37 s in the authors' own table.

Problem

A NeRF is optimised for one scene and gives no handle for editing. Editing needs a region-level change that carries across views, and it has to run interactively. Changing every weight is slow and can alter unedited parts §1 §3.2. A single-instance NeRF cannot generalise an edit to other views, and 2D GAN editing does not keep 3D shape consistent §4.2 §4.3.

Method

The method has two parts: a network built to separate shape from colour, and an edit procedure that picks which weights to update and what loss to minimise.

Network §3.1 fig. 2. A positional encoding of the 3D point feeds three MLPs. A shared shape MLP sees only the encoding. An instance shape MLP also sees the instance shape code. Their outputs are summed and passed to a fusion MLP. A linear head reads density from the fusion output §3.1. The colour path reduces the fusion output to an 8-dimensional bottleneck, which is combined with the encoded view direction and an embedding of the colour code, then passed through a two-layer MLP app. A.

# Conditional radiance field, one point and view  [§3.1] [app. A]
# Inputs: x (3,), d (2,), shape_code (32,), color_code (32,)
# Outputs: sigma (scalar), rgb (3,)

def conditional_field(x, d, shape_code, color_code):
    # Encode the position once; every shape branch reads it  [§3.1]
    pos = positional_encoding(x)
    # Shape-independent features, shared by all instances  [§3.1]
    shared = shared_shape_mlp(pos)
    # Instance features, conditioned on this instance's shape code  [§3.1]
    instance = instance_shape_mlp(concat(pos, shape_code))
    instance = instance_to_fusion(concat(pos, instance))
    # Add the shared features, then fuse them  [§3.1]
    # [code: models/base_model.py#L74-L125]
    fused = fusion_mlp(relu(instance + shared))
    # Density is a linear read-out of the fused features  [§3.1]
    sigma = density_linear(fused)
    # 8-dimensional colour feature for the colour branch  [app. A]
    colour_feature = bottleneck_linear(fused)
    # Colour branch sees the view direction and a colour-code embedding  [§3.1]
    colour_embedding = colour_code_mlp(color_code)
    rgb = rgb_mlp(concat(colour_feature, encode(d), colour_embedding))
    return sigma, rgb

Rendering §3 eq. 1. Each pixel is rendered by alpha compositing N_c samples along the ray, with alpha_i = 1 - exp(-sigma_i delta_i) eq. 1. Samples are ordered from near to far.

Training app. A. A coarse and a fine network are trained together on batches of rays drawn from one instance at a time. The loss is squared error on both outputs app. A. Codes and networks are optimised jointly. The authors train until convergence, about 1M iterations on average.

Edit choice §3.2 fig. 2. Three options are compared. Updating only the codes is cheap but gives low-quality edits. Updating the whole network is slow and changes unedited regions. The hybrid the authors keep updates only the colour branch and colour code for colour edits, and only the fusion and density layers for shape edits. The paper reports 972 s to 260 s for colour (3.7x) and 1,081 s to 342 s for shape (3.2x) §3.2; see Pitfalls for the conflict with Table 3.

Ray subsampling and feature caching §3.2. Each edit step uses a random subset of user rays: 64 for colour, 8,192 for shape. Outputs of layers that the edit does not touch are cached: per displayed view for rendering, which cuts 256x256 rendering from 6.2 s to under 0.8 s, and during optimisation as well §3.2.

Colour edit loss §3.3 eq. 2 eq. 3. Given target colours on a foreground scribble and original colours on a background scribble, the reconstruction loss is the summed squared error between the composited colours and those targets eq. 2. A regulariser penalises the squared change of weights from the original model, and the colour loss adds it with lambda_reg = 10 eq. 3.

Shape removal §3.4 eq. 4 eq. 5. The user scribbles over the part to remove. The foreground rays are whitened. A density entropy term pushes predicted densities on those rays toward zero. The density vector is first normalised to unit length eq. 4. The total loss adds the reconstruction, the density term with lambda_dens = 0.01, and the regulariser with lambda_reg = 10 eq. 5. Only the density and fusion layers are optimised.

Shape addition §3.4 eq. 6. The user picks an instance to copy a part from and pastes its shape code into a region of the target view. The composite image is the training target. The loss is reconstruction plus the regulariser with lambda_reg = 10 eq. 6. Appendix C gives a faster density matching variant.

Real image §4.5. Editing a single photograph first fits the shape and colour codes with the MLP weights fixed, then fits all parameters jointly. The viewpoint is picked by hand from the training set §4.5.

Hyperparameters

Results

Author-reported. The main metrics are PSNR and LPIPS on held-out views.

View reconstruction and ablation tab. 1, PhotoShapes and Aubry chairs:

The text adds: training on all 626 PhotoShape instances gives 35.79 PSNR, or 33.91 without the shared branch §4.1. SSIM versions of the tables are in app. B.

Colour editing tab. 2, averaged over three edits:

Shape editing tab. 3, averaged over three edits:

Code swapping §4.4 fig. 5 shows that changing one code leaves the other property unchanged. The architecture enforces this for shape, because density does not depend on the colour code §4.4. Real-photo editing §4.5 fig. 6 is qualitative only.

Paper vs code

The official repository is stevliu/editnerf (MIT), linked in the paper's introduction §1.

• Editing learning rate. Paper: 1e-2 app. A §4. Code: LR = 0.001 editingapp.py.
• Training learning rate. Paper: 1e-4 app. A. Code default: 5e-4 inputs.py. None of the repo's configs sets it, and I did not find a training script that does.
• Addition iterations. Appendix says editing uses 100 iterations throughout app. A. Code uses 10,000 for addition editingapp.py. The paper gives no count for addition (it calls the method slow §3.4).
• Regulariser scale. Paper: the sum of squared weight differences, times lambda_reg = 10 §3.3 eq. 3. Code: 10 times the sum over weight tensors of the per-tensor mean squared difference editingapp.py.
• Reconstruction scale. Paper: a sum over rays eq. 2. Code: a mean over pixels run_nerf_helpers.py editingapp.py. Only the relative weight of the terms changes.
• Removal entropy. Paper: densities normalised to unit length §3.4. Code: the compositing weights are used directly, with no normalisation step I could find, averaged over the batch editingapp.py.
• Colour input. Paper: the fused features are concatenated with the colour code §3.1. Code: an 8-dim bottleneck from the fused features is combined with an embedding of the colour half of the code base_model.py.
• Code size. Consistent: style_dim = 64 is split into two 32-dim codes base_model.py config.txt.
• Coarse and fine samples. Paper gives no counts. Configs use 64 coarse and 128 fine samples config.txt, while the argparse default for fine samples is 0 inputs.py.
• Real images. The repository's script uses LBFGS for code-only fitting, 1,000 steps, then 10,000 steps, with 512 rays reconstruction_experiments.sh. The paper gives none of these numbers.
• View-direction regulariser. Paper: K = 64 app. A. The CARLA config sets N_viewdirs_reg = 64 config.txt. I did not trace the loss term itself.

Pitfalls

• Shape edit timing conflicts. Section 3.2 gives 342 s for the hybrid shape edit and 1,081 s for the whole network §3.2. Table 3 gives 37.4 s and 66.4 s for the same two conditions. The paper does not explain the difference, so do not rely on either as a speed target without measuring.
• Learning rate for edits. The paper says 1e-2, and the UI code uses 1e-3. If you port the method, pick one and test it. The code value is the one used by the interactive app.
• Subsampling and quality. The text says subsampling "preserves edit quality" §3.2. Appendix B says that adding a full-region optimisation gives an average PSNR of 34.49 over the three colour edits app. B. It does not say which setting that number belongs to, and no table reports it.
• Addition is unbounded in the paper. Section 3.4 says the method is slow and gives no count. Appendix A says edits take 100 steps. The code uses 10,000. Budget for it.
• Removal under occlusion. The main removal method assumes nothing sits behind the part. Appendix C handles occlusion by zeroing the first density mode along each foreground ray, which the authors say can be slow and may leave artefacts app. C.
• Interactivity. The discussion says a shape edit takes over a minute to give feedback, and that rendering dominates the cost §5. Table 3 gives 37.4 s for the same kind of edit. The paper does not reconcile the two.
• Generalisation. The authors say the method fails on instances that are very different from the rest of the class §5. The paper does not test edits beyond its three datasets.
• Single-view data. On CARLA, one view per instance can cause colour inconsistency across views. The fix is the view-direction regulariser app. A.
• Near and far bounds. The paper does not state them. Configs use 0 and 10 config.txt. Check them against your data.
• Fine network state. Edits keep the coarse network fixed and change only the fine network app. A. The paper does not test changing the coarse network too.
• Faster addition variant. The density-matching cross-entropy loss appears only in the appendix app. C. The main shape-addition loss has no density term.

Lineage

• Builds on: NeRF (Mildenhall et al., 2020), for the volume-rendering and positional-encoding design §3 app. A. Also GRAF and pi-GAN for category-level radiance fields with latent codes §2.
• Compared against: a single-instance NeRF; Model Rewriting (Bau et al., 2020), a 2D GAN editing method §4.2; and naive whole-generator finetuning app. B.
• Superseded by: not known from content already in this wiki.