[Roadmap] torch_geometric.nn.aggr 🚀

[rusty1s]

🚀 The feature, motivation and pitch

The goal of this roadmap is to unify the concepts of aggregation inside GNNs across both MessagePassing and global readouts. Currently, these concepts are separated, e.g., via MessagePassing.aggr = "mean" and global_mean_pool(...) while the underlying implementation is the same. In addition, some aggregations are only available as global pooling operators (global_sort_pool, Set2Set, ...), while, in theory, they are also applicable during MessagePassing (and vice versa, e.g., SAGEConv.aggr = "lstm"). One additional feature is the combination of aggregations, which is a useful feature both in MessagePassing (PNAConv, EGConv, ...) and global readouts.

As such, we want to provide re-usable aggregations as part of a newly defined torch_geometric.nn.aggr.* package. Unifying these concepts also helps us to perform optimization and specialized implementations in a single place (e.g., fused kernels for multiple aggregations). After integration, the following functionality is applicable:

class MyConv(MessagePassing):
    def __init__(self):
        super().__init__(aggr="mean")

class MyConv(MessagePassing):
    def __init__(self):
        super().__init__(aggr=LSTMAggr(channels=...))

class MyConv(MessagePassing):
    def __init__(self):
        super().__init__(aggr=MultiAggr("mean", "max", Set2Set(channels=...))

Roadmap

The general roadmap looks as follows (at best, each implemented in a separate PR):

• Define torch_geometric.nn.aggr.* and implement a BaseAggr abstract class (torch_geometric.nn.aggr package with base class #4687)
• Add new aggregators:
  • DeepGCN aggregations SoftmaxAggr, PowerMeanAggr, cf., GENConv (torch_geometric.nn.aggr package with base class #4687)
  • LSTMAggr, cf., SAGEConv (LSTMAggregation #4731)
• Allow for multiple aggregations as part of a MultiAggr class (MultiAggregation and aggregation_resolver #4749)
• Add support for class-resolver, similar to here (MultiAggregation and aggregation_resolver #4749, Add class-resolver for Aggregation #4716)
• Ensure torch.jit.script support (Integration of nn.aggr within MessagePassing #4779)
• Integrate new aggregations into MessagePassing interface (Integration of nn.aggr within MessagePassing #4779)
• Move aggregators from torch_geometric.nn.glob to torch_geometric.nn.aggr (respecting the new interface), deprecate old implementation:
  • MeanAggr, SumAggr, MaxAggr, MinAggr, MulAggr, VarAggr, StdAggr (torch_geometric.nn.aggr package with base class #4687, MultiAggregation and aggregation_resolver #4749)
  • MedianAggr (Add MedianAggregation and QuantileAggregation #5098)
  • AttentionalAggr (Move nn.glob.attention.GlobalAttention to nn.aggr.attention.AttentionalAggregation #4986)
  • Set2Set (nn.aggr.Set2Set #4762)
  • GlobalSortAggr (Add GlobalSortAggr to nn.aggr package. #4957)
  • GraphMultiSetTransformer (Move GraphMultisetTransformer to nn.aggr package #4973)
  • EquilibriumAggr (EquilibriumAggregation global aggregation layer #4522)
  • Deprecate torch_geometric.nn.glob (Deprecate nn.glob package #5039)
• Update existing GNN layers to make use of new interface, e.g.:
  • SAGEConv (Refactor SAGEConv to use LSTMAggregation #4863)
  • PNAConv (Refactor PNAConv to rely on new Aggregation #4864)
  • GravNetConv (Refactor GravNetConv to rely on new Aggregation #4865)
  • GENConv (Refactor GENConv to rely on new Aggregation #4866)
• Add support for "reverse" aggregation resolver to keep message_and_aggregate functionality intact (Add reverse support in aggregation_resolver #5084)
• Support for multiple aggregations in SAGEConv (Support for multiple aggregations in SAGEConv #5033)
• MultiAggregation: Support for concat, concat+transform, sum, mean, max, attention (Add combine support to MultiAggregation #5000, Raise error for one single aggregation when proj or attn mode is used in MultiAggregation #5034)
• Add semi_grad functionality to SoftmaxAggregation (Add semi_grad and docs to SoftmaxAggregtion #4995)
• Add and verify torch_geometric.nn.aggr.* documentation (Update SoftmaxAggregation and PowerMeanAggregation doc strings #5036, Fix documentation and README.md regarding new nn.aggr package #5097, Update nn.aggr documentation #5099, Add missing doc-strings to nn.aggr #5104)
• Add a tutorial on the new concepts (Add notebook tutorial for torch_geometric.nn.aggr package #4927)
• Kernel fusion: Optimize aggregations, e.g., by computing multiple aggregations in parallel (at best discussed in a separate issue)

Any feedback and help from the community is highly appreciated!

cc: @lightaime @Padarn

[Padarn]

Looks great @rusty1s! I'll try to pick up some of the smaller tasks this weekend

[lightaime]

Added some simple ones. MaxAggr, MinAggr, SumAggr, SoftmaxAggr, PowermeanAggr, VarAggr, StdAggr.

[Padarn]

I planned to pick up a couple of the tasks - hope you guys don't mind me editing the issue to make it clear what I plan on doing (I'll sick to smaller PRs if possible since I typically don't have much time during weeks)

[Padarn]

Kernel fusion: Optimize aggregations, e.g., by computing multiple aggregations in parallel (at best discussed in a separate issue)

Do we have an open issue on this? I'd be interested to understand a bit more about what we're thinking here. Is it mostly for the case where we want to (for example) compute both a sum and mean?

[rusty1s]

Yes, indeed. There do not exist clear plans for implementation yet though. It will likely depend on PyTorch fusing this ops as part of TorchScript, or on us providing special CUDA kernels.

[rusty1s]

Thanks everyone for the hard work @lightaime @Padarn. I think that the final outcome looks fantastic - many cool things to promote in our upcoming release :)