# Kernel methods inside of Flux

**URL:** <https://discourse.julialang.org/t/kernel-methods-inside-of-flux/122820>\
**Category:** Machine Learning\
**Tags:** question, flux\
**Created:** [November 19, 2024, 10:06pm UTC](https://discourse.julialang.org/t/kernel-methods-inside-of-flux/122820 "2024-11-19T22:06:16Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![gideonsimpson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gideonsimpson/32/1928_2.png) [@gideonsimpson](https://discourse.julialang.org/u/gideonsimpson)\
**Post date:** [November 19, 2024, 10:06pm UTC](https://discourse.julialang.org/t/kernel-methods-inside-of-flux/122820/1 "2024-11-19T22:06:16Z")

</div>

I’m experimenting with setting up a classical RBF kernel method wtihin `Flux`, i.e., I want to train the model:

\sum\_{i=1}^n a\_i k(x, X\_i),

where, for simplicity, take the kernel k to be a Gaussian. n is fixed, and the X\_i are fixed; this is really just a regression problem. But thinking of it in a network architecture, what we want to have happen is:

x\mapsto (k(x, X\_i))\mapsto (\sum\_{i=1}^n a\_i k(x, X\_i))

which nis to say we go from `d=>n=>1`, where `d` is the dimension of the input argument.

What I am struggling with is getting the kernel layer to properly vectorize its output (i.e., if i input a `d x 10` array, I should get a `n x 10` output array), and if i push it all the way through i should end up with a `1 x 10` output array. This is what I’ve been trying:

```julia
Random.seed!(100)
n_pts = 100;
X_pts = [randn(Float32, 2,1) for _ in 1:n_pts];
X_train = randn(Float32, 2, 1000)
h = 1.0f0;
kernel(x, y) = exp(-sum(abs2, x - y) / (2 * h));

function k_layer(x)
    return [kernel(x, X_) for X_ in X_pts]
end

f_kernel = Chain(k_layer, Dense(n_pts => 1, bias=false));

```

and

```julia
f_kernel(rand(Float32,2))

```

works fine, but not

```julia
f_kernel(rand(Float32,2,10))

```

returns

```julia
DimensionMismatch: a has size (2, 10), b has size (2, 1), mismatch at dim 2

```

EDIT: The following works, but I have no sense of how clever/efficient an implementation it is:

```julia
f_kernel = Chain(x -> MapCols{2}(kernel_layers, x), final_layer)

```

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 20, 2024, 4:01am UTC](https://discourse.julialang.org/t/kernel-methods-inside-of-flux/122820/2 "2024-11-20T04:01:51Z")

</div>

Often the efficient way is not to make slices, and write everything to deal with solid arrays, keeping batch etc. dimensions the whole way through:

```julia
n_pts = 3;
X_pts = [randn(Float32, 2) for p in 1:n_pts];
X_fixed = hcat(X_pts...) # 2×3 Matrix, call it X[v,p] vector and point

X_train = randn(Float32, 2, 7) # x[v,b]: vector and batch indices
X_1 = X_train[:,1]

h = 1.0f0; # non-const global, not ideal
kernel(x, y) = exp(-sum(abs2, x - y) / (2 * h)); # z = exp( sum_v (x[v] - y[v])^2 / (-2h) )
kernel(X_pts[1], X_1) # Vector, Vector -> Float32

function k_layer(x) # z[p] = exp( sum_v (x[v] - X[v,p])^2 / (-2h) )
    return [kernel(x, X_) for X_ in X_pts]
end

k_layer(X_1) # 3-element Vector, z[p]
# k_layer(X_train) # error, but you want this:
# z[p,b] = exp( sum_v (x[v,b] - X[v,p])^2 / (-2h) )
# = exp(-1/2h sum_v (x[v,b])^2 - 2 x[v,b] * X[v,p] + (X[v,p])^2 )

using TensorCast
function k_layer2(x, X=X_fixed)
    @reduce tmp[p,b] := sum(v) (x[v,b] - X[v,p])^2 / (-2h)
    exp.(tmp)
end

k_layer2(X_train) # 3×7 Matrix, first col matches k_layer(X_1)

function k_layer3(x, X=X_fixed)
    term1 = diag(x'*x)' # t[b] = sum_v (x[v,b])^2
    term2 = X'*x # t[p,b] = sum_v x[v,b] * X[v,p]
    term3 = diag(X'*X) # t[p] = sum_v (X[v,p])^2
    @. exp((term1 -2*term2 + term3)/(-2h))
end

k_layer3(X_train) # 3×7 Matrix, first col matches k_layer(X_1)

using Flux
f_kernel = Chain(k_layer3, Dense(n_pts => 1, bias=false));

f_kernel(X_train) # 1×7 Matrix

```