Add norm flows

src/Bijectors.jl

@@ -5,8 +5,9 @@ using Reexport, Requires
 using StatsFuns
 using LinearAlgebra
 using MappedArrays
+using Roots
 
-export  TransformDistribution, 
+export  TransformDistribution,
         RealDistribution,
         PositiveDistribution,
         UnitDistribution,
@@ -28,7 +29,9 @@ export  TransformDistribution,
         bijector,
         transformed,
         UnivariateTransformed,
-        MultivariateTransformed
+        MultivariateTransformed,
+        PlanarLayer,
+        RadialLayer
 
 const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_BIJECTORS", "0")))
 
@@ -177,8 +180,8 @@ function _clamp(x::T, dist::SimplexDistribution) where T
 end
 
 function link(
-    d::SimplexDistribution, 
-    x::AbstractVector{T}, 
+    d::SimplexDistribution,
+    x::AbstractVector{T},
     ::Type{Val{proj}} = Val{true}
 ) where {T<:Real, proj}
     y, K = similar(x), length(x)
@@ -206,8 +209,8 @@ end
 
 # Vectorised implementation of the above.
 function link(
-    d::SimplexDistribution, 
-    X::AbstractMatrix{T}, 
+    d::SimplexDistribution,
+    X::AbstractMatrix{T},
     ::Type{Val{proj}} = Val{true}
 ) where {T<:Real, proj}
     Y, K, N = similar(X), size(X, 1), size(X, 2)
@@ -234,8 +237,8 @@ function link(
 end
 
 function invlink(
-    d::SimplexDistribution, 
-    y::AbstractVector{T}, 
+    d::SimplexDistribution,
+    y::AbstractVector{T},
     ::Type{Val{proj}} = Val{true}
 ) where {T<:Real, proj}
     x, K = similar(y), length(y)
@@ -260,8 +263,8 @@ end
 
 # Vectorised implementation of the above.
 function invlink(
-    d::SimplexDistribution, 
-    Y::AbstractMatrix{T}, 
+    d::SimplexDistribution,
+    Y::AbstractMatrix{T},
     ::Type{Val{proj}} = Val{true}
 ) where {T<:Real, proj}
     X, K, N = similar(Y), size(Y, 1), size(Y, 2)
@@ -355,8 +358,8 @@ function invlink(d::PDMatDistribution, Y::AbstractMatrix{T}) where {T<:Real}
 end
 
 function logpdf_with_trans(
-    d::PDMatDistribution, 
-    X::AbstractMatrix{<:Real}, 
+    d::PDMatDistribution,
+    X::AbstractMatrix{<:Real},
     transform::Bool
 )
     T = eltype(X)
@@ -436,4 +439,6 @@ end
 
 include("interface.jl")
 
+include("norm_flows.jl")
+
 end # module

src/interface.jl

@@ -46,15 +46,15 @@ end
 Broadcast.broadcastable(b::Bijector) = Ref(b)
 
 "Computes the log(abs(det(J(x)))) where J is the jacobian of the transform."
-logabsdetjac(b::T1, y::T2) where {T<:Bijector,T1<:Inversed{T},T2} = 
+logabsdetjac(b::T1, y::T2) where {T<:Bijector,T1<:Inversed{T},T2} =
     error("`logabsdetjac(b::$T1, y::$T2)` is not implemented.")
 
 "Transforms the input using the bijector."
 transform(b::T1, y::T2) where {T<:Bijector,T1<:Inversed{T},T2} =
     error("`transform(b::$T1, y::$T2)` is not implemented.")
 
 "Computes both `transform` and `logabsdetjac` in one forward pass."
-forward(b::T1, y::T2) where {T<:Bijector,T1<:Inversed{T},T2} = 
+forward(b::T1, y::T2) where {T<:Bijector,T1<:Inversed{T},T2} =
     error("`forward(b::$T1, y::$T2)` is not implemented.")
 
 
@@ -112,7 +112,7 @@ end
 
 function compose(ts...)
     res = []
-    
+
     for b ∈ ts
         if b isa Composed
             # "lift" the transformations
@@ -151,6 +151,14 @@ function transform(cb::Composed{<: Bijector}, x)
     return res
 end
 
+function inv(cb::Composed{<: Bijector}, y)
+    res = y
+    for b ∈ reverse(cb.ts)
+        res = inv(b, res)
+    end
+    return res
+end
+
 (cb::Composed{<: Bijector})(x) = transform(cb, x)
 
 function forward(cb::Composed{<:Bijector}, x)
@@ -162,6 +170,12 @@ function forward(cb::Composed{<:Bijector}, x)
     return res
 end
 
+function rand(flow::Composed, dims::Integer, shape::Integer=1)
+    dims = [dims]
+    append!(dims, shape)
+    print(dims)
+    return transform(flow, randn(dims...))
+end
 ##############################
 # Example bijector: Identity #
 ##############################

src/norm_flows.jl

@@ -0,0 +1,130 @@
+using Distributions
+using LinearAlgebra
+using Random
+using StatsFuns: softplus
+using Roots # for inverse
+
+################################################################################
+#                            Planar and Radial Flows                           #
+#             Ref: Variational Inference with Normalizing Flows,               #
+#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
+################################################################################
+
+(b::Bijector)(x) = transform(b, x)
+
+mutable struct PlanarLayer{T1,T2} <: Bijector
+    w::T1
+    u::T1
+    b::T2
+end
+
+function get_u_hat(u, w)
+    # To preserve invertibility
+    return (
+        u + (planar_flow_m(transpose(w) * u) - transpose(w) * u)[1]
+        * w / (norm(w[:,1],2) ^ 2)
+    ) # from A.1
+end
+
+function PlanarLayer(dims::Int, container=Array)
+    w = container(randn(dims, 1))
+    u = container(randn(dims, 1))
+    b = container(randn(1))
+    return PlanarLayer(w, u, b)
+end
+
+planar_flow_m(x) = -1 .+ softplus.(x) # for planar flow from A.1
+dtanh(x) = 1 .- (tanh.(x)) .^ 2 # for planar flow
+ψ(z, w, b) = dtanh(transpose(w) * z .+ b) .* w # for planar flow from eq(11)
+
+function transform(flow::PlanarLayer, z)
+    u_hat = get_u_hat(flow.u, flow.w)
+    return z + u_hat * tanh.(transpose(flow.w) * z .+ flow.b) # from eq(10)
+end
+
+function forward(flow::T, z) where {T<:PlanarLayer}
+    u_hat = get_u_hat(flow.u, flow.w)
+    # Compute log_det_jacobian
+    psi = ψ(z, flow.w, flow.b)
+    log_det_jacobian = log.(abs.(1.0 .+ transpose(psi) * u_hat)) # from eq(12)
+    transformed = z + u_hat * tanh.(transpose(flow.w) * z .+ flow.b)
+    return (rv=transformed, logabsdetjac=log_det_jacobian) # from eq(10)
+end
+
+function inv(flow::PlanarLayer, y)
+    u_hat = get_u_hat(flow.u, flow.w)
+    # Implemented with reference from A.1
+    function f(y)
+        return loss(alpha) = (
+                (transpose(flow.w) * y)[1] - alpha
+                - (transpose(flow.w) * u_hat)[1]
+                * tanh(alpha+flow.b[1])
+            )
+    end
+    alphas_ = [find_zero(f(y[:,i:i]), 0.0, Order16()) for i in 1:size(y, 2)]
+    alphas = transpose(alphas_)
+    z_para = (flow.w ./ norm(flow.w,2)) * alphas
+    z_per = (
+            y - z_para - u_hat * tanh.(
+                                    transpose(flow.w) * z_para
+                                    .+ flow.b
+            )
+    )
+
+    return z_para + z_per
+end
+
+mutable struct RadialLayer{T1,T2} <: Bijector
+    α_::T1
+    β::T1
+    z_0::T2
+end
+
+function RadialLayer(dims::Int, container=Array)
+    α_ = container(randn(1))
+    β = container(randn(1))
+    z_0 = container(randn(dims, 1))
+    return RadialLayer(α_, β, z_0)
+end
+
+h(α, r) = 1 ./ (α .+ r) # for radial flow from eq(14)
+dh(α, r) = - h(α, r) .^ 2 # for radial flow, derivative of h()
+
+function transform(flow::RadialLayer, z)
+    α = softplus(flow.α_[1]) # from A.2
+    β_hat = -α + softplus(flow.β[1]) # from A.2
+    r = transpose(norm.([z[:,i] .- flow.z_0 for i in 1:size(z, 2)], 2))
+    return z + β_hat .* h(α, r) .* (z .- flow.z_0) # from eq(14)
+end
+
+function forward(flow::T, z) where {T<:RadialLayer}
+    α = softplus(flow.α_[1]) # from A.2
+    β_hat = -α + softplus(flow.β[1]) # from A.2
+    r = transpose(norm.([z[:,i] .- flow.z_0 for i in 1:size(z, 2)], 2))
+    transformed = z + β_hat .* h(α, r) .* (z .- flow.z_0) # from eq(14)
+    # Compute log_det_jacobian
+    d = size(flow.z_0, 1)
+    h_ = h(α, r)
+    log_det_jacobian = @. (
+        (d-1) * log(1.0 + β_hat * h_)
+        + log(1.0 +  β_hat * h_ + β_hat * (- h_ ^ 2) * r)
+    ) # from eq(14)
+    return (rv=transformed, logabsdetjac=log_det_jacobian)
+end
+
+function inv(flow::RadialLayer, y)
+    α = softplus(flow.α_[1]) # from A.2
+    β_hat = - α + softplus(flow.β[1]) # from A.2
+    function f(y)
+        # From eq(26)
+        return loss(r) = (
+                        norm(y - flow.z_0, 2)
+                        - r * (1 + β_hat / (α + r))
+                        )
+    end
+    rs_ = [find_zero(f(y[:,i:i]), 0.0, Order16()) for i in 1:size(y, 2)] # A.2
+    rs = transpose(rs_)
+    z_hat = (y .- flow.z_0) ./ (rs .* (1 .+ β_hat ./ (α .+ rs)) ) # from eq(25)
+    z = flow.z_0 .+ rs .* z_hat # from A.2
+    return z
+end

test/norm_flows.jl

@@ -0,0 +1,39 @@
+using Test
+using Bijectors, ForwardDiff, LinearAlgebra
+
+@testset "planar flows" begin
+    for i in 1:10
+        flow = PlanarLayer(10)
+        z = randn(10, 100)
+        forward_diff = log(abs(det(ForwardDiff.jacobian(t -> transform(flow, t), z))))
+        our_method = sum(forward(flow, z).logabsdetjacob)
+        @test our_method ≈ forward_diff
+
+        # Inverse not accurate enough to pass with `≈` operator.
+        @test_broken inv(flow, transform(flow, z)) ≈ z
+    end
+
+    w = ones(10, 1)
+    u = zeros(10, 1)
+    b = ones(1)
+    flow = PlanarLayer(w, u, b)
+    z = ones(10, 100)
+    @test inv(flow, transform(flow, z)) ≈ z
+end
+
+@testset "radial flows" begin
+    for i in 1:10
+        flow = RadialLayer(2)
+        z = randn(2, 100)
+        forward_diff = log(abs(det(ForwardDiff.jacobian(t -> transform(flow, t), z))))
+        our_method = sum(forward(flow, z).logabsdetjacob)
+        @test our_method ≈ forward_diff
+        @test inv(flow, transform(flow, z)) ≈ z
+    end
+    α_ = ones(1)
+    β = ones(1)
+    z_0 = zeros(10, 1)
+    z = ones(10, 100)
+    flow = RadialLayer(α_, β, z_0)
+    @test inv(flow, transform(flow, z)) ≈ z
+end