KiloNeRF: Speeding up Neural Radiance Fields with Thousands of Tiny MLPs

In plain words

A neural radiance field (NeRF) stores a 3D scene inside a neural network: give it a point in space and a viewing direction, and it returns a colour and how solid that point is. To draw one image you ask the network about hundreds of points along the ray of every pixel. With a big network, that takes close to a minute per image, far too slow for anything interactive.

KiloNeRF cuts the scene into a 3D grid of up to a few thousand boxes and gives each box its own very small network. A small network only has to remember a small piece of the scene, so it can be tiny and still look right. On top of that, the renderer skips points it already knows are empty air and stops walking along a ray once the ray has hit something opaque.

Training thousands of tiny networks straight from photos leaves floating smudges in empty space. So the authors first train a normal NeRF, teach each tiny network to copy that "teacher" inside its box, and only then fine-tune everything on the photos. A small penalty on the part of each network that handles view-dependent colour removes the remaining smudges.

The result, as reported by the authors, is images of about the same quality as the original NeRF, rendered roughly two thousand times faster on a consumer GPU, which turns minute-long renders into tens of milliseconds at 800x800.

TL;DR

KiloNeRF partitions a bounded scene's AABB into a uniform grid (largest side 16 cells) and assigns an independent tiny MLP (4 hidden layers of 32 units) to each cell §3.2 §4.1. Training is three-stage: a NeRF teacher without hierarchical sampling, per-cell distillation on α and RGB, then photometric fine-tuning with L2 weight regularization on the view-dependent layers §3.3 §4.1. Rendering uses an occupancy grid extracted from the teacher for empty space skipping plus early ray termination, and a custom kernel for evaluating thousands of networks on variable-sized batches §3.4 §3.5. Author-reported: comparable PSNR/SSIM/LPIPS to NeRF and NSVF on four datasets, with render times of 26-91 ms versus 37-183 s for NeRF on a GTX 1080 Ti tab. 1. Relevant if you need fast NeRF inference for bounded scenes and can afford a long training pipeline and custom CUDA.

Problem

NeRF evaluates a deep MLP (10 hidden layers, mostly 256 wide) for every one of the W \times H \times K sample points of an image, which makes rendering very slow even on modern GPUs §1 §3.1. Interactive uses such as VR require real-time rendering on the consumer's own device, whereas training can be parallelized on a cluster, so the paper targets rendering speed only §1. Simply shrinking the MLP loses quality because a small network cannot fit a complex scene §1. Assumptions: the scene is bounded and its AABB is known §3.2 §5.

Method

Background (NeRF) §3.1. A network f_\theta maps a positionally encoded 3D point \mathbf{x} and view direction \mathbf{d} to colour \mathbf{c} and density \sigma; only \mathbf{c} depends on \mathbf{d}. For K samples along a ray:

(\mathbf{c}_i, \sigma_i) = f_\theta(\mathbf{x}_i, \mathbf{d}), \quad i = 1, \ldots, K eq. 1

\hat{\mathbf{c}} = \sum_{i=1}^{K} T_i \alpha_i \mathbf{c}_i eq. 2

T_i = \prod_{j=1}^{i-1} (1 - \alpha_j) eq. 3

\alpha_i = 1 - \exp(\sigma_i \delta_i) eq. 4

with \delta_j = \lVert \mathbf{x}_{j+1} - \mathbf{x}_j \rVert §3.1. Eq. 4 is printed without the minus sign in the exponent; the code uses 1 - \exp(-\mathrm{ReLU}(\sigma)\,\delta) run_nerf.py (see Pitfalls). Training samples one image and B of its pixels per step and minimizes

\mathcal{L} = \frac{1}{B} \sum_{b}^{B} \lVert \mathbf{c}_b - \hat{\mathbf{c}}_b \rVert_2^2 eq. 5

Grid of tiny MLPs §3.2. Given AABB bounds \mathbf{b}_{min}, \mathbf{b}_{max} and grid resolution \mathbf{r} = (r_x, r_y, r_z), each cell index \mathbf{i} owns parameters \theta(\mathbf{i}). A point is routed by element-wise spatial binning:

g(\mathbf{x}) = \left\lfloor (\mathbf{x} - \mathbf{b}_{min}) / ((\mathbf{b}_{max} - \mathbf{b}_{min}) / \mathbf{r}) \right\rfloor eq. 6

(\mathbf{c}, \sigma) = f_{\theta(g(\mathbf{x}))}(\mathbf{x}, \mathbf{d}) eq. 7

The grid's largest dimension is 16 and the others are chosen so cells are cubic, giving at most 4096 networks §4.1. In the code, each network receives the point in local coordinates of its own cell, rescaled to [-1, 1] per axis, before positional encoding local_distill.py multi_modules.py; the paper does not mention this.

Tiny network architecture §3.2 fig. 2. A downscaled NeRF MLP: 4 hidden layers of 32 units, no skip connection, view direction fed into the last hidden layer, density independent of direction. All affine layers use ReLU except the RGB output (sigmoid) and the feature vector passed into the penultimate hidden layer (no activation) §3.2. A forward pass needs 1/87 of NeRF's FLOPs fig. 2; one such MLP has about 6k parameters §4.3. Concretely, in the code (num_hidden_layers: 2, late_feed_direction: true) the layer stack is multi_modules.py Synthetic_NeRF_Lego.yaml:

# One KiloNeRF cell network, as built in code  [§3.2] [fig. 2]
# [code: multi_modules.py#L395-L417] [code: multi_modules.py#L459-L503]
# Input:  pos_enc (N, 63) from 10 frequencies; dir_enc (N, 27) from 4
# Output: rgb (N, 3), raw density (N, 1)

def tiny_mlp(pos_enc, dir_enc):
    # first two hidden layers on the encoded position, ReLU  [§3.2]
    hidden = relu(linear_in(pos_enc))                    # (N, 32)
    hidden = relu(linear_hidden(hidden))                 # (N, 32)
    # density head reads the position features only  [§3.2]
    raw_density = linear_alpha(hidden)                   # (N, 1)
    # feature layer has no activation  [§3.2]
    feature = linear_feature(hidden)                     # (N, 32)
    # view direction joins here; this and rgb are "view-dependent"  [§3.2]
    view_hidden = relu(linear_direction(concat(feature, dir_enc)))
    # sigmoid output colour  [§3.2]
    rgb = sigmoid(linear_rgb(view_hidden))               # (N, 3)
    return rgb, raw_density

Three-stage training §3.3 §4.1.

1. Train a teacher NeRF without hierarchical sampling for 600k iterations; NeRF's coarse model is too inaccurate and its fine model too biased towards surfaces to distil from §4.1.
2. Distil: for every cell network, sample random 3D points inside its cell and random view directions on the unit sphere, query student and teacher, convert both densities to α with eq. 4 (to de-emphasize differences between large densities), and minimize an L2 loss on α and RGB §3.3. No rendering and no training images are used. 150k iterations §4.1.
3. Fine-tune on the training images with the photometric loss (eq. 5) for 1000k iterations, using the teacher-derived occupancy grid to skip empty space §3.3 §3.4 §4.1. Add L2 regularization to the weights and biases of the last two layers (direction layer and RGB layer), weight 1e-6 §3.3 §4.1.

# Distillation stage  [§3.3]
# [code: local_distill.py#L300-L347] [code: local_distill.py#L131-L166]

def distill(teacher, students, cells, iterations):
    # pre-sample a fixed pool of points and directions per cell
    # (code: 1.02M per network, 1M train / 20k held out)
    for cell in cells:
        points[cell] = uniform_in_box(cell.min, cell.max, pool_size)
        dirs[cell] = normalize(gaussian(pool_size, 3))
        # teacher sees global coords; colour via sigmoid, density via ReLU
        rgb_t, sigma_t = teacher(points[cell], dirs[cell])
        # fixed "typical" spacing; code: alpha_distance = 0.0211
        alpha_t = 1 - exp(-relu(sigma_t) * alpha_distance)
        targets[cell] = concat(rgb_t, alpha_t)
    for step in range(iterations):
        # code: 128 examples per network per step, all networks at once
        batch = sample_indices(pool_size, per_network_batch)
        local = to_local_coords(points[:, batch], cells)  # in [-1, 1]
        rgb_s, sigma_s = students(encode(local), encode(dirs[:, batch]))
        # code uses leaky_relu on the student density here
        alpha_s = 1 - exp(-leaky_relu(sigma_s) * alpha_distance)
        # L2 on (rgb, alpha); mean per network, summed over networks
        loss = sum_over_networks(
            mean(square(concat(rgb_s, alpha_s) - targets[:, batch]))
        )
        adam_step(loss)

# Fine-tuning stage  [§3.1] [§3.3] [§3.4] [eq. 5]
# [code: run_nerf.py#L981-L1019]

def finetune_step(image, students, occupancy, step):
    # 8192 pixels of one training image per step  [§4.1]
    rays, target_rgb = sample_pixels(image, 8192)
    # K = 384 stratified equidistant samples, no hierarchical pass  [§3.4]
    points, deltas = stratified_samples(rays, 384)
    # empty space skipping: evaluate only points in occupied cells  [§3.4]
    active = occupancy[voxel_index(points)]
    rgb, sigma = query_networks(students, points[active], rays.dirs)
    # alpha compositing of eq. 2-4; no early termination in training  [§3.4]
    pred_rgb = composite(rgb, sigma, deltas, active)
    loss = mean(square(pred_rgb - target_rgb))
    # L2 on direction layer + rgb layer, weight 1e-6  [§3.3] [§4.1]
    # code: sum of (non-squared) L2 norms of those tensors
    loss = loss + 1e-6 * sum(norm(p) for p in view_dependent_params)
    adam_step(loss)
    # NeRF-style exponential decay; code: 0.1 ** (step / 500k)
    set_lr(initial_lr * 0.1 ** (step / 500_000))

Sampling at render time §3.4.

• At most K equidistant points per ray, with stratified sampling; K effectively sets the spacing \delta. No hierarchical sampling §3.4.
• Empty space skipping (ESS): a second, finer uniform occupancy grid (network grid resolution × 16, so at most 256³) stores one bit per cell. Each occupancy cell is filled by querying teacher densities on a 3×3×3 sub-grid inside it and marking it occupied if any density exceeds \tau = 10 §3.4 §4.1 build_occupancy_tree.py. Used for fine-tuning and inference §3.4.
• Early ray termination (ERT): stop sampling a ray once T_i < \epsilon = 0.01; inference only. Implemented by evaluating samples in passes ordered by distance from the camera and dropping terminated rays between passes §3.4 §4.1. In the benchmark render config each pass advances at most 8 samples per ray benchmark.yaml fast_kilonerf_renderer.py.

# Inference with ESS + ERT  [§3.4] [§3.5]
# [code: fast_kilonerf_renderer.py#L76-L174]

def render(rays, networks, occupancy, max_samples=384, eps=0.01):
    transmittance = ones(num_rays)
    color = zeros(num_rays, 3)
    active_rays = all_rays
    depth_index = 0
    while active_rays and depth_index < max_samples:
        # next group of samples, nearest first; code: 8 per pass
        points = next_samples(active_rays, depth_index, per_pass)
        # ESS: drop points in empty occupancy cells  [§3.4]
        points = points[occupancy[voxel_index(points)]]
        # route each point to its cell network  [eq. 6]
        net_ids = floor((points - b_min) / ((b_max - b_min) / grid_res))
        # sort so each network's inputs are contiguous  [§3.5]
        points, batch_sizes = group_by_network(points, net_ids)
        # fused kernel: all tiny MLPs, variable batch sizes  [§3.5]
        rgb, sigma = fused_multi_mlp(networks, points, batch_sizes)
        color, transmittance = composite_inplace(
            color, transmittance, rgb, sigma
        )
        # ERT: retire rays whose transmittance fell below eps  [§3.4]
        active_rays = active_rays[transmittance[active_rays] >= eps]
        depth_index += per_pass
    return color

Implementation §3.5. PyTorch, MAGMA and Thrust plus custom CUDA kernels for sampling, ESS, ERT, positional encoding, network evaluation and α-blending. Each linear layer becomes N independent products of a B_i \times I matrix with an I \times O matrix where B_i varies per network, so ordinary batched matmul does not apply; training uses MAGMA's variable-size batched GEMM, inference a self-written kernel that fuses the whole network. Inputs must first be sorted by network, which is cheap once ESS and ERT have removed most samples §3.5.

Hyperparameters

Results

Author-reported. Setup: NSVF's four bounded-scene datasets with NSVF's bounding boxes, 25 scenes total; baselines are original NeRF (with hierarchical sampling) trained for a long schedule, and NSVF; render time averaged over test images on a GTX 1080 Ti §4.1 tab. 1.

Speedup breakdown on Lego tab. 2:

So ESS/ERT alone give 71×; the smaller networks contribute the remaining factor of roughly 36 tab. 2 §4.3. Ablations on Lego, judged by LPIPS on close-ups fig. 5 §4.3: a single tiny MLP for the whole scene degrades quality dramatically; halving the grid (5×8×5) hurts more than halving hidden units; skipping fine-tuning hurts substantially; training without distillation gives similar detail but free-space artifacts fig. 3; removing L2 regularization also causes free-space artifacts fig. 6. Training time: teacher NeRF ~2 days on a 1080 Ti (for the NeRF baseline), distillation up to 8 h, fine-tuning 17 h §4.1. All evaluated scenes need less than 100 MB of storage §5.

Paper vs code

Checked against creiser/kilonerf at e79af85. The repo has no LICENSE file, so nothing is quoted from it.

• Fine-tuning learning rate: paper 5e-4 §4.1; the paper's fine-tune configs use 1e-3 Synthetic_NeRF_Lego.yaml (the global default is 5e-4 default.yaml, which the teacher config matches).
• Lego network grid: paper 10×16×10 §4.3; config 9×16×10 with occupancy 144×256×160 Synthetic_NeRF_Lego.yaml Synthetic_NeRF_Lego.yaml. Both have 16 on the largest side.
• L2 regularization: the paper says L2 on weights and biases §3.3; code adds λ times the sum of the non-squared 2-norms of each tensor of the direction and RGB layers run_nerf.py.
• Distillation density activation: teacher targets use ReLU before α conversion local_distill.py, the student uses leaky ReLU during distillation training local_distill.py. Not mentioned in the paper.
• Distillation α conversion uses a fixed spacing alpha_distance: 0.0211, not the actual per-ray δ Synthetic_NeRF_Lego.yaml; the paper only refers to eq. 4.
• Local coordinates: every cell network sees its input rescaled to [-1, 1] within its cell local_distill.py multi_modules.py; the paper does not say.
• The distillation script is a more general adaptive k-d tree fitter; with fixed_resolution and max_error: 100000 no cell is ever split, which reduces it to the paper's uniform grid local_distill.py local_distill.py.
• Teacher training uses 1024 rays per batch and center-cropping for the first 10k iterations Synthetic_NeRF_Lego.yaml; the paper only says "NeRF without hierarchical sampling" §4.1.

Pitfalls

• Sign error in eq. 4. As printed, \alpha_i = 1 - \exp(\sigma_i\delta_i) is negative for positive density. Use 1 - \exp(-\mathrm{ReLU}(\sigma_i)\,\delta_i), as the code does eq. 4 run_nerf.py.
• Local vs global coordinates. Students take cell-local coordinates in [-1, 1] while the teacher takes global coordinates local_distill.py local_distill.py. Feeding global coordinates to the tiny MLPs wastes positional encoding frequencies on the whole scene extent.
• "4 hidden layers" counting. The four are two position layers, the linear feature layer and the direction layer; the density head branches off before the feature layer §3.2 multi_modules.py. In the code's config this is num_hidden_layers: 2.
• α for distillation needs a δ. The paper converts densities to α but does not say which δ; the code uses a constant 0.0211 per scene config. The choice decides how much large densities are compressed §3.3 Synthetic_NeRF_Lego.yaml.
• Distillation alone is not enough. Fine-tuning is required for quality; distillation only initializes the weights §3.3 §4.3. Skipping distillation gives free-space artifacts fig. 3, as does skipping the L2 term fig. 6. Reducing the width of the view-dependent layer, as NeRF does, was tried and lost quality §3.3.
• Teacher choice matters. Use a NeRF without hierarchical sampling as the teacher; the coarse model is too inaccurate and the fine model is surface-biased §4.1.
• Occupancy grid comes from the teacher and is frozen. Geometry the teacher misses at density ≤ 10 is skipped during fine-tuning and can never be recovered §3.4 §4.1.
• ERT is inference-only; training renders all non-empty samples §3.4.
• Hidden costs. The full pipeline is ~2 days of teacher training plus up to 8 h distillation plus 17 h fine-tuning on one 1080 Ti §4.1. The code pre-samples about a million points per network for distillation Synthetic_NeRF_Lego.yaml, which takes a lot of host memory when hundreds of networks are processed together.
• Speed depends on custom kernels. Thousands of networks with variable batch sizes cannot use standard batched matmul; without a grouped GEMM (MAGMA vbatched) or a fused kernel and per-network sorting, much of the speedup is lost §3.5.
• Bounded scenes only. A known AABB is required; unbounded scenes need many more networks and memory §3.2 §5.
• Baseline caveat. NeRF render times are for the original NeRF with hierarchical sampling, and the authors' NeRF baseline scores are higher than those NSVF reported, which they attribute to longer training §4.1.

Lineage

• Builds on: NeRF (arXiv:2003.08934): same field, volume rendering and photometric loss, with one large MLP replaced by a grid of tiny ones §3.1 §3.2.
• Compared against: NeRF and NSVF (arXiv:2007.11571); NSVF also uses ESS and ERT, so the speedup over it is mainly due to the smaller networks §4.1.
• Related: complementary methods that reduce samples per ray, AutoInt (arXiv:2012.01714) and DONeRF (arXiv:2103.03231) §5.