WIRE: Wavelet Implicit Neural Representations

In plain words

A neural network can be trained to be a picture, or a 3D shape, or a radiance field: you give it a coordinate such as "pixel (x, y)" and it answers with the colour or density at that point. Such networks are called implicit neural representations. How well they fit depends heavily on the small non-linear function that each layer applies to its inputs.

The common choices are a sine wave (SIREN) or a plain bell curve (Gaussian). A sine fits fine detail well, but it also fits random noise just as well as a real picture, and it needs careful starting weights. A bell curve is more forgiving but blurs edges.

WIRE replaces these with a wavelet: a wave that is multiplied by a bell-shaped window, so it is localised in both space and frequency. Images are naturally described well by such small localised wiggles, so a network built from them fits real pictures quickly and accurately. It also resists noise better, which makes it useful for denoising, filling in missing data and reconstructing images from partial measurements. The paper reports that it is the most accurate of the compared options and that it trains faster, and it works in 3D and for neural radiance fields as well.

TL;DR

• WIRE is an MLP for implicit neural representations whose activation is a complex Gabor wavelet, exp(j w0 x) exp(-|s0 x|^2) §3.2 eq. 2.
• It is meant to fit images well while fitting noise poorly, so early training favours the image §3.3 fig. 2.
• It needs no SIREN-style initialisation and is less sensitive to w0 and s0 than SIREN or Gaussian activations §3.2 §3.4 fig. 3.
• Author-reported wins: image fitting, occupancy volumes, denoising, super-resolution, CT and NeRF with few views §4.
• Relevant if you need a robust continuous signal model for an inverse problem; not relevant if you need a fast general-purpose network for discrete data.

Problem

Implicit neural representations (INRs) fit a signal with an MLP and reach high accuracy with sine or Gaussian activations, but those activations also overfit noise, so the fit is brittle §1 §3.1. Existing INRs also take tens of seconds to fit high-dimensional data, are weak at denoising and super-resolution, and need many measurements §1. Prior work (deep image prior, SIREN, Gaussian and positional-encoding variants, MFN) is reviewed in §2.

The paper's assumption is that the signal is an image-like function, so a basis of spatially and spectrally compact atoms (wavelets) should give a better inductive bias than sines §1 §3.2.

Method

Setup. An INR is an M-layer MLP F_theta mapping coordinates to signal values. Each hidden layer applies a linear map and a nonlinearity sigma §3.1 eq. 1:

\mathbf{y}_m = \sigma(W_m \mathbf{y}_{m-1} + \mathbf{b}_m) eq. 1

Activation. WIRE uses the continuous complex Gabor wavelet, with frequency w0 and width s0 §3.2 eq. 2:

\sigma(x) = \psi(x;\omega_0,s_0) = e^{j\omega_0 x}\, e^{-|s_0 x|^2} eq. 2

The first-layer output is a set of shifted and scaled copies of this mother wavelet, with the scales and shifts set by W_1 and b_1 §3.2 eq. 3.

Complex arithmetic. Weights and activations are complex, so that phase is kept through the network. The real signal is the real part of the final output; the imaginary part is discarded §3.2.

Why a Gabor wavelet. The complex exponential gives the frequency selectivity of SIREN, and the Gaussian window gives the spatial compactness of the Gaussian activation. Because the window already makes each layer spatially compact, the default uniform random initialisation works for any w0, s0 §3.2 app. A.1.

Implicit bias. In the lazy-training (NTK) regime, gradient flow with early stopping fits the image at time t as in eq. 4 §3.3. The authors compare the NTK eigenfunctions of several activations on noisy Tiny ImageNet images (64 x 64 x 3, noise std 0.05) and on 24 Kodak images (768 x 512 x 3) trained with ordinary gradient descent §3.3 fig. 2. The claim is that WIRE converges to the image before the noise.

F_{\theta_t}(\mathbf{x}) = [(I - e^{-tK}) g](\mathbf{x}) eq. 4

Multi-dimensional WIRE. For a 2D or higher spatial localisation, the authors add D_m - 1 Gaussian windows, one per extra coordinate, applied to the same layer §3.5 eq. 5:

\begin{aligned} \mathbf{y}_m = {} & \psi(W_m^{(1)}\mathbf{y}_{m-1} + \mathbf{b}_m^{(1)};\, \omega_0, s_0) \\ & \cdot \exp\!\Big(-\sum_{k=2}^{D_m} \big|s_0 (W_m^{(k)}\mathbf{y}_{m-1} + \mathbf{b}_m^{(k)})\big|^2\Big) \end{aligned} eq. 5

The resulting first-layer units are compact along both axes §3.5 fig. 11.

Alternate real form. When complex weights are not available, the real part (or imaginary part) of the wavelet, sin(w0 x) exp(-(s0 x)^2), can be used. Setting s0 = 0 gives the SIREN sine, and w0 = 0 gives the Gaussian §3.4, "Alternate forms of WIRE". A constant-Q variant sets w0 s0 = Q §3.4.

Training. Adam, PyTorch, an l2 loss between the forward-modelled output and the measurements, no other regularisation, and a learning rate that is decayed to 0.1 times its initial value at the end of training §4.

# WIRE activation, complex-valued  [§3.2] [eq. 2]
# Input:  lin (N, H) complex or real pre-activations
# Output: activations (N, H) complex

def gabor(lin, omega_0, scale_0):
    # e^{j w0 x} times the Gaussian window e^{-|s0 x|^2}  [§3.2] [eq. 2]
    return exp(1j * omega_0 * lin - abs(scale_0 * lin) ** 2)

# WIRE forward pass  [§3.1] [§3.2] [code: modules/wire.py#L86-L91]
# Input:  coords (N, D_i) in [-1, 1]
# Output: real signal (N, D_o)

def wire_forward(coords, first, hidden_layers, final, omega_0, scale_0):
    # First layer: real weights, complex output  [§3.2] [eq. 3]
    hidden = gabor(first(coords), omega_0, scale_0)      # (N, H)
    # Hidden layers: complex weights and activations  [§3.2] [eq. 1]
    for layer in hidden_layers:
        hidden = gabor(layer(hidden), omega_0, scale_0)  # (N, H)
    # Complex linear read-out; keep only the real part  [§3.2]
    output = final(hidden)                               # (N, D_o)
    return real(output)

# Fitting one signal by gradient descent  [§4]
# Input:  measurements y, forward operator A, coords, model
# Output: fitted parameters

def fit_wire(measurements, coords, model, steps, lr):
    # Adam, l2 loss, no other regularisation  [§4]
    optimizer = Adam(model.parameters(), lr=lr)
    for step in range(steps):
        # Forward model: the INR is the unknown signal  [§4.2]
        prediction = A(model(coords))
        loss = mean((prediction - measurements) ** 2)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # Learning rate decays to 0.1 x its start by the last step  [§4]
        scheduler.step()

Step count, batch size, the learning-rate schedule shape and the coordinate grid are not stated in the paper; the code choices are given under Paper vs code.

Hyperparameters

Results

Author-reported. Setup: Kodak images, a 3D occupancy volume (Thai statue), DIV2K images, the synthetic NeRF drums scene, and Tiny ImageNet for the NTK experiment; baselines are SIREN, Gaussian, ReLU with positional encoding, MFN, and DIP for 2D inverse problems §3.3 §4 §4.1 §4.2.

Paper vs code

Official repository: https://github.com/vishwa91/wire, at code_commit bf95232e0f60434bcbd9b4398ef4c11490832526. It is linked from the authors' project page (https://vishwa91.github.io/wire, given on the paper's title block p. 1), and the owner name matches the first author. The code is licensed under MIT (copyright 2022 Vishwanath Saragadam) LICENSE. The NeRF training itself uses torch-ngp, a third-party repository that is not in this code; it is not part of this card's code.

Searched: the paper's project-page link (found the repository), then the paper's own links. No GitHub search was needed.

Differences between paper and code:

• Complex weights. The paper says the weights and outputs are complex §3.2. The code makes the first layer real-valued and only the hidden and output layers complex wire.py.
• Width reduction. The paper reduces the width by sqrt(2) for WIRE §4.1 app. A.5. The 1D code uses int(width / sqrt(2)) wire.py, so the 300-feature setting becomes 212 features. The 2D code halves the width instead wire2d.py, and its comment says the parameter count is reduced by 4 wire2d.py. The paper gives only the sqrt(2) rule for 1D and says the 2D and 1D parameter counts match §4.4.
• Occupancy hyperparameters. The paper gives w0 = 20 and s0 = 10 for WIRE in both image and occupancy experiments §4.1. The occupancy script uses w0 = 10 and s0 = 40, with a comment that says these work best for 3D occupancies wire_occupancy.py. The script also decays the learning rate to 0.2x, not 0.1x wire_occupancy.py, and clips the volume to its bounding box wire_occupancy.py.
• Denoising w0, s0. The paper does not give w0 and s0 for the denoising experiment §4.2. The script uses w0 = s0 = 5 wire_image_denoise.py, with a comment that suggests 4 and 4 for denoising, and 20 and 30 for representation.
• Denoising photon count. The paper gives a maximum of 30 photons §4.2 and 50 in the sweep §3.4. The script uses tau = 30 wire_image_denoise.py.
• Readout noise units. The paper calls the readout term a "readout count of 2" §4.2. The script comment calls it dB wire_image_denoise.py, but measure adds Gaussian noise in electron counts utils.py. The code matches the paper's reading.
• CT. The paper says 3 hidden layers and 256 features app. A.4. The script uses 2 hidden layers, 300 nominal features, w0 = s0 = 10 and 5000 iterations wire_ct.py. Its 2D model halves the width.
• Multi-image SR. The paper says 3 hidden layers and 256 features app. A.4. The script uses 2 hidden layers and 256 features, with w0 = 10, s0 = 5, tau = 1000 and 2000 iterations wire_multi_sr.py.
• Single-image SR. The paper gives no w0, s0 or learning rate §4.2. The script uses w0 = 8, s0 = 6, learning rate 1e-2 and 2000 iterations wire_SISR.py. These values are code-only.
• NeRF width. The paper says 4 layers of width 128 for all nonlinearities §4.3 and 4 layers of 182 features, 128 for WIRE, in the appendix app. A.5. The NeRF code is not in this repository.
• Initialisation. modules/wire.py has no custom initialisation, so the default PyTorch uniform init applies wire.py. This matches the default in app. A.1.

No discrepancies found in: the complex Gabor formula wire.py eq. 2; the real-part output wire.py §3.2; the 2D window wire2d.py eq. 5.

Pitfalls

• Unstated settings. The paper gives no learning rate, step count, coordinate normalisation or w0 and s0 for denoising, super-resolution and CT. The code supplies values (above), so reproduce from the code or sweep them yourself.
• Tuning on the test image. Parameters were "chosen to obtain fastest approximation rate" §4.1 §4.3. No validation split is stated for the image experiments. Treat the numbers as best-case.
• Single images. Most inverse-problem results use one image or one volume §4.2 app. A.4. The paper does not show generalisation to other images.
• Large s0. The Gaussian window exp(-|s0 x|^2) is nearly zero far from the origin, so with large s0 (e.g. 40 in the occupancy script) gradients may vanish in those regions. This is inferred from the formula and the code; the paper does not report it.
• Cost. The paper notes that many layers are computationally expensive and unstable app. A.3. Complex arithmetic and the 2D variant (two linear maps per unit) cost more per step; no timing is reported for this.
• NTK argument. The implicit-bias evidence uses the empirical NTK on small images (64 x 64) and gradient flow §3.3. The argument is that early training decides the bias, which the authors argue but do not prove.
• Figure and equation citations. Figure numbers are in source order; appendix numbers follow the appendix section labels (A.1 to A.5).

Lineage

• Builds on: SIREN (sine activation) §2 §3.1; Gaussian INR nonlinearity §2; positional encoding and ReLU §2; MFN (Gabor-filter multiplicative networks) §2 §4.1; Gabor and wavelet theory §1 §2.
• Compared against: SIREN, Gaussian, ReLU with positional encoding, MFN, and DIP for 2D inverse problems §4 §4.2.
• Superseded by: not known from content already in this wiki.