`EquilibriumAggregation` global aggregation layer #4522

Padarn · 2022-04-24T02:09:18Z

Implements a simple version of the paper "Equilibrium Aggregation: Encoding Sets via Optimization". Note that exact details of the paper have not been implemented yet, and I plan to leave specific details about the optimizer out of this MR, it should only add a new readout layer than implements the method.

TODO:

Example of the 'exact' methods converging (median)
Cleanup of the readout layer
Support batched datasets (similar to other glob layers)
Tests

Addresses #4447

codecov · 2022-04-24T02:12:06Z

Codecov Report

Merging #4522 (a3a3a96) into master (0c24277) will increase coverage by 0.07%.
The diff coverage is 97.80%.

@@            Coverage Diff             @@
##           master    #4522      +/-   ##
==========================================
+ Coverage   82.87%   82.95%   +0.07%     
==========================================
  Files         331      332       +1     
  Lines       18197    18288      +91     
==========================================
+ Hits        15081    15171      +90     
- Misses       3116     3117       +1

Impacted Files	Coverage Δ
torch_geometric/nn/aggr/equilibrium.py	`97.77% <97.77%> (ø)`
torch_geometric/nn/aggr/__init__.py	`100.00% <100.00%> (ø)`
torch_geometric/nn/aggr/base.py	`95.74% <0.00%> (+2.12%)`	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us.

Padarn · 2022-04-24T02:12:53Z

This is still pretty heavy WIP (and I'm not even sure if it is working 100% yet).

@rusty1s do you think there is a good example graph classification problem to try and compare accuracy to? The paper uses the MOLPCBA benchmark (would I could not see in the repo yet, though we could add it). I currently just modified the MUTAG example, there is no performance improvement (or degradation).

I plan also to have the median example from the paper as an example (acts as a reasonable test it works).

torch_geometric/nn/glob/equilibrium_aggregation.py

examples/mutag_gin_equib.py

torch_geometric/nn/glob/__init__.py

torch_geometric/nn/glob/equilibrium_aggregation.py

Padarn · 2022-04-29T14:23:18Z

I've updated this, and just added a simple example of the median readout learning. I didn't have time to do a larger example yet sorry.

Padarn · 2022-05-06T01:39:23Z

Hey @rusty1s this is mostly ready for review (though the example I use is of the median aggregation).

I'm currently trying to play around to get this to converge well, but I haven't had much luck. I'm not sure if its because of parameters I've used or details missing in the paper (like the alpha used in the inner optimization). But perhaps convergence is just slow - they run for 10 million steps which I'm trying to do now but will take quite some time on my resources :-)

Padarn · 2022-05-06T07:39:32Z

After ~1 million iterations no signs of convergence, guessing I've got something wrong. Maybe @FabianFuchsML has some input?

Padarn · 2022-05-14T07:31:39Z

torch_geometric/nn/glob/equilibrium_aggregation.py

+        momentum = torch.zeros_like(y)
+        for _ in range(iterations):
+            val = func(x, y, batch)
+            grad = torch.autograd.grad(val, y, create_graph=True,


The authors of the paper suggest adding an axillary loss on the grad here (grad.square().sum()/iterations) to be precise.

The least intrusive way I could think to do this would be to accumulate the grad in a property like:

grad = ... self.aux_loss += grad.square().sum()/iterations

then to call backwards on it by using a hook? I'm not eventually sure this makes sense.

Any suggestions?

Padarn · 2022-05-28T05:50:42Z

I've moved this to the aggregation module but there are a few things missing which I'll need opinions on. Leaving comments in the code.

Padarn · 2022-05-28T05:51:18Z

torch_geometric/nn/aggr/equilibrium.py

+    def energy(self, x: Tensor, y: Tensor, index: Optional[Tensor]):
+        return self.potential(x, y, index) + self.reg(y)
+
+    def forward(self, x: Tensor, index: Optional[Tensor] = None, *,


Just a note this was all written before the Aggregation module, so its not using reduce yet - need to rethink some of the implementation to make use of it.

No need to make use of reduce. The LSTMAggregation module doesn't make use of it as well.

Padarn · 2022-05-28T05:52:48Z

torch_geometric/nn/aggr/equilibrium.py

+        if ptr is not None:
+            raise ValueError(f"{self.__class__} doesn't support `ptr`")
+
+        if dim_size is not None:


I could technically support this, but I don't know what you'd expect the behavior to be: Unlike sum/mean etc, we can't just assume the input is zero (or if we do, the output would just be random).

Thoughts?

Please have a look at https://pytorch-scatter.readthedocs.io/en/latest/functions/scatter.html. The dim_size argument ideally is identical to index_size = index.max() + 1 in your case. If passed, there is no need to compute it in the first place.

I didn't fully get that sorry. If dim_size is passed

if dim_size < index_size I'd expect error,

if dim_size > index_size I'm not sure what to expect - I guess zero in for those entries?

Yes, that is correct.

lightaime

Great work @Padarn! Added some minor comments.

test/nn/aggr/test_equilibrium.py

lightaime · 2022-05-29T14:40:52Z

examples/equilibrium_median.py

+`EquilibriumAggregation` to learn to take the median of
+a set of numbers
+"""
+


It would be great to add a few comments about the convergence of EquilibriumAggregation.

torch_geometric/nn/aggr/__init__.py

lightaime · 2022-05-29T15:00:30Z

examples/equilibrium_median.py

+    dist = np.random.choice([norm, gamma, uniform])
+    x = dist.sample((input_size, 1))
+    y = model(x)
+    loss = (y - x.median()).norm(2)


It seems the loss is not normalized by the input size.

Suggested change

loss = (y - x.median()).norm(2)

loss = (y - x.median()).norm(2) / input_size

torch_geometric/nn/aggr/equilibrium.py

Padarn · 2022-06-01T01:01:22Z

thanks @lightaime - I've got a busy week, I'll address your comments ASAP :-)

Padarn · 2022-06-04T05:31:49Z

Thanks for the reviews - I've updated based on most comments. I still have an uncertainty here: #4522 (comment) but it maybe for another PR

rusty1s · 2022-06-06T06:30:16Z

Thanks @Padarn. Can you also resolve the merge conflicts? I will take a look soon, and address the unresolved comment as well.

Padarn · 2022-06-25T04:25:40Z

Hey @rusty1s, any thoughts on this one? Happy to keep working on it or break it up, but might be good to not leave it here unfinished.

rusty1s · 2022-06-25T11:17:25Z

Yes, now that #4779 is merged, let's try to integrate this next :)

Padarn · 2022-06-25T13:24:20Z

Cool. I'll make another pass to see if there is anything I can clean up based on #4779

for more information, see https://pre-commit.ci

Co-authored-by: Guohao Li <lightaime@gmail.com>

for more information, see https://pre-commit.ci

Padarn · 2022-07-26T13:29:18Z

hey @rusty1s and @lightaime what do you guys think about merging this one?

lightaime

Thanks, LGTM!

Padarn · 2022-07-27T00:12:54Z

Thanks @lightaime - I'm going to merge this one, but will try and build some more example use cases before we release the new version so make sure its good.

Padarn commented Apr 24, 2022

View reviewed changes

torch_geometric/nn/glob/equilibrium_aggregation.py Outdated Show resolved Hide resolved

Padarn added the feature label Apr 24, 2022

Padarn self-assigned this Apr 24, 2022

rusty1s added the 0 - Priority P0 label Apr 25, 2022

rusty1s reviewed Apr 25, 2022

View reviewed changes

Padarn changed the title ~~[WIP] EquilibriumAggregation global aggregation layer~~ EquilibriumAggregation global aggregation layer Apr 29, 2022

Padarn commented May 14, 2022

View reviewed changes

rusty1s mentioned this pull request May 25, 2022

[Roadmap] torch_geometric.nn.aggr 🚀 #4712

Closed

26 tasks

Padarn force-pushed the padarn/optim-embedding branch from b32b0fa to b3438f2 Compare May 28, 2022 05:11

Padarn requested review from lightaime and RexYing May 28, 2022 05:50

Padarn commented May 28, 2022

View reviewed changes

lightaime reviewed May 29, 2022

View reviewed changes

Padarn and others added 4 commits July 9, 2022 17:49

wip equilibrium aggregation

fbafde5

update with docs and batch

3c1d0ba

[pre-commit.ci] auto fixes from pre-commit.com hooks

9b66232

for more information, see https://pre-commit.ci

example of median

fcef12c

Padarn and others added 19 commits July 9, 2022 17:49

update init

b21760c

update changelog

58e1678

add details to median example

3b1bc45

match parameters from paper

a02eea8

match parameters from paper

3f0a23b

wip

8d79424

add nesterov loss

3f1216a

change loss to be a mean across a single batch

894a50f

move to aggreation package

5091ffe

move example import

ed94b33

add assertions for unsupported

3b3aaca

update changelog

99f020f

update changelog

8561c8b

add suggestions and test update

8460438

Update test/nn/aggr/test_equilibrium.py

6a38ff9

Co-authored-by: Guohao Li <lightaime@gmail.com>

Update test/nn/aggr/test_equilibrium.py

b34b3a5

Co-authored-by: Guohao Li <lightaime@gmail.com>

fix error message

4171372

fix indent error

cc0eac1

rebase to fix tests

fa47205

Padarn force-pushed the padarn/optim-embedding branch from bd4bb93 to fa47205 Compare July 9, 2022 09:52

minor update for agg class

7359edc

Padarn force-pushed the padarn/optim-embedding branch from 6656ea1 to 7359edc Compare July 9, 2022 09:58

pre-commit-ci bot and others added 2 commits July 9, 2022 09:59

[pre-commit.ci] auto fixes from pre-commit.com hooks

2236dc1

for more information, see https://pre-commit.ci

Merge branch 'master' into padarn/optim-embedding

d504053

lightaime approved these changes Jul 26, 2022

View reviewed changes

Merge branch 'master' into padarn/optim-embedding

a3a3a96

Padarn merged commit 333d3d3 into pyg-team:master Jul 27, 2022

Padarn deleted the padarn/optim-embedding branch July 27, 2022 07:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

`EquilibriumAggregation` global aggregation layer #4522

`EquilibriumAggregation` global aggregation layer #4522

Padarn commented Apr 24, 2022 •

edited

Loading

codecov bot commented Apr 24, 2022 •

edited

Loading

Padarn commented Apr 24, 2022 •

edited

Loading

Padarn commented Apr 29, 2022

Padarn commented May 6, 2022

Padarn commented May 6, 2022 •

edited

Loading

Padarn May 14, 2022 •

edited

Loading

Padarn commented May 28, 2022

Padarn May 28, 2022

rusty1s May 29, 2022

Padarn May 28, 2022

rusty1s May 29, 2022

Padarn May 29, 2022

rusty1s May 29, 2022

lightaime left a comment •

edited

Loading

lightaime May 29, 2022

lightaime May 29, 2022

Padarn commented Jun 1, 2022

Padarn commented Jun 4, 2022

rusty1s commented Jun 6, 2022

Padarn commented Jun 25, 2022

rusty1s commented Jun 25, 2022 •

edited

Loading

Padarn commented Jun 25, 2022

Padarn commented Jul 26, 2022

lightaime left a comment

Padarn commented Jul 27, 2022

	loss = (y - x.median()).norm(2)
	loss = (y - x.median()).norm(2) / input_size

EquilibriumAggregation global aggregation layer #4522

EquilibriumAggregation global aggregation layer #4522

Conversation

Padarn commented Apr 24, 2022 • edited Loading

codecov bot commented Apr 24, 2022 • edited Loading

Codecov Report

Padarn commented Apr 24, 2022 • edited Loading

Padarn commented Apr 29, 2022

Padarn commented May 6, 2022

Padarn commented May 6, 2022 • edited Loading

Padarn May 14, 2022 • edited Loading

Choose a reason for hiding this comment

Padarn commented May 28, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lightaime left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Padarn commented Jun 1, 2022

Padarn commented Jun 4, 2022

rusty1s commented Jun 6, 2022

Padarn commented Jun 25, 2022

rusty1s commented Jun 25, 2022 • edited Loading

Padarn commented Jun 25, 2022

Padarn commented Jul 26, 2022

lightaime left a comment

Choose a reason for hiding this comment

Padarn commented Jul 27, 2022

`EquilibriumAggregation` global aggregation layer #4522

`EquilibriumAggregation` global aggregation layer #4522

Padarn commented Apr 24, 2022 •

edited

Loading

codecov bot commented Apr 24, 2022 •

edited

Loading

Padarn commented Apr 24, 2022 •

edited

Loading

Padarn commented May 6, 2022 •

edited

Loading

Padarn May 14, 2022 •

edited

Loading

lightaime left a comment •

edited

Loading

rusty1s commented Jun 25, 2022 •

edited

Loading