GNNLux MEGNetConv drops the state of stateful sublayers

Open Beginner friendly
  • Is something wrong?
#732 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
76/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
julia
Domain
machine-learning

Research direction

Read the MEGNetConv forward pass and compare how EdgeConv, GINConv, and NNConv return updated sublayer state. Add a regression test using BatchNorm in ϕe that checks whether its running statistics change after a training-mode call, then confirm the test passes with the updated state returned.

Written by the indexing model from the issue text.

Description

lux

GNNLux.MEGNetConv wraps ϕe and ϕv in StatefulLuxLayer{true}, but its forward pass returns the input state unchanged:

function (l::MEGNetConv)(g, x, e, ps, st)
    ϕe = StatefulLuxLayer{true}(l.ϕe, ps.ϕe, _getstate(st, :ϕe))
    ϕv = StatefulLuxLayer{true}(l.ϕv, ps.ϕv, _getstate(st, :ϕv))
    m = (; ϕe, ϕv, aggr=l.aggr)
    return GNNlib.megnet_conv(m, g, x, e), st
end

Any state update performed by a sublayer is therefore dropped. The default constructor only builds Dense chains, which are stateless, but the MEGNetConv(in_dims, out_dims, ϕe, ϕv; aggr) constructor accepts arbitrary layers. With a stateful one, e.g. BatchNorm or Dropout inside ϕe/ϕv, the running statistics (or RNG) never advance in training mode.

The same pattern is already used by EdgeConv, GINConv and, since #728, NNConv:

stnew = (; ϕe = _getstate(ϕe), ϕv = _getstate(ϕv))
return GNNlib.megnet_conv(m, g, x, e), stnew

This was found by reading the code while reviewing #728 and has not been reproduced yet; a regression test with a BatchNorm in ϕe (checking that the running stats change after a training-mode call) would pin it down. Related: #727.

The other layers with sublayer fields (CGConv, EGNNConv, GATConv, GATv2Conv, GMMConv, GatedGraphConv) have the same return-st pattern, but their constructors only create Dense/Chain of Dense sublayers (plus a GRUCell in GatedGraphConv, whose rng state doesn't appear to be used there), so nothing seems to be lost in practice.

🧔‍♂️🤖 Assisted by AI

Dominant language
Julia
Stars
310
Forks
75
Avg merge
2d 15h
Merged PRs (30d)
10

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from JuliaGraphs/GraphNeuralNetworks.jl

All issues in JuliaGraphs/GraphNeuralNetworks.jl

Similar issues

More Julia issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.