Learnable epsilon in GIN

GNNLux/src/layers/conv.jl

@@ -1287,7 +1287,7 @@ function Base.show(io::IO, l::GatedGraphConv)
 end
 
 @doc raw"""
-    GINConv(f, ϵ; aggr=+)
+    GINConv(f, [ϵ]; aggr=+, train_eps=false)
 
 Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf).
 
@@ -1300,7 +1300,9 @@ where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer o
 # Arguments
 
 - `f`: A (possibly learnable) function acting on node features. 
-- `ϵ`: Weighting factor.
+- `ϵ`: Initial value of the weighting factor ``\epsilon``. Default `0.0f0`.
+- `train_eps`: If `true`, ``\epsilon`` is trainable. Default `false`.
+- `aggr`: Neighborhood aggregation function. Default `+`.
 
 # Examples:
 
@@ -1322,7 +1324,7 @@ x = randn(rng, Float32, in_channel, g.num_nodes)
 nn = Dense(in_channel, out_channel)
 
 # create layer
-l = GINConv(nn, 0.01f0, aggr = mean)
+l = GINConv(nn; ϵ = 0.01f0, train_eps = true, aggr = mean)
 
 # setup layer
 ps, st = LuxCore.setup(rng, l)
@@ -1333,23 +1335,41 @@ y, st = l(g, x, ps, st)       # size:  out_channel × num_nodes
 """
 @concrete struct GINConv <: GNNContainerLayer{(:nn,)}
     nn <: AbstractLuxLayer
-    ϵ <: Real
+    ϵ
     aggr
+    train_eps::Bool
 end
 
-GINConv(nn, ϵ; aggr = +) = GINConv(nn, ϵ, aggr)
+GINConv(nn, ϵ::Real; aggr = +, train_eps::Bool = false) =
+    GINConv(nn, train_eps ? [ϵ] : ϵ, aggr, train_eps)
+
+GINConv(nn, ϵ::Real, aggr) = GINConv(nn, ϵ; aggr)
+
+GINConv(nn; ϵ::Real = 0.0f0, aggr = +, train_eps::Bool = false) =
+    GINConv(nn, ϵ; aggr, train_eps)
+
+function LuxCore.initialparameters(rng::AbstractRNG, l::GINConv)
+    nn = LuxCore.initialparameters(rng, l.nn)
+    return l.train_eps ? (; nn, ϵ = l.ϵ) : (; nn)
+end
+
+LuxCore.parameterlength(l::GINConv) = parameterlength(l.nn) + l.train_eps
+LuxCore.statelength(l::GINConv) = statelength(l.nn)
 
 function (l::GINConv)(g, x, ps, st)
     nn = StatefulLuxLayer{true}(l.nn, ps.nn, st.nn)
-    m = (; nn, l.ϵ, l.aggr)
+    ϵ = l.train_eps ? ps.ϵ : l.ϵ
+    m = (; nn, ϵ, l.aggr)
     y = GNNlib.gin_conv(m, g, x)
     stnew = (; nn = _getstate(nn))
     return y, stnew
 end
 
 function Base.show(io::IO, l::GINConv)
     print(io, "GINConv($(l.nn)")
-    print(io, ", $(l.ϵ)")
+    print(io, ", ϵ=$(l.ϵ)")
+    l.train_eps && print(io, ", train_eps=true")
+    l.aggr == (+) || print(io, ", aggr=$(l.aggr)")
     print(io, ")")
 end
 
@@ -1879,4 +1899,4 @@ function (l::SAGEConv)(g, x, ps, st)
     m = (; ps.weight, bias = _getbias(ps), 
           l.σ, l.aggr)
     return GNNlib.sage_conv(m, g, x), st
-end
+end

GNNLux/test/layers/conv.jl

@@ -92,8 +92,17 @@
 
     @testset "GINConv" begin
         nn = Chain(Dense(in_dims => out_dims, tanh), Dense(out_dims => out_dims))
-        l = GINConv(nn, 0.5)
+        l = GINConv(nn; ϵ = 0.5)
         test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true)
+        ps = LuxCore.initialparameters(rng, l)
+        @test !hasproperty(ps, :ϵ)
+
+        l = GINConv(nn; train_eps = true)
+        test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true)
+        ps = LuxCore.initialparameters(rng, l)
+        @test ps.ϵ == [0.0f0]
+
+        @test GINConv(nn, 0.5).ϵ == 0.5
     end
 
     @testset "MEGNetConv" begin

GNNlib/src/layers/conv.jl

@@ -252,7 +252,8 @@ function gin_conv(l, g::AbstractGNNGraph, x)
     xj, xi = expand_srcdst(g, x) 
 
     m = propagate(copy_xj, g, l.aggr, xj = xj)
-    return l.nn((1 .+ ofeltype(xi, l.ϵ)) .* xi .+ m)
+    ϵ = l.ϵ isa Real ? ofeltype(xi, l.ϵ) : l.ϵ
+    return l.nn((1 .+ ϵ) .* xi .+ m)
 end
 
 ####################### NNConv ######################################

GraphNeuralNetworks/src/layers/conv.jl

@@ -590,7 +590,7 @@ function Base.show(io::IO, l::EdgeConv)
 end
 
 @doc raw"""
-    GINConv(f, ϵ; aggr=+)
+    GINConv(f, [ϵ]; aggr=+, train_eps=false)
 
 Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf).
 
@@ -603,7 +603,9 @@ where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer o
 # Arguments
 
 - `f`: A (possibly learnable) function acting on node features. 
-- `ϵ`: Weighting factor.
+- `ϵ`: Initial value of the weighting factor ``\epsilon``. Default `0.0f0`.
+- `train_eps`: If `true`, ``\epsilon`` is trainable. Default `false`.
+- `aggr`: Neighborhood aggregation function. Default `+`.
 
 # Examples:
 
@@ -619,28 +621,37 @@ g = GNNGraph(s, t)
 nn = Dense(in_channel, out_channel)
 
 # create layer
-l = GINConv(nn, 0.01f0, aggr = mean)
+l = GINConv(nn; ϵ = 0.01f0, train_eps = true, aggr = mean)
 
 # forward pass
 y = l(g, x)  
 ```
 """
-struct GINConv{R <: Real, NN, A} <: GNNLayer
+struct GINConv{E, NN, A} <: GNNLayer
     nn::NN
-    ϵ::R
+    ϵ::E
     aggr::A
+    train_eps::Bool
 end
 
 Flux.@layer :expand GINConv
-Flux.trainable(l::GINConv) = (nn = l.nn,)
+Flux.trainable(l::GINConv) = l.train_eps ? (; l.nn, l.ϵ) : (; l.nn)
 
-GINConv(nn, ϵ; aggr = +) = GINConv(nn, ϵ, aggr)
+GINConv(nn, ϵ::Real; aggr = +, train_eps::Bool = false) =
+    GINConv(nn, train_eps ? [ϵ] : ϵ, aggr, train_eps)
+
+GINConv(nn, ϵ::Real, aggr) = GINConv(nn, ϵ; aggr)
+
+GINConv(nn; ϵ::Real = 0.0f0, aggr = +, train_eps::Bool = false) =
+    GINConv(nn, ϵ; aggr, train_eps)
 
 (l::GINConv)(g, x) = GNNlib.gin_conv(l, g, x)
 
 function Base.show(io::IO, l::GINConv)
     print(io, "GINConv($(l.nn)")
-    print(io, ", $(l.ϵ)")
+    print(io, ", ϵ=$(l.ϵ)")
+    l.train_eps && print(io, ", train_eps=true")
+    l.aggr == (+) || print(io, ", aggr=$(l.aggr)")
     print(io, ")")
 end
 

GraphNeuralNetworks/test/layers/conv.jl

@@ -278,19 +278,31 @@ end
     using .TestModule
     nn = Dense(D_IN, D_OUT)
 
-    l = GINConv(nn, 0.01, aggr = mean)
+    l = GINConv(nn; ϵ = 0.01, aggr = mean)
     for g in TEST_GRAPHS
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
         test_gradients(l, g, g.x, rtol = RTOL_HIGH)
     end
+    @test keys(Flux.trainable(l)) == (:nn,)
 
-    @test !in(:eps, Flux.trainable(l))
+    l = GINConv(Dense(D_IN, D_OUT); train_eps = true, aggr = mean)
+    @test keys(Flux.trainable(l)) == (:nn, :ϵ)
+    @test l.ϵ == [0.0f0]
+    for g in TEST_GRAPHS
+        @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
+    end
+    g = first(TEST_GRAPHS)
+    grad = Flux.gradient(l -> mean(l(g, g.x)), l)[1]
+    @test grad.ϵ isa AbstractVector
+    @test length(grad.ϵ) == 1
+
+    @test GINConv(nn, 0.5).ϵ == 0.5
 end
 
 @testitem "GINConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
     using .TestModule
     nn = Dense(D_IN, D_OUT)
-    l = GINConv(nn, 0.01, aggr = mean)
+    l = GINConv(nn; train_eps = true, aggr = mean)
     for g in TEST_GRAPHS
         g.graph isa AbstractSparseMatrix && continue
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)