rescale_weights!

Author: cossio

Project.toml

@@ -1,7 +1,7 @@
 name = "RestrictedBoltzmannMachines"
 uuid = "12e6b396-7db5-4506-8cb6-664a4fe1e50e"
 authors = ["Jorge Fernandez-de-Cossio-Diaz <j.cossio.diaz@gmail.com>"]
-version = "1.0.0"
+version = "2.0.0-DEV"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/from_grad.jl

@@ -2,45 +2,3 @@
 
 grad2ave(::Union{Binary,Spin,Potts,Gaussian,ReLU,pReLU,xReLU}, ∂::AbstractArray) = ∂[1, ..]
 grad2ave(::dReLU, ∂::AbstractArray) = ∂[1, ..] + ∂[2, ..]
-
-grad2var(::Union{Binary,Potts}, ∂::AbstractArray) = ∂[1, ..] .* (1 .- ∂[1, ..])
-grad2var(::Spin, ∂::AbstractArray) = (1 .- ∂[1, ..]) .* (1 .+ ∂[1, ..])
-
-function grad2var(l::Union{Gaussian,ReLU}, ∂::AbstractArray)
-    ∂θ = @view ∂[1, ..]
-    ∂γ = @view ∂[2, ..]
-    return -2∂γ .* sign.(l.γ) - ∂θ.^2
-end
-
-function grad2var(l::dReLU, ∂::AbstractArray)
-    ∂θp = ∂[1, ..]
-    ∂θn = ∂[2, ..]
-    ∂γp = ∂[3, ..]
-    ∂γn = ∂[4, ..]
-    return -2 * (∂γp .* sign.(l.γp) + ∂γn .* sign.(l.γn)) - (∂θp + ∂θn).^2
-end
-
-function grad2var(l::pReLU, ∂::AbstractArray)
-    ∂θ = -∂[1, ..]
-    ∂γ = -∂[2, ..]
-    ∂Δ = -∂[3, ..]
-    ∂η = -∂[4, ..]
-
-    abs_γ = abs.(l.γ)
-    ∂absγ = ∂γ .* sign.(l.γ)
-
-    return @. 2l.η/abs_γ * ((2l.Δ * ∂Δ + l.η * ∂η) * l.η - ∂η - l.Δ * ∂θ) + 2∂absγ * (1 + l.η^2) - ∂θ^2
-end
-
-function grad2var(l::xReLU, ∂::AbstractArray)
-    ∂θ = -∂[1, ..]
-    ∂γ = -∂[2, ..]
-    ∂Δ = -∂[3, ..]
-    ∂ξ = -∂[4, ..]
-
-    abs_γ = abs.(l.γ)
-    ∂absγ = ∂γ .* sign.(l.γ)
-
-    ν = @. 2∂absγ - ∂θ^2
-    return @. (ν * abs_γ - 2 * (∂ξ + ∂θ * l.Δ) * l.ξ + ((ν + 2∂absγ) * abs_γ + 4 * ∂Δ * l.Δ) * l.ξ^2 - 4∂ξ * l.ξ^3 + 2abs(l.ξ) * (ν * abs_γ - 3∂ξ * l.ξ - ∂θ * l.Δ * l.ξ)) / (abs_γ * (1 + abs(l.ξ))^2)
-end

src/gauge/rescale_hidden.jl

@@ -1,11 +1,9 @@
 """
     rescale_hidden!(rbm, λ::AbstractArray)
 
-For continuous hidden units with a scale parameter,
-scales parameters such that hidden unit activations
-are divided by `λ`. For other hidden units does
-nothing. The resulting RBM is equivalent to the
-original one.
+For continuous hidden units with a scale parameter, scales parameters such that hidden
+unit activations are divided by `λ`. For other hidden units does nothing. The resulting RBM
+is equivalent to the original one.
 """
 function rescale_hidden!(rbm::RBM, λ::AbstractArray)
     @assert size(rbm.hidden) == size(λ)
@@ -15,12 +13,28 @@ function rescale_hidden!(rbm::RBM, λ::AbstractArray)
     return rbm
 end
 
+"""
+    rescale_weights!(rbm, λ::AbstractArray)
+
+For continuous hidden units with a scale parameter, scales parameters such that the weights
+attached to each hidden unit have norm 1.
+"""
+function rescale_weights!(rbm::RBM)
+    ω = weight_norms(rbm)
+    λ = inv.(ω)
+    return rescale_hidden!(rbm, λ)
+end
+
+function weight_norms(rbm::RBM)
+    w2 = sum(abs2, rbm.w; dims=1:ndims(rbm.visible))
+    return reshape(sqrt.(w2), size(rbm.hidden))
+end
+
 """
     rescale_activations!(layer, λ::AbstractArray)
 
-For continuous layers with scale parameters, re-parameterizes
-such that unit activations are divided by `λ`, and returns `true`.
-For other layers just returns `false`.
+For continuous layers with scale parameters, re-parameterizes such that unit activations
+are divided by `λ`, and returns `true`. For other layers, does nothing and returns `false`.
 """
 rescale_activations!(layer::Union{Binary,Spin,Potts}, λ::AbstractArray) = false
 
@@ -29,15 +43,15 @@ must have positive activations. So we dissallow it below. =#
 
 function rescale_activations!(layer::Union{Gaussian,ReLU}, λ::AbstractArray)
     @assert size(layer) == size(λ)
-    @assert all(λ .> 0)
+    @assert all(>(0), λ)
     layer.θ .*= λ
     layer.γ .*= λ.^2
     return true
 end
 
 function rescale_activations!(layer::dReLU, λ::AbstractArray)
     @assert size(layer) == size(λ)
-    @assert all(λ .> 0)
+    @assert all(>(0), λ)
     layer.θp .*= λ
     layer.θn .*= λ
     layer.γp .*= λ.^2
@@ -47,7 +61,7 @@ end
 
 function rescale_activations!(layer::Union{pReLU,xReLU}, λ::AbstractArray)
     @assert size(layer) == size(λ)
-    @assert all(λ .> 0) # makes life simpler
+    @assert all(>(0), λ) # it's just simpler
     layer.θ .*= λ
     layer.Δ .*= λ
     layer.γ .*= λ.^2

src/train/pcd.jl

@@ -21,11 +21,7 @@ function pcd!(
 
     # gauge
     zerosum::Bool = true, # zerosum gauge for Potts layers
-    rescale::Bool = true, # normalize continuous hidden units to var(h) = 1
-
-    # momentum for hidden activity statistics tracking
-    ρh::Real = 99//100,
-    ϵh::Real = 1//100, # prevent vanishing var(h) estimate
+    rescale::Bool = true, # normalize weights to unit norm (for continuous hidden units only)
 
     callback = Returns(nothing), # called for every batch
 
@@ -36,15 +32,10 @@ function pcd!(
 )
     @assert size(data) == (size(rbm.visible)..., size(data)[end])
     @assert isnothing(wts) || size(data)[end] == length(wts)
-    @assert ϵh ≥ 0
-
-    # used to scale hidden unit activities
-    var_h = total_var_from_inputs(rbm.hidden, inputs_h_from_v(rbm, data); wts)
-    @assert all(var_h .+ ϵh .> 0)
 
     # gauge constraints
     zerosum && zerosum!(rbm)
-    rescale && rescale_hidden!(rbm, sqrt.(var_h .+ ϵh))
+    rescale && rescale_weights!(rbm)
 
     # store average weight of each data point
     wts_mean = isnothing(wts) ? 1 : mean(wts)
@@ -73,14 +64,8 @@ function pcd!(
         batch_weight = isnothing(wts) ? 1 : mean(wd) / wts_mean
         ∂ *= batch_weight
 
-        # Exponential moving average of variance of hidden unit activations.
-        ρh_eff = ρh ^ batch_weight # effective damp after 'batch_weight' updates
-        var_h_batch = grad2var(rbm.hidden, -∂d.hidden) # extract hidden unit statistics from gradient
-        var_h .= ρh_eff * var_h .+ (1 - ρh_eff) * var_h_batch
-        @assert all(var_h .+ ϵh .> 0)
-
         # reset gauge
-        rescale && rescale_hidden!(rbm, sqrt.(var_h .+ ϵh))
+        rescale && rescale_weights!(rbm)
         zerosum && zerosum!(rbm)
 
         callback(; rbm, optim, iter, vm, vd, wd)

test/gauge/rescale_hidden.jl

@@ -2,9 +2,10 @@ import Random
 using Test: @test, @testset, @inferred
 using Statistics: mean, var
 using Random: bitrand, rand!, randn!
+using LinearAlgebra: norm
 using RestrictedBoltzmannMachines: RBM, Binary, free_energy, Gaussian, ReLU, dReLU, pReLU, xReLU,
     sample_v_from_v, sample_h_from_h, mean_from_inputs, var_from_inputs,
-    rescale_hidden!, rescale_activations!
+    rescale_hidden!, rescale_activations!, rescale_weights!,  weight_norms
 
 Random.seed!(23)
 
@@ -55,3 +56,19 @@ end
     @test var(v; dims=2) ≈ var_v rtol=0.1
     @test var(h; dims=2) ≈ var_h ./ λ.^2 rtol=0.1
 end
+
+@testset "rescale_weights!" begin
+    rbm = RBM(Binary((2,)), ReLU((1,)), randn(2,1))
+    randn!(rbm.visible.θ)
+    randn!(rbm.hidden.θ)
+    rand!(rbm.hidden.γ)
+    rbm.hidden.γ .+= 0.5
+
+    v = sample_v_from_v(rbm, bitrand(size(rbm.visible)..., 1000); steps=100)
+    F = free_energy(rbm, v)
+
+    ω = @inferred weight_norms(rbm)
+    @test ω ≈ [norm(rbm.w)]
+    @inferred rescale_weights!(rbm)
+    @test free_energy(rbm, v) ≈ F .- sum(log, ω)
+end

test/layers.jl

@@ -8,7 +8,7 @@ using LogExpFunctions: logistic
 using EllipsisNotation: (..)
 using QuadGK: quadgk
 using RestrictedBoltzmannMachines: RBM, Binary, Spin, Potts, Gaussian, ReLU, dReLU, xReLU, pReLU,
-    flatten, batch_size, batchmean, batchvar, batchcov, grad2ave, grad2var, drelu_energy,
+    flatten, batch_size, batchmean, batchvar, batchcov, grad2ave, drelu_energy,
     mean_from_inputs, var_from_inputs, meanvar_from_inputs, batchdims, gauss_energy, relu_energy,
     std_from_inputs, mean_abs_from_inputs, sample_from_inputs, mode_from_inputs,
     energy, cgf, free_energy, cgfs, energies, ∂cgf, vstack, ∂energy, ∂free_energy, binary_rand,
@@ -156,7 +156,6 @@ end
     ∂ = ∂cgf(layer)
     @test ∂ ≈ only(gs).par ≈ vstack((mean_from_inputs(layer),))
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "Spin" begin
@@ -170,7 +169,6 @@ end
     ∂ = ∂cgf(layer)
     @test ∂ ≈ only(gs).par ≈ vstack((mean_from_inputs(layer),))
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "Potts" begin
@@ -189,7 +187,6 @@ end
     ∂ = ∂cgf(layer)
     @test ∂ ≈ only(gs).par ≈ vstack((mean_from_inputs(layer),))
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "Gaussian" begin
@@ -220,7 +217,6 @@ end
     @test ∂[1, ..] ≈ μ
     @test ∂[2, ..] ≈ -sign.(layer.γ) .* μ2/2
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "ReLU" begin
@@ -251,7 +247,6 @@ end
     @test ∂[1, ..] ≈ μ
     @test ∂[2, ..] ≈ -sign.(layer.γ) .* μ2/2
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "pReLU / xReLU / dReLU convert" begin
@@ -372,7 +367,6 @@ end
     ∂ = @inferred ∂cgf(layer)
     @test ∂ ≈ only(gs).par
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 
     # check law of total variance
     inputs = randn(size(layer)..., 1000)
@@ -384,7 +378,6 @@ end
     ν_ext = batchvar(layer, h_ave; mean = μ)
     ν = ν_int + ν_ext # law of total variance
     @test grad2ave(layer, ∂) ≈ μ
-    @test grad2var(layer, ∂) ≈ ν
     μ1, ν1 = total_meanvar_from_inputs(layer, inputs)
     @test μ1 ≈ μ ≈ total_mean_from_inputs(layer, inputs)
     @test ν1 ≈ ν ≈ total_var_from_inputs(layer, inputs)
@@ -399,7 +392,6 @@ end
     ∂ = ∂cgf(layer)
     @test ∂ ≈ only(gs).par
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "xReLU" begin
@@ -411,7 +403,6 @@ end
     ∂ = ∂cgf(layer)
     @test ∂ ≈ only(gs).par
     @test grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
-    @test grad2var(layer, ∂) ≈ var_from_inputs(layer)
 end
 
 @testset "grad2ave $Layer" for Layer in _layers