Learned Initializations for Optimizing Coordinate-Based Neural Representations

Matthew Tancik; Ben Mildenhall; Terrance Wang; Divi Schmidt; Pratul P. Srinivasan; Jonathan T. Barron; Ren Ng

Source-linked AI summary

Learned Initializations for Optimizing Coordinate-Based Neural Representations

Matthew Tancik, Ben Mildenhall, Terrance Wang, Divi Schmidt, Pratul P. Srinivasan, Jonathan T. Barron, Ren Ng

arXiv:2012.02189v2Dec 3, 2020cs.CV

TL;DR

Coordinate-based neural representations require costly optimization for each new signal, especially when observations are partial or indirect. The paper uses MAML or Reptile to learn signal-class-specific initial weights, finding faster convergence and better reconstruction from fewer observations across several tasks. The approach retains the underlying representation and test-time optimization procedure, but requires example signals and still needs test-time optimization.

Problem
Fitting a coordinate-based network from scratch for every signal is inefficient, and sparse or indirect observations can make accurate recovery difficult.
Method
The paper applies optimization-based meta-learning to learn initial weights for fully connected coordinate-based networks across distributions of signals.
Results
Across 2D images, CT reconstruction, 3D shapes, and scenes, learned initializations improve convergence or reconstruction from fewer observations, including single-view 3D shape recovery.
Takeaways & Limitations
Learned initial weights act as a strong prior while preserving the architecture and test-time optimization procedure.
Takeaways & Limitations
The approach requires a sizable dataset of example signals from the target distribution and still requires some test-time optimization.

Abstract

from arXiv · show

Coordinate-based neural representations have shown significant promise as an alternative to discrete, array-based representations for complex low dimensional signals. However, optimizing a coordinate-based network from randomly initialized weights for each new signal is inefficient. We propose applying standard meta-learning algorithms to learn the initial weight parameters for these fully-connected networks based on the underlying class of signals being represented (e.g., images of faces or 3D models of chairs). Despite requiring only a minor change in implementation, using these learned initial weights enables faster convergence during optimization and can serve as a strong prior over the signal class being modeled, resulting in better generalization when only partial observations of a given signal are available. We explore these benefits across a variety of tasks, including representing 2D images, reconstructing CT scans, and recovering 3D shapes and scenes from 2D image observations.

1. Introduction

Coordinate-based MLPs compactly represent continuous low-dimensional signals, but fitting each new signal from scratch is costly. The paper proposes meta-learned initial weights that improve convergence and generalization while preserving the representation’s flexibility.

1. Introduction: Coordinate-based MLPs map coordinates to signal values, such as pixel coordinates to RGB colors, providing continuous representations without fixed spatial resolution.They can also represent 3D shapes without the cubic storage complexity of voxel grids.
1. Introduction: Optimizing weights for each new signal requires many gradient-descent steps, ranging from seconds for small images to hours for some inverse problems.Existing latent-vector and hypernetwork approaches may restrict representations to a learned latent space.
1. Introduction: Meta-learning learns initial weights across signal classes so new representations converge faster and generalize better from partial observations.The approach applies MAML or Reptile to observations sampled from a fixed underlying class.
1. Introduction: The implementation adds an outer MAML or Reptile loop and later reloads the learned weights in place of standard initialization.The paper emphasizes that this requires only a few extra lines of code and a dataset of training examples.

2. Related Work

Prior work established neural representations for shapes and image-based reconstruction, while meta-learning learned initializations for efficient adaptation. This paper extends learned initialization across images, volumetric data, and 3D scenes.

Neural Representations: Neural representations encode 3D shapes as implicit surfaces and can reconstruct objects from incomplete point clouds, depth scans, or 2D images.Differentiable rendering enables recovery from image observations.
Neural Representations: Coordinate-based networks have been adapted with sine activations or Fourier features, and this work instead learns how fitting many category signals shapes initialization.The learned initialization is intended to reduce optimization steps and observations for new signals.
Meta-learning: Optimization-based meta-learning uses an outer gradient-based loop to find initial weights that efficiently optimize new task instances with standard optimizers.MAML and Reptile are representative methods for this setup.
Meta-learning: MetaSDF showed rapid convergence for signed distance fields; this paper applies learned initialization to a wider range of signal types and studies initialization as a prior.The broader scope includes images, volumetric data, and 3D scenes.

3. Overview

The paper formulates coordinate-based representation fitting as gradient-based optimization and learns initial weights that minimize post-optimization loss across a signal distribution. At test time, standard optimization adapts these weights to previously unseen signals, including inverse problems with indirect observations.

3. Overview: A coordinate-based representation is a fully connected network mapping d-dimensional coordinates to n-channel signal values, optimized to match the target over its coordinate domain.Images map 2D coordinates to colors, while volumetric representations can map 3D locations to color and density.
3. Overview: Direct pointwise observations supervise fitting with an L2 loss, and gradient descent updates the network weights from their initialization.The initial weights are θ0, while θi denotes the weights after i optimization steps.
3. Overview: Different initial weights produce different final weights and approximation errors under a fixed optimization-step budget.The final weights can be written as θm(θ0, T) when emphasizing dependence on initialization and signal.
3. Overview: Indirect observations use a forward model M, such as rendering a 3D object from camera pose p, and require optimizing a loss that compares predicted and observed measurements.This formulation turns recovery into an inverse problem.
3. Overview: Sparse or information-losing observations can prevent close recovery, making strong prior knowledge important for tasks such as reconstructing a 3D object from one view.The limitation arises when M discards too much information or too few observations are supplied.
3.1. Optimizing initial weights: The objective is to learn initial weights that minimize final loss after m optimization steps for unseen signals drawn from the same distribution.The setup assumes a dataset of observations from a particular signal distribution.
3.1. Optimizing initial weights: MAML wraps an outer loop around m-step inner optimization, while Reptile uses a simpler update that avoids second-order gradients.Both methods learn an initialization for new signal instances.
3.1. Optimizing initial weights: The learned weights do not alter the network architecture or test-time optimization procedure and can be optimized for more steps than used during meta-training.MAML is generally stronger at a fixed inner-loop budget, whereas Reptile is less memory-intensive and supports more unrolled steps.

4. Results

Across image, CT, shape, and scene reconstruction tasks, meta-learned initial weights improve test-time optimization by accelerating convergence and enabling reconstruction from partial observations.

Meta-learned initial weights improve performance across 2D image regression, CT reconstruction, 3D object reconstruction, and 3D scene reconstruction tasks.The benefits include faster convergence, better reconstruction from fewer views, single-view geometry recovery, and novel-view rendering from one photograph.
4.1. Tasks: A face-specialized initialization reaches target image quality after 2 gradient steps, while non-meta initializations require 10 to 20 times as many iterations.The learned initialization also speeds convergence on other natural images.
The method requires only a small implementation change: an outer MAML or Reptile loop learns reusable initial weights for later test-time optimization.The approach addresses shortcomings of prior systems to some degree without architectural changes.
4.4. Generalizing from partial observations: For ShapeNet chairs observed from 25 views, meta-learned weights converge faster than Standard initialization, although both methods reach similar quality after many steps.The result indicates a convergence advantage rather than a lasting final-quality difference in this setting.
4.1. Tasks: Learned initializations induce dataset-dependent priors, generalizing best when the test distribution matches the dataset used for meta-learning.Transfer is stronger between similar natural-image datasets and weaker for dissimilar text images.
4.2. Baselines: In sparse-view CT, meta-learned initializations match Standard reconstruction PSNR with half as many views, outperforming the Mean initialization.The learned weights preserve reconstruction quality when supervision is limited.
4.4. Generalizing from partial observations: Meta-learned weights enable single-view ShapeNet reconstruction, whereas a Standard randomly initialized model cannot recover meaningful shape from one view.The prior persists when meta-training uses either multiple views or a single reference view per object.

5. Conclusion

Changing only the initial weights can guide coordinate-based neural representations toward better optimization trajectories, enabling faster convergence or a strong prior without limiting representational power. The approach requires a sizable dataset of example signals and still needs test-time optimization.

Learned initial weights guide networks along better optimization trajectories without changing the architecture or test-time optimization procedure.
The initial weights can provide faster convergence or a strong prior for signals from a target distribution.
The method partially addresses separately optimizing a network for every new signal without limiting representational power.
Beneficial initial weights require a sizable dataset of example signals from the target distribution.
The approach still requires some amount of test-time optimization.

A. Implementation details

The implementation adds an outer meta-learning loop around an existing test-time optimization framework, while experiments tune optimization choices to compare initializations fairly.

Optimization methods and hyperparameters were tuned separately for each experiment using held-out validation sets.
For Adam experiments, the standard parameters were β1 = 0.9, β2 = 0.999, and ϵ = 10−8.

A.1. Image regression

For image regression, MAML and Reptile provide learned initializations for test-time network optimization, with both meta-learned approaches outperforming non-meta-learned weights.

The image-regression network uses five layers with 256 channels, with sine nonlinearities as in SIREN.
MAML meta-learns image-regression weights using two unrolled gradient steps for each image category.
Reptile also outperforms the non-meta-learned weights on image reconstruction.
MAML test-time optimization uses gradient descent at 10−2, whereas Standard, Mean, Matched, and Shuffled baselines use Adam at 10−4.

A.2. CT reconstruction

The CT reconstruction setup uses a ReLU MLP with random Fourier features, and test-time experiments vary supervision views while tuning optimization to address overfitting.

The CT reconstruction MLP has five layers of 256 channels, ReLU nonlinearities, and random Fourier features on input coordinates.
Reptile uses 12 unrolled gradient steps for CT reconstruction and outperformed MAML, which was limited to three steps by memory constraints.
Test-time experiments compare reconstruction quality across different numbers of supervision views.
Models are more prone to overfitting when fewer views are provided, so learning rates and gradient steps are tuned using 16 validation images.

A.3. ShapeNet [3] view synthesis

ShapeNet view synthesis uses a simplified NeRF-style MLP and specifies separate meta-training and test-time optimization settings for scene reconstruction.

A.3. ShapeNet [3] view synthesis: The view-synthesis model is a six-layer MLP with 256 channels per layer, ReLU activations, and positional encoding for each input coordinate.The simplified model uses one network instead of separate coarse and fine networks and omits view-direction inputs.
A.3. ShapeNet [3] view synthesis: The positional encoding uses N = 20 encodings with log-max frequency f = 8, and rendering accumulates 128 samples per ray.
A.3. ShapeNet [3] view synthesis: Reptile meta-training runs for 100K iterations with batches of 128 rays, 32 inner steps per outer step, and scene-dependent Adam learning rates.The outer learning rate is 5 × 10−4 for Chairs and 5 × 10−5 for Lamps and Cars.
A.3. ShapeNet [3] view synthesis: At test time, each experiment uses an inner batch of 64 rays, while optimization algorithms and learning rates vary across initialization and reconstruction settings.For 25-view chair reconstruction, Reptile uses SGD with learning rate 10−1, whereas standard initialization uses Adam with learning rate 10−4.

A.4. Phototourism [16] view synthesis

Phototourism view synthesis retains the simplified NeRF architecture while using an extended Reptile training schedule and comparing against Basic NeRF with one inner step.

A.4. Phototourism [16] view synthesis: Reptile is trained for 150K iterations with an outer batch size of 1 and 64 inner-loop steps using 64-ray batches.Each ray uses 128 volume-rendering samples, and the outer optimizer is Adam.
A.4. Phototourism [16] view synthesis: The Phototourism setup compares Reptile with Basic NeRF using the same configuration, except Basic NeRF performs only one inner step.Trevi is trained for 60K iterations under the Basic NeRF comparison.

B. Weight space interpolation

Weight-space interpolation produces meaningful outputs for networks optimized from meta-learned weights, including image representations and Phototourism landmark reconstructions.

B. Weight space interpolation: Linear interpolation between networks trained on images produces meaningful outputs when the networks use meta-learned weights.The comparison is between networks optimized from standard random initialization and meta-learned initialization.
B. Weight space interpolation: The same interpolation behavior appears for networks reconstructing a Phototourism landmark from different images.The interpolated networks render the scene from a fixed camera pose.
B. Weight space interpolation: The center of the Phototourism interpolation grid is rendered directly from the meta-learned initial weights without test-time optimization, representing an average scene appearance.