[Paper Review] Fast Neural Kernel Embeddings for General Activations

In‐Su Han, Amir Zandieh+4 more|arXiv (Cornell University)|Sep 9, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper introduces a fast, general-purpose method to compute Neural Network Gaussian Process (NNGP) and Neural Tangent Kernel (NTK) matrices for arbitrary smooth activation functions using truncated Hermite expansions. By expressing kernels via series expansions of activations and applying sketching techniques, the approach achieves up to 106× speedup over exact computation while maintaining high accuracy, enabling scalable kernel-based analysis across diverse architectures and datasets beyond ReLU.

ABSTRACT

Infinite width limit has shed light on generalization and optimization aspects of deep learning by establishing connections between neural networks and kernel methods. Despite their importance, the utility of these kernel methods was limited in large-scale learning settings due to their (super-)quadratic runtime and memory complexities. Moreover, most prior works on neural kernels have focused on the ReLU activation, mainly due to its popularity but also due to the difficulty of computing such kernels for general activations. In this work, we overcome such difficulties by providing methods to work with general activations. First, we compile and expand the list of activation functions admitting exact dual activation expressions to compute neural kernels. When the exact computation is unknown, we present methods to effectively approximate them. We propose a fast sketching method that approximates any multi-layered Neural Network Gaussian Process (NNGP) kernel and Neural Tangent Kernel (NTK) matrices for a wide range of activation functions, going beyond the commonly analyzed ReLU activation. This is done by showing how to approximate the neural kernels using the truncated Hermite expansion of any desired activation functions. While most prior works require data points on the unit sphere, our methods do not suffer from such limitations and are applicable to any dataset of points in $\mathbb{R}^d$. Furthermore, we provide a subspace embedding for NNGP and NTK matrices with near input-sparsity runtime and near-optimal target dimension which applies to any \emph{homogeneous} dual activation functions with rapidly convergent Taylor expansion. Empirically, with respect to exact convolutional NTK (CNTK) computation, our method achieves $106 imes$ speedup for approximate CNTK of a 5-layer Myrtle network on CIFAR-10 dataset.

Motivation & Objective

• To overcome the computational bottleneck of exact NNGP and NTK computation for general activation functions, which are limited to ReLU and a few others due to intractable kernel expressions.
• To extend infinite-width kernel theory beyond ReLU by enabling exact and approximate computation of neural kernels for arbitrary smooth activations.
• To develop a sketching-based method that achieves near-input-sparsity runtime and near-optimal target dimension for NNGP and NTK matrices across general homogeneous dual activation functions.
• To empirically validate the method on real-world datasets and show that kernel-informed activations can improve finite-width network performance.

Proposed method

• The method expresses any smooth activation function as a truncated Hermite series, enabling exact computation of the dual kernel via coefficients of the expansion.
• It derives closed-form expressions for the dual kernel of monomial, Gabor, and ELU activations using hypergeometric functions and special function identities.
• For general activations with unknown exact duals, it uses Hermite expansion to approximate the kernel with controllable error.
• A sketching algorithm is proposed to approximate full NNGP and NTK matrices with runtime near-input-sparsity and target dimension close to optimal.
• The approach does not require data to lie on the unit sphere, making it applicable to arbitrary datasets in R^d.
• The method leverages the fact that Hermite polynomials serve as random features for monomial kernels, enabling efficient kernel computation.

Experimental results

Research questions

• RQ1 Can we compute exact NNGP and NTK kernels for a broad class of non-Relu activation functions such as GeLU, sinusoidal, and ELU?
• RQ2 How can we efficiently approximate neural kernels for general activations without relying on Monte Carlo sampling or exact recursive computation?
• RQ3 Can we design a sketching-based algorithm that achieves near-input-sparsity runtime and near-optimal dimension for NNGP and NTK matrices across diverse activation functions?
• RQ4 To what extent do kernel-informed activations derived from known kernel forms improve generalization in finite-width networks?

Key findings

• The proposed method achieves a 106× speedup over exact convolutional NTK (CNTK) computation for a 5-layer Myrtle network on CIFAR-10, with minimal accuracy loss.
• The method enables exact dual kernel computation for monomials, Gabor, and ELU activations via Hermite series and special function identities.
• A new activation, ABReLU, derived from the normalized Gaussian kernel via reverse engineering, achieved the highest test accuracy (94.8%) on CIFAR-10 among ReLU, GeLU, Erf, and PReLU in a 5-layer Myrtle-5 network.
• The sketching algorithm provides a subspace embedding for NNGP and NTK matrices with near-input-sparsity runtime and near-optimal target dimension for homogeneous dual activation functions with rapidly convergent Taylor expansions.
• The method is applicable to any dataset in R^d and does not require data to be normalized to the unit sphere, unlike prior works.