[Feature Request] Support local (cluster) pooling for multi-level graph coarsening task

[lygztq]

🚀 Feature

Support graph pooling operation on cliques.

Motivation

Graph coarsening is common in tasks where we try to extract graph features from different scales. Usually, the input graph is partitioned into many small cliques and global pooling operations are performed in these cliques to obtain nodes in the output (coarser) graph. This global pooling on cliques, or local pooling on graphs, is also common in multi-level graph coarsening methods like Metis and Graclus. However, currently DGL does not support this operation.

An example (in PyG) can be found here

Alternatives

It is true that one can first convert graph in sparse representation into dense representation, then such a problem can be redefined as matrix multiplication problem: X' = C^{T}X and A' = C^{T}AC where C \in \mathcal{R}^{N_{old} \times N_{new}} is the assignment matrix. However, this is expensive for large and sparse graphs.

An implementation for PyG can also be found here

Pitch

Given two tensor c and x, where element in c is a clique label for each node (which clique does this node belongs to) and x is node feature tensor. And also given the graph g. Output the coarser graph and node features.

[jermainewang]

Hi @lygztq , thanks for the request and I agree that this looks like an good feature to have. Could you further answer the following questions?

• As an initial effort, what graph pooling functions should be included? Could you also justify a little bit on the selection (e.g., which papers are using them?)
• For each of these functions, how will the API/interface look like?
• For each of these functions, are there any existing implementations? Are they efficient enough? Is it possible (or worthwhile) to have GPU implementations?
• Are they differentiable?
• Since graph pooling functions create new (coarsened) graphs, it is similar to the graph transformation routines in DGL. It is then necessary to think about (1) how the node/edge features are transformed? (2) how the batch information is modified?

[lygztq]

Thanks for the reply. The following is my opinion:

What graph pooling functions should be included and their APIs

I think these functions should be included: cluster_max_pooling, cluster_avg_pooling, and cluster_sum_pooling (to distinguish from global pooling functions). These functions perform a "global graph pooling" in each cluster. And their interfaces should be similar:

def cluster_xxx_pooling(graph: DGLGraph, node_feats: tensor, cluster_ids: tensor) -> Tuple[DGLGraph, tensor]:
    pass

these functions return the pooled graph together with pooled node features.

Why these functions? Which papers are using them?

At first, I thought these functions could be useful in graph classification tasks where graphs are coarsened layer by layer, with a cluster assignment for each node, to obtain a final graph embedding, just like DiffPool does. However, under this scenario, pooled graphs are usually small and dense. Applying sparse version cluster pooling to these graphs is expensive. Actually both implementations of DiffPool in PyG and DGL are based on dense implementation. However, I find these cluster pooling functions are useful in the voxel grid pooling of point cloud data where we overlay a regular grid of user-defined size over a point cloud and clusters all points within the same voxel. Point cloud graphs are sparse and usually large. The paper can be found here and an example in PyG can be found here

About implementation

To the best of my knowledge, there's no implementation in DGL so far. One possible way for the implementation of these functions is based on scatter operations (for example, scatter_add for sum/avg pooling, scatter_max for max pooling). However, currently DGL has pool support on these operations. Instead, we can use SpMM (message passing) to achieve the same target. I explain this method here: #2718. The difficult parts of this method are how to construct the heterogeneous graph in an efficient and user-friendly manner, and how to get new edges in coarser graphs. All these possible implementations can be done via GPU since the underlying operations have GPU versions. I think the "differentiable" part can be done via writing autograd functions for them. For the batch information, I think we need to track the batch ID for each cluster (i.e. the node in coarser graphs).

[jermainewang]

Thanks for the answer. Some more thoughts:

• I can see two parts in the whole process: (1) aggregating node features based on assignment and (2) generate edges in the result graph (i.e., constructing the coarsened graph).
• For (1), like you said, we could use scatter_reduce or sort them by the assignment and then use spmm.
• For (2), there are many choices. For example, DiffPool creates a clique graph with weights calculated by pair-wise distance of nodes. The VoxelGrid algorithm creates sparse graph by KNN (K is usually small). Overall, this is the place we probably should leave design flexibility.
• Is there any cases where edge features also need to be aggregated?

[lygztq]

Maybe there exists works that need to aggregate edge features, but currently I have no idea.

[MozhganPourKeshavarz]

@jermainewang @lygztq Is this feature supported in DGL?

[jermainewang]

Hi @MozhganPourKeshavarz , the feature request is currently tracked in our project but the team is working on other higher priority tasks. We will post updates to this thread.