Conditional probability distributions powered by Flux.jl and DistributionsAD.jl.
The conditional PDFs that are defined in this package can be used in
conjunction with Flux models to provide trainable mappings. As an example,
assume you want to learn the mapping from a conditional to an MvNormal. The
mapping m takes a vector x and maps it to a mean μ and a variance σ,
which can be achieved by using a ConditionalDists.SplitLayer as the last
layer in the network.
julia> m = Chain(Dense(2,2,σ), SplitLayer(2, [3,4]))
julia> m(rand(2))
(Float32[0.07946974, 0.13797458, 0.03939067], Float32[0.7006321, 0.37641272, 0.3586885, 0.82230335])With the mapping m we can create a conditional distribution with trainable
mapping parameters:
julia> using ConditionalDists, Distributions
julia> using ConditionalDists: SplitLayer
julia> xlength = 3
julia> zlength = 2
julia> batchsize = 10
julia> m = SplitLayer(zlength, [xlength,xlength])
julia> p = ConditionalMvNormal(m)
julia> res = condition(p, rand(zlength)) # this also works for batches!
julia> μ = mean(res)
julia> σ2 = var(res)
julia> @assert res isa DistributionsAD.TuringDiagMvNormal
julia> x = rand(xlength, batchsize)
julia> z = rand(zlength, batchsize)
julia> logpdf(p,x,z)
julia> rand(p, randn(zlength, 10))The trainable parameters (of the SplitLayer) are accessible as usual through
Flux.params. For different variance configurations (i.e. fixed/unit variance,
etc) check the doc strings with julia>? ConditionalMvNormal/julia>? SplitLayer.
The next few lines show how to optimize p to match a
given Gaussian by using the kl_divergence defined in
IPMeasures.jl.
julia> using IPMeasures
julia> d = MvNormal(zeros(xlength), ones(xlength))
julia> loss(x) = sum(kl_divergence(p, d, x))
julia> opt = ADAM(0.1)
julia> data = [(randn(zlength),) for i in 1:2000]
julia> Flux.train!(loss, Flux.params(p), data, opt)The learnt mean and variance are close to a standard normal:
julia> @assert condition(p, randn(zlength)) isa DistributionsAD.TuringDiagMvNormal
julia> mean(p,randn(zlength))
3-element Array{Float32,1}:
0.003194877
1.7912015f-32
-1.6101733f-6
julia> var(p,randn(zlength))
3-element Array{Float32,1}:
1.000585
1.0021493
1.0000007