Neural network Gaussian process

Overview

A neural network with randomly drawn weights and biases, in the limit of very many channels per layer, defines a Gaussian process over its inputs. The network's kernel, the covariance between the outputs for two inputs, can then be computed without training any weights 1808.05587 §2 1808.05587 §3. The shared invariant is that, under the prior, only the variances at corresponding locations and the covariance between two inputs at each layer are needed. The kernel is a recursion over layers, and the Gaussian process posterior then gives an exact prediction 1808.05587 §3.2 1808.05587 §4.

The original argument covers fully connected networks 1808.05587 §1. The paper under this slug extends it to convolutional and residual networks, which use translation-invariant filters and shortcut connections. The kernel for these needs only the three kernel entries at each location, not the covariance between all pairs of locations 1808.05587 §3.2.

Variants

All three variants come from one paper and share the same recursion. They differ in layer type, depth and skip pattern, and the paper's hyperparameter table sets each one 1808.05587 tab. 2.

Convolutional kernel (ConvNet GP)

A stack of convolutional layers with a ReLU nonlinearity and no skip connections. The paper's setting uses 7 layers with 7×7 filters 1808.05587 tab. 2. Its reported MNIST test error is 1.03% with 27 kernels in the random search 1808.05587 tab. 1. Card: Deep Convolutional Networks as shallow Gaussian Processes.

Residual convolutional kernel (Residual CNN GP)

The same convolutional recursion with a skip connection at every layer, with 4×4 filters and 9 layers 1808.05587 tab. 2. The paper reports 0.93% test error on MNIST 1808.05587 tab. 1. Card: Deep Convolutional Networks as shallow Gaussian Processes.

Deep residual kernel (ResNet GP, 32 layers)

A 32-layer residual kernel with 3×3 filters, a skip every two layers, and mostly stride-1 layers 1808.05587 tab. 2. Its bias variance is zero and its weight variance is fixed rather than searched 1808.05587 tab. 2 1808.05587 §4. This is the most accurate kernel in the paper: 0.84% test error and 0.68% validation error on MNIST 1808.05587 tab. 1. The paper reports the cost of computing its kernel matrix and the validation and test kernels as 3 h 40 min on two Tesla P100 GPUs 1808.05587 §4. Card: Deep Convolutional Networks as shallow Gaussian Processes.

Efficiency

The cost of computing the kernel for one pair of inputs is about one forward pass of a network with one filter per layer 1808.05587 §3.3. For N training examples, L layers and input dimension D, the paper gives O(N²·L·D) for the kernel matrix and O(N³) for inverting it 1808.05587 §4. For the MNIST runs, the kernel matrix dominates the cost: inversion and error evaluation took about 43 s on one GPU against the 3 h 40 min for kernel computation 1808.05587 §4.

When to use which

Implementation sketch

Pseudocode of the common core for the convolutional recursion. The three kernel entries per location are propagated through the layers. The ReLU expectation is written as a helper so that the constant can be chosen explicitly (see Contradictions).

# Convolutional NNGP kernel for one pair of inputs x, x2
# Input: images (C0, H0*W0), number of layers L, sigma_w2, sigma_b2
# Output: scalar kernel value k(x, x2)

def nngp_conv_kernel(x, x2, L, sigma_w2, sigma_b2):
    # First layer: average over input channels, then sum over each filter
    # patch; this gives the three kernel entries at each location [arxiv:1808.05587 eq. 10]
    k_xx = sigma_b2 + sigma_w2 * patch_sum(mean_channels(x * x))
    k_yy = sigma_b2 + sigma_w2 * patch_sum(mean_channels(x2 * x2))
    k_xy = sigma_b2 + sigma_w2 * patch_sum(mean_channels(x * x2))
    for _ in range(L):
        # Expected product of post-nonlinearity activations for the pair  [arxiv:1808.05587 eq. 13]
        v_xy = relu_expectation(k_xy, k_xx, k_yy, scale=scale_choice)
        # Propagate each entry through the next filter  [arxiv:1808.05587 eq. 11]
        k_xx = sigma_b2 + sigma_w2 * patch_sum(k_xx)
        k_yy = sigma_b2 + sigma_w2 * patch_sum(k_yy)
        k_xy = sigma_b2 + sigma_w2 * patch_sum(v_xy)
    # Read-out: one number per pair; normalisation not stated [arxiv:1808.05587 alg. 1]
    return mean(k_xy)

For residual kernels, a skip connection adds the kernel stored from the layer it spans back, in place of the usual plain update 1808.05587 eq. 14 1808.05587 eq. 15. Pooling is not part of the method. Average pooling would need covariances between locations within a filter, which the efficient form avoids, so the paper ends with a dense layer instead 1808.05587 §4.

Prediction is exact Gaussian process regression on the training kernel matrix, with one-hot targets in {−1, +1} and the largest output as the class 1808.05587 §4. The kernel itself has no gradient-trained parameters, so no optimiser or learning rate applies 1808.05587 §4. Hyperparameters were chosen on a validation split, with no retraining afterwards 1808.05587 §4.

Contradictions

• Nonlinearity constant. The text defines the nonlinearity as max(0, x) 1808.05587 §3.2. The closed form in 1808.05587 eq. 13 matches √2·max(0, x) instead, and a card check finds the paper does not say which the ResNet kernel uses. Plugging plain ReLU into 1808.05587 eq. 13 halves the expected products. See Deep Convolutional Networks as shallow Gaussian Processes for the code's choice. Unresolved.
• Weight-scale note for the ResNet. Table 2 says the ResNet weight variance matters only for numerical stability, because a plain ReLU stack can absorb it 1808.05587 tab. 2. With residual sums in 1808.05587 eq. 15, the branch scale changes the balance against the shortcut, so this claim does not hold for the residual recursion. The card finds the same. Unresolved.
• Output normalisation. The paper's read-out is one location with no stated normalisation 1808.05587 alg. 1. A constant factor on the whole kernel leaves a noiseless GP argmax unchanged, but a per-layer factor does not. The paper does not say which is used.
• Search size. The paper reports 27 kernels per row 1808.05587 tab. 1. The card notes the run configuration in the public code does not obviously match this count; see Deep Convolutional Networks as shallow Gaussian Processes.

Papers

• Deep Convolutional Networks as shallow Gaussian Processes: gives the convolutional and residual kernels, their recursion, and MNIST results with the ResNet kernel at 0.84% test error.