Flux.jl: Gaussian process

[kkmann]

Hi,

is there currently any way to implement GPs in Flux? inv() is available on the current master, but I cannot get the log determinant of my kernel matrix to work with Flux. I.e.,

det(K)

for a two dimensional array K complains about lu! not being implemented for TrackedArrays :(
Is there currently a way to do this? Are you planning to support these functions in the future? I was really hoping to get rid of tensorflow x)

[MikeInnes]

I believe @willtebbutt was doing some GP work with Flux, so he my have an idea of what we need to do. Adding a gradient for det doesn't seem like it should be too difficult.

[willtebbutt]

@kkmann we're currently missing a couple of key AD-related features. The key ones being cholesky and logdet for TrackedArrays (you don't really want to implement GPs in terms of the lu factorisation, inv and det for numeric stability and performance reasons). I'm working on adding them at the minute (I need them for my own research, so am quite motivated to get this sorted), and will have a PR open soon.

Once we have these, Flux should "just work" with Stheno.jl , and let us do lots of interesting things to combine GPs with Deep Learning in Julia :)

[kkmann]

Excellent news, I'll stay tuned - thanks for the great work!

[MoonCoral]

Workaround:

julia> a = Flux.param(rand(3,3))
Tracked 3×3 Array{Float64,2}:
 0.963785  0.568489  0.314755
 0.171466  0.892723  0.0310552
 0.113449  0.647206  0.682914

julia> b = Matrix{eltype(a)}(undef, 3, 3)
3×3 Array{Flux.Tracker.TrackedReal{Float64},2}:
 #undef  #undef  #undef
 #undef  #undef  #undef
 #undef  #undef  #undef

julia> b .= a
3×3 Array{Flux.Tracker.TrackedReal{Float64},2}:
 0.963785 (tracked)  0.568489 (tracked)  0.314755 (tracked)
 0.171466 (tracked)  0.892723 (tracked)  0.0310552 (tracked)
 0.113449 (tracked)  0.647206 (tracked)  0.682914 (tracked)

julia> det(b)
0.5066896404099932 (tracked)

[chriscoey]

@willtebbutt any progress on cholesky and logdet?

[willtebbutt]

@willtebbutt any progress on cholesky and logdet?

I've now implemented them in Zygote. Perhaps someone could open a PR to translate them to Flux? @MikeInnes perhaps?

[chriscoey]

Ah great, I'll try Zygote then instead, thanks

[MikeInnes]

Yes, Zygote is still beta so if you run into issues with that I'm happy to port those gradients back to flux.

[aterenin]

+1 to please port cholesky back to Flux.

I'd imagine this is much more efficient than converting TrackedMatrix{Float64} to Matrix{Flux.Tracker.TrackedReal{Float64}} and then calling cholesky on that.

[kkmann]

been a long time since I played around with this. looks like you can do that now:

using Flux
using Flux.Tracker
import LinearAlgebra.cholesky
import LinearAlgebra.logdet

k(x::T1, y::T1, l::T2) where {T1,T2<:Real} = exp(-((x - y)/l)^2)

function k(x::Vector{T1}, y::Vector{T1}, l::T2) where{T1,T2<:Real}
    m   = length(x); n = length(y)
    res = zeros(T2, m, n)
    for i in 1:m, j in 1:n
        res[i, j] = k(x[i], y[j], l)
    end
    return res
end

l = param(1.)

function loglikelihood(x, y)
    L    = cholesky(k(x, x, l))
    Linv = inv(L)
    res  = -2*logdet(L) - transpose(y)*transpose(Linv)*Linv*y - length(x)*log(pi)
end

x = [1.; 2; 3]
y = [0.; 1; 0]

grads = Tracker.gradient(() -> loglikelihood(x, y), params(l))
grads[l]

(I know I could vectorize k() directly, but supporting this kind of no-brainer element-wise loop code efficiently is kinda what makes Flux so much more enjoyable than TF ;) )

[kkmann]

strange though, manual gradient descent works fine but none of the provided optimizers

loss(x, y) = -loglikelihood(x, y)

for i in 1:10000
    grads = gradient(() -> loss(x, y), params(l))
    update!(l, -.1 * grads[l])
end
l
grads = gradient(() -> loss(x, y), params(l))
grads[l]

data = [(x, y)]
Flux.train!(loss, params(l), data, Descent())

train! throws an error:

MethodError: no method matching copyto!(::Float64, ::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Tuple{},typeof(*),Tuple{Float64,Float64}})
Closest candidates are:
copyto!(!Matched::AbstractArray, ::Base.Broadcast.Broadcasted{#s561,Axes,F,Args} where Args<:Tuple where F where Axes where #s561<:Base.Broadcast.AbstractArrayStyle{0}) at broadcast.jl:798
copyto!(!Matched::AbstractArray, ::Base.Broadcast.Broadcasted) at broadcast.jl:792
copyto!(!Matched::AbstractArray, ::Any) at abstractarray.jl:644
...

[nxtruong]

@kkmann If you define l to be a vector instead of a scalar (l = param([1.])) and adapt loglikelihood accordingly, you won't have the error. Apparently, Flux doesn't like scalar parameters.

[kkmann]

makes sense performance wise. it's sometimes a bit difficult to predict performance impact since you are not restricted to a DSL as in TF...