Transfer NeuralKernelNetwork over from Stheno (#283)

* NeuralKernelNetwork

* Bump patch

* Remove redundant line

* Improve docs slightly

* Fix formatting

* Fix formatting

* Remove Flux dep

* Remove Flux

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Add compat for LogExpFunctions

* Stop exporting NKN

* Fix formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.9.5"
+version = "0.9.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/KernelFunctions.jl

@@ -45,15 +45,16 @@ using Compat
 using ChainRulesCore: ChainRulesCore, Composite, Zero, One, DoesNotExist, NO_FIELDS
 using ChainRulesCore: @thunk, InplaceableThunk
 using CompositionsBase
-using Requires
-using Distances, LinearAlgebra
+using Distances
+using FillArrays
 using Functors
+using LinearAlgebra
+using Requires
 using SpecialFunctions: loggamma, besselk, polygamma
-using ZygoteRules: ZygoteRules
-using StatsFuns: logtwo, twoπ
+using StatsFuns: logtwo, twoπ, softplus
 using StatsBase
 using TensorCore
-using FillArrays
+using ZygoteRules: ZygoteRules
 
 abstract type Kernel end
 abstract type SimpleKernel <: Kernel end
@@ -96,6 +97,7 @@ include(joinpath("kernels", "kernelsum.jl"))
 include(joinpath("kernels", "kernelproduct.jl"))
 include(joinpath("kernels", "kerneltensorproduct.jl"))
 include(joinpath("kernels", "overloads.jl"))
+include(joinpath("kernels", "neuralkernelnetwork.jl"))
 include(joinpath("approximations", "nystrom.jl"))
 include("generic.jl")
 

src/distances/pairwise.jl

@@ -36,6 +36,14 @@ function colwise(d::PreMetric, x::AbstractVector)
     return zeros(Distances.result_type(d, x, x), length(x)) # Valid since d(x,x) == 0 by definition
 end
 
+function colwise(d::PreMetric, x::ColVecs)
+    return zeros(Distances.result_type(d, x.X, x.X), length(x)) # Valid since d(x,x) == 0 by definition
+end
+
+function colwise(d::PreMetric, x::RowVecs)
+    return zeros(Distances.result_type(d, x.X, x.X), length(x)) # Valid since d(x,x) == 0 by definition
+end
+
 ## The following is a hack for DotProduct and Delta to still work
 function colwise(d::Distances.UnionPreMetric, x::ColVecs)
     return Distances.colwise(d, x.X, x.X)

src/kernels/neuralkernelnetwork.jl

@@ -0,0 +1,125 @@
+# Linear layer, perform linear transformation to input array
+# x₁ = softplus.(W) * x₀
+struct LinearLayer{T,MT<:AbstractArray{T}}
+    W::MT
+end
+@functor LinearLayer
+
+LinearLayer(in_dim, out_dim) = LinearLayer(randn(out_dim, in_dim))
+
+(lin::LinearLayer)(x) = softplus.(lin.W) * x
+
+function Base.show(io::IO, layer::LinearLayer)
+    return print(io, "LinearLayer(", size(layer.W, 2), ", ", size(layer.W, 1), ")")
+end
+
+# Product function, given an 2d array whose size is M×N, product layer will
+# multiply every m neighboring rows of the array elementwisely to obtain
+# an new array of size (M÷m)×N
+function product(x, step=2)
+    m, n = size(x)
+    m % step == 0 || error("the first dimension of inputs must be multiple of step")
+    new_x = reshape(x, step, m ÷ step, n)
+    return .*([new_x[i, :, :] for i in 1:step]...)
+end
+
+# Primitive layer, mainly act as a container to hold basic kernels for the neural kernel network
+struct Primitive{T}
+    kernels::T
+    Primitive(ks...) = new{typeof(ks)}(ks)
+end
+@functor Primitive
+
+# flatten k kernel matrices of size Mk×Nk, and concatenate these 1d array into a k×(Mk*Nk) 2d array
+_cat_kernel_array(x) = vcat([reshape(x[i], 1, :) for i in 1:length(x)]...)
+
+# NOTE, though we implement `ew` & `pw` function for Primitive, it isn't a subtype of Kernel
+# type, I do this because it will facilitate writing NeuralKernelNetwork
+ew(p::Primitive, x) = _cat_kernel_array(map(k -> kernelmatrix_diag(k, x), p.kernels))
+pw(p::Primitive, x) = _cat_kernel_array(map(k -> kernelmatrix(k, x), p.kernels))
+
+function ew(p::Primitive, x, x′)
+    return _cat_kernel_array(map(k -> kernelmatrix_diag(k, x, x′), p.kernels))
+end
+pw(p::Primitive, x, x′) = _cat_kernel_array(map(k -> kernelmatrix(k, x, x′), p.kernels))
+
+function Base.show(io::IO, layer::Primitive)
+    print(io, "Primitive(")
+    join(io, layer.kernels, ", ")
+    return print(io, ")")
+end
+
+"""
+    NeuralKernelNetwork(primitives, nn)
+
+Constructs a Neural Kernel Network (NKN) [1].
+
+`primitives` are the based kernels, combined by `nn`.
+
+```julia
+k1 = 0.6 * (SEKernel() ∘ ScaleTransform(0.5))
+k2 = 0.4 * (Matern32Kernel() ∘ ScaleTransform(0.1))
+primitives = Primitive(k1, k2)
+nkn = NeuralKernelNetwork(primitives, Chain(LinearLayer(2, 2), product))
+```
+
+[1] - Sun, Shengyang, et al. "Differentiable compositional kernel learning for Gaussian
+    processes." International Conference on Machine Learning. PMLR, 2018.
+"""
+struct NeuralKernelNetwork{PT,NNT} <: Kernel
+    primitives::PT
+    nn::NNT
+end
+@functor NeuralKernelNetwork
+
+# use this function to reshape the 1d array back to kernel matrix
+_rebuild_kernel(x, n, m) = reshape(x, n, m)
+_rebuild_diag(x) = reshape(x, :)
+
+(κ::NeuralKernelNetwork)(x, y) = only(kernelmatrix(κ, [x], [y]))
+
+function kernelmatrix_diag(nkn::NeuralKernelNetwork, x::AbstractVector)
+    return _rebuild_diag(nkn.nn(ew(nkn.primitives, x)))
+end
+
+function kernelmatrix(nkn::NeuralKernelNetwork, x::AbstractVector)
+    return _rebuild_kernel(nkn.nn(pw(nkn.primitives, x)), length(x), length(x))
+end
+
+function kernelmatrix_diag(nkn::NeuralKernelNetwork, x::AbstractVector, x′::AbstractVector)
+    return _rebuild_diag(nkn.nn(ew(nkn.primitives, x, x′)))
+end
+
+function kernelmatrix(nkn::NeuralKernelNetwork, x::AbstractVector, x′::AbstractVector)
+    return _rebuild_kernel(nkn.nn(pw(nkn.primitives, x, x′)), length(x), length(x′))
+end
+
+function kernelmatrix_diag!(K::AbstractVector, nkn::NeuralKernelNetwork, x::AbstractVector)
+    K .= kernelmatrix_diag(nkn, x)
+    return K
+end
+
+function kernelmatrix!(K::AbstractMatrix, nkn::NeuralKernelNetwork, x::AbstractVector)
+    K .= kernelmatrix(nkn, x)
+    return K
+end
+
+function kernelmatrix_diag!(
+    K::AbstractVector, nkn::NeuralKernelNetwork, x::AbstractVector, x′::AbstractVector
+)
+    K .= kernelmatrix_diag(nkn, x, x′)
+    return K
+end
+
+function kernelmatrix!(
+    K::AbstractMatrix, nkn::NeuralKernelNetwork, x::AbstractVector, x′::AbstractVector
+)
+    K .= kernelmatrix(nkn, x, x′)
+    return K
+end
+
+function Base.show(io::IO, kernel::NeuralKernelNetwork)
+    print(io, "NeuralKernelNetwork(")
+    join(io, [kernel.primitives, kernel.nn], ", ")
+    return print(io, ")")
+end

test/Project.toml

@@ -7,6 +7,7 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Kronecker = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -21,6 +22,7 @@ Documenter = "0.25, 0.26"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 Flux = "0.10, 0.11, 0.12"
 ForwardDiff = "0.10"
+LogExpFunctions = "0.2"
 Kronecker = "0.4"
 PDMats = "0.9, 0.10, 0.11"
 ReverseDiff = "1.2"

test/kernels/neuralkernelnetwork.jl

@@ -0,0 +1,66 @@
+using KernelFunctions: NeuralKernelNetwork, LinearLayer, product, Primitive
+
+@testset "neural_kernel_network" begin
+    rng, N, N′, D = MersenneTwister(123456), 5, 6, 2
+    x0 = collect(range(-2.0, 2.0; length=N)) .+ 1e-3 .* randn(rng, N)
+    x1 = collect(range(-1.7, 2.3; length=N)) .+ 1e-3 .* randn(rng, N)
+    x2 = collect(range(-1.7, 3.3; length=N′)) .+ 1e-3 .* randn(rng, N′)
+
+    X0 = ColVecs(randn(rng, D, N))
+    X1 = ColVecs(randn(rng, D, N))
+    X2 = ColVecs(randn(rng, D, N′))
+
+    # Most of the NeuralKernelNetwork tests are currently broken.
+    @testset "general test" begin
+
+        # Specify primitives.
+        k1 = 0.6 * (SEKernel() ∘ ScaleTransform(0.5))
+        k2 = 0.4 * (Matern32Kernel() ∘ ScaleTransform(0.1))
+        primitives = Primitive(k1, k2)
+
+        # Build NKN Kernel.
+        nkn = NeuralKernelNetwork(primitives, Chain(LinearLayer(2, 2), product))
+
+        # Apply standard test suite.
+        TestUtils.test_interface(nkn, Float64)
+    end
+    @testset "kernel composition test" begin
+        rng = MersenneTwister(123456)
+
+        # Specify primitives.
+        k1 = rand(rng) * transform(SEKernel(), randn(rng))
+        k2 = rand(rng) * transform(Matern32Kernel(), randn(rng))
+        primitives = Primitive(k1, k2)
+
+        @testset "LinearLayer" begin
+            # Specify linear NKN and equivalent composite kernel.
+            weights = rand(rng, 1, 2)
+            nkn_add_kernel = NeuralKernelNetwork(primitives, LinearLayer(weights))
+            sum_k =
+                LogExpFunctions.softplus(weights[1]) * k1 +
+                LogExpFunctions.softplus(weights[2]) * k2
+
+            # Vector input.
+            @test kernelmatrix_diag(nkn_add_kernel, x0) ≈ kernelmatrix_diag(sum_k, x0)
+            @test kernelmatrix_diag(nkn_add_kernel, x0, x1) ≈
+                  kernelmatrix_diag(sum_k, x0, x1)
+
+            # ColVecs input.
+            @test kernelmatrix_diag(nkn_add_kernel, X0) ≈ kernelmatrix_diag(sum_k, X0)
+            @test kernelmatrix_diag(nkn_add_kernel, X0, X1) ≈
+                  kernelmatrix_diag(sum_k, X0, X1)
+        end
+        @testset "product" begin
+            nkn_prod_kernel = NeuralKernelNetwork(primitives, product)
+            prod_k = k1 * k2
+
+            # Vector input.
+            @test kernelmatrix(nkn_prod_kernel, x0) ≈ kernelmatrix(prod_k, x0)
+            @test kernelmatrix(nkn_prod_kernel, x0, x1) ≈ kernelmatrix(prod_k, x0, x1)
+
+            # ColVecs input.
+            @test kernelmatrix(nkn_prod_kernel, X0) ≈ kernelmatrix(prod_k, X0)
+            @test kernelmatrix(nkn_prod_kernel, X0, X1) ≈ kernelmatrix(prod_k, X0, X1)
+        end
+    end
+end

test/runtests.jl

@@ -4,6 +4,7 @@ using Distances
 using Documenter
 using Kronecker: Kronecker
 using LinearAlgebra
+using LogExpFunctions
 using PDMats
 using Random
 using SpecialFunctions
@@ -123,6 +124,7 @@ include("test_utils.jl")
         include(joinpath("kernels", "scaledkernel.jl"))
         include(joinpath("kernels", "transformedkernel.jl"))
         include(joinpath("kernels", "normalizedkernel.jl"))
+        include(joinpath("kernels", "neuralkernelnetwork.jl"))
     end
     @info "Ran tests on Kernel"