Neural net kernels

September 16, 2019 — May 24, 2021

Hilbert space

kernel tricks

machine learning

metrics

probabilistic algorithms

signal processing

spheres

statistics

stochastic processes

Suspiciously similar content

Figure 1: How I imagine the hyperspherical regularity of an NN kernel.

Random infinite-width NN induce covariances which are nearly dot product kernels in the input parameters. Say we wish to compare the outputs given two input examples. They depend on several dot products: $x^{⊤} x$ , $x^{⊤} y$ , and $y^{⊤} y$ . Often it is convenient to discuss the angle $θ$ between the inputs: $θ = \cos^{- 1} (\frac{x^{⊤} y}{∥ x ∥ ∥ y ∥})$

The classic result is that in a single-layer wide neural net, $\begin{aligned} K (x, y) & = E [ψ (Z_{x}) ψ (Z_{y})], where \\ (\begin{array}{c} Z_{x} \\ Z_{y} \end{array}) & \sim N (0, \underset{:= Σ}{\underset{⏟}{(\begin{array}{c} x^{⊤} x & x^{⊤} y \\ y^{⊤} x & y^{⊤} y \end{array})}}) . \end{aligned}$ It is sometimes useful to note that $(\begin{matrix} Z_{x} \\ Z_{y} \end{matrix}) \overset{d}{=} Chol (Σ) Z_{1},$ where $Z_{1} \sim N (0, (\begin{matrix} 1 & 0 \\ 0 & 1 \end{matrix}))$ and $Chol (Σ) = (\begin{matrix} ∥ x ∥ & ∥ y ∥ \cos θ \\ 0 & ∥ y ∥ \sqrt{1 - \cos^{2} θ} \end{matrix}) .$

These $Z_{x}$ terms arise from the (appropriately scaled limit of) the random weight matrix $\begin{aligned} Z_{x} & = W^{⊤} x \\ Z_{y} & = W^{⊤} y . \end{aligned}$ Now, define $\begin{aligned} Z_{x i} : & = W_{i} x_{i}, \\ Z_{y j} : & = W_{j} y_{j}, \\ Z_{x i}^{'} : & = W_{i}, \\ Z_{y j}^{'} : & = W_{j} . \end{aligned}$ We have that $\begin{aligned} κ & = E [ψ (Z_{x}) ψ (Z_{y})] \\ \frac{\partial κ}{\partial x_{i}} x_{i} & = E [ψ^{'} (Z_{x}) ψ (Z_{y}) Z_{x i}] \\ \frac{\partial^{2} κ}{\partial x_{i} \partial y_{j}} x_{i} y_{j} & = E [ψ^{'} (Z_{x}) ψ^{'} (Z_{y}) Z_{x i} Z_{y j}] \\ \frac{\partial^{2} κ}{\partial x_{i} \partial x_{j}} x_{i} x_{j} & = E [ψ^{″} (Z_{x}) ψ (Z_{y}) Z_{x i} Z_{x j}] \end{aligned}$ and thus $\begin{aligned} \frac{\partial κ}{\partial x_{i}} & = E [ψ^{'} (Z_{x}) ψ (Z_{y}) Z_{x i}^{'}] \\ \frac{\partial^{2} κ}{\partial x_{i} \partial y_{j}} & = E [ψ^{'} (Z_{x}) ψ^{'} (Z_{y}) Z_{x i}^{'} Z_{y j}^{'}] \\ \frac{\partial^{2} κ}{\partial x_{i} \partial x_{j}} & = E [ψ^{″} (Z_{x}) ψ (Z_{y}) Z_{x i}^{'} Z_{x j}^{'}] . \end{aligned}$

1 Erf kernel

Williams (1996) recover a kernel that corresponds to the Erf sigmoidal activation in the infinite width limit. Let $\tilde{x} = (1, x_{1}, \dots, x_{d})$ be an augmented copy of the inputs with a 1 prepended so that it includes the bias, and let $Σ$ be the covariance matrix of the weights (which are usually isotropic, $Σ = I$ ). Then $K_{\erf} (x, y)$ can be written as $K_{\erf} (x, y) = \frac{1}{(2 π)^{\frac{d + 1}{2}} | Σ |^{1 / 2}} \int Φ (w^{⊤} \tilde{x}) Φ (w^{⊤} \tilde{y}) \exp (- \frac{1}{2} w^{⊤} Σ^{- 1} w) d w .$ This integral can be evaluated analytically to give

$K_{\erf} (x, y) = \frac{2}{π} \sin^{- 1} \frac{2 {\tilde{x}}^{⊤} Σ \tilde{y}}{\sqrt{(1 + 2 {\tilde{x}}^{⊤} Σ \tilde{x}) (1 + 2 {\tilde{y}}^{⊤} Σ \tilde{y})}} .$

If there is no bias term, you can lop those tildes off and a factor of $\sqrt{2 π}$ and the result should still hold. If the weights are isotropic, the $Σ$ s vanish also.

2 Arc-cosine kernel

An interesting dot-product kernel is the arc-cosine kernel (Cho and Saul 2009):

$K_{n} (x, y) = \frac{2}{(2 π)^{\frac{d}{2}}} \int Θ (w^{⊤} x) Θ (w^{⊤} y) (w^{⊤} x)^{n} (w^{⊤} y)^{n} \exp (- \frac{1}{2} w^{⊤} w) d w$

Specifically, $K_{n} (x, y) = \frac{1}{π} ∥ x ∥^{n} ∥ y ∥^{n} J_{n} (θ)$ where $J_{n} (θ)$ is given by: $J_{n} (θ) = (- 1)^{n} (\sin θ)^{2 n + 1} {(\frac{1}{\sin θ} \frac{\partial}{\partial θ})}^{n} (\frac{π - θ}{\sin θ})$ The first few $J_{n}$ are $\begin{array}{l} J_{0} (θ) = π - θ \\ J_{1} (θ) = \sin θ + (π - θ) \cos θ . \end{array}$ $J_{1}$ recovers the ReLU activation in the infinite width limit. i.e. The arc-cosine kernel of order $1$ corresponding to the case where $ψ$ is the ReLU is $\begin{aligned} k (x, y) & = \frac{σ_{w}^{2} ‖ x ‖ ‖ y ‖}{2 π} (\sin | θ | + (π - | θ |) \cos θ) \end{aligned}$

Observation: This appears related to Grothendieck’s identity, that any fixed vectors $u, v \in S^{n - 1},$ we have $E sign X_{u} sign X_{v} = \frac{2}{π} \arcsin u^{⊤} v .$ I don’t have any use for that, it is just a cool identity I wanted to note down. In an aside Djalil Chafaï observes that the Rademacher RV is the distribution over the 1-dimensional sphere, $\in S^{0} .$ Is that what makes this go?

3 Absolutely homogenous

Activation functions which are absolutely homogeneous of degree $r$ satisfying $ψ (| a | z) = | a |^{r} ψ (z)$ have additional structure. This class includes the ReLU and leaky ReLU activations (which are also included as the first-order arc-cosine kernel above). It follows from the definition that functions $f$ drawn from an NN with such an activation a.s. satisfy $f (| a | x) = | a |^{r} f (x)$ .

For absolutely homogeneous activation we can sum the derivatives over the coordinate indices $\begin{aligned} \sum_{i, j = 1}^{d} \frac{\partial^{2} κ}{\partial x_{i} \partial x_{j}} x_{i} x_{j} & = E [ψ^{″} (Z_{x}) ψ (Z_{y}) (Z_{x})^{2}] = 0 \\ \sum_{i, j = 1}^{d} \frac{\partial^{2} κ}{\partial y_{i} \partial y_{j}} y_{i} y_{j} & = E [ψ^{″} (Z_{y}) ψ (Z_{x}) (Z_{y})^{2}] = 0 \\ \sum_{i, j = 1}^{d} \frac{\partial^{2} κ}{\partial x_{i} \partial y_{j}} x_{i} y_{j} & = κ . \end{aligned}$ i.e. $\begin{aligned} x \frac{\partial^{2} κ}{\partial x_{p} \partial x_{q}^{⊤}} y^{⊤} & = κ \\ x \frac{\partial^{2} κ}{\partial x_{p} \partial x_{p}^{⊤}} x^{⊤} & = 0 \\ y \frac{\partial^{2} κ}{\partial x_{q} \partial x_{q}^{⊤}} y^{⊤} & = 0. \end{aligned}$

4 References

Adlam, Lee, Xiao, et al. 2020. “Exploring the Uncertainty Properties of Neural Networks’ Implicit Priors in the Infinite-Width Limit.” arXiv:2010.07355 [Cs, Stat].

Arora, Du, Hu, et al. 2019. “On Exact Computation with an Infinitely Wide Neural Net.” In Advances in Neural Information Processing Systems.

Belkin, Ma, and Mandal. 2018. “To Understand Deep Learning We Need to Understand Kernel Learning.” In International Conference on Machine Learning.

Chen, and Xu. 2020. “Deep Neural Tangent Kernel and Laplace Kernel Have the Same RKHS.” arXiv:2009.10683 [Cs, Math, Stat].

Cho, and Saul. 2009. “Kernel Methods for Deep Learning.” In Proceedings of the 22nd International Conference on Neural Information Processing Systems. NIPS’09.

Domingos. 2020. “Every Model Learned by Gradient Descent Is Approximately a Kernel Machine.” arXiv:2012.00152 [Cs, Stat].

Fan, and Wang. 2020. “Spectra of the Conjugate Kernel and Neural Tangent Kernel for Linear-Width Neural Networks.” In Advances in Neural Information Processing Systems.

Fort, Dziugaite, Paul, et al. 2020. “Deep Learning Versus Kernel Learning: An Empirical Study of Loss Landscape Geometry and the Time Evolution of the Neural Tangent Kernel.” In Advances in Neural Information Processing Systems.

Geifman, Yadav, Kasten, et al. 2020. “On the Similarity Between the Laplace and Neural Tangent Kernels.” In arXiv:2007.01580 [Cs, Stat].

He, Lakshminarayanan, and Teh. 2020. “Bayesian Deep Ensembles via the Neural Tangent Kernel.” In Advances in Neural Information Processing Systems.

Jacot, Gabriel, and Hongler. 2018. “Neural Tangent Kernel: Convergence and Generalization in Neural Networks.” In Advances in Neural Information Processing Systems. NIPS’18.

Liu, Zhu, and Belkin. 2020. “On the Linearity of Large Non-Linear Models: When and Why the Tangent Kernel Is Constant.” In Advances in Neural Information Processing Systems.

Marteau-Ferey, Bach, and Rudi. 2020. “Non-Parametric Models for Non-Negative Functions.” In Proceedings of the 34th International Conference on Neural Information Processing Systems. NIPS ’20.

Neal. 1996. “Priors for Infinite Networks.” In Bayesian Learning for Neural Networks. Lecture Notes in Statistics.

Pearce, Tsuchida, Zaki, et al. 2019. “Expressive Priors in Bayesian Neural Networks: Kernel Combinations and Periodic Functions.” In Uncertainty in Artificial Intelligence.

Simon, Anand, and DeWeese. 2022. “Reverse Engineering the Neural Tangent Kernel.”

Tsuchida, Ong, and Sejdinovic. 2023. “Squared Neural Families: A New Class of Tractable Density Models.”

Tsuchida, Roosta, and Gallagher. 2018. “Invariance of Weight Distributions in Rectified MLPs.” In International Conference on Machine Learning.

Williams. 1996. “Computing with Infinite Networks.” In Proceedings of the 9th International Conference on Neural Information Processing Systems. NIPS’96.