chr26195 / pmlp Goto Github PK

This is the official code repository for "Graph Neural Networks are Inherently Good Generalizers: Insights by Bridging GNNs and MLPs", which is accepted to ICLR 2023.

Related materials: paper, slides, poster

• 0. What Could We Do with PMLP
• 1. Quick Guide of Implementation
• 2. PMLP Extensions and FAQ
• 3. Run the Code

• Accelerate GNN training by modifying only a few lines of codes.
• Empower MLP (or any other backbone models) by incorporating message passing / graph convolution in inference.
• Empower GNN in some scenarios, e.g., datasets with many noisy structures.
• Simple and useful tool for research and scientific discovery.

The implementation of PMLP is very simple, and can be plugged into one's own pipeline by modifying only a few lines of codes. The key idea of PMLP is to just remove message passing modules in GNNs during training. We allow the model after removal of message passing layers to be other models, such as ResNet (corresponding to GCNII) and MLP+JK (corresponding to JKNet). Here we introduce several different ways to implement PMLP and discuss their advantages and limitations.

This is the default way and go-to choice to implement PMLP, which combines three models (MLP, PMLP, GNN) in one single class. The key idea of this implementation is to add a use_conv = True/False parameter in the self.forward() function for any GNN classes. To implement PMLP, just set this parameter to be False in training and validation, and then reset it to be True in testing. For example

# version A: three models (mlp, pmlp, gnn) in one class
class My_GNN(nn.Module):
    ...
    def forward(self, x, edges, use_conv = True):
        ...
        x = self.feed_forward_layer(x) 
        if use_conv: 
            x = self.message_passing_layer(x, edges)  
        ...
        return x

my_gnn = My_GNN()

# in the training loop
my_gnn.train()
for epoch in range(args.epochs):
    prediction = my_gnn(x, edge, use_conv = False)

# for inference
my_gnn.eval()
prediction = my_gnn(x, edge, use_conv = True)

This implementation is very flexible. If use_conv = True is both training and testing, the model is equivalent to the original GNN. And if use_conv = False is both training and testing, the model is equivalent to the original MLP (or other 'backbone' models). One can optionally let use_conv = True for validation by modifying the evaluate() function in data_utils.py, which could lead to better or worse performance depending on the specific task and could be treated as a hyperparameter to tune. In default, we do not use message passing in validation such that the model selection process of PMLP is also equivalent to MLP. Many extensions of PMLP can be developed upon this implementation for adapting to new tasks or further boosting performance. For example, we can specify def if_use_conv(*args) as a function of different inputs (e.g., whether the model is training, whether the loss function has reached a certain threshold, current training epoch, etc.) to control concretely when and where to use message passing layers.

Here is the most simple way to implement PMLP, which only requires modifying only one line of code in your own GNN pipeline. This implementation leverages the (PyTorch) built-in parameter self.training to automatically specify when to use graph convolution.

# version B: one line of code is all you need
class My_GNN(nn.Module):
    ...
    def forward(self, x, edges):
        ...
        x = self.feed_forward_layer(x) 
        if not self.training: # only modified part
            x = self.message_passing_layer(x, edges)  
        ...
        return x

my_gnn = My_GNN()

# in the training loop
my_gnn.train()
for epoch in range(args.epochs):
    prediction = my_gnn(x, edges)

# for inference
my_gnn.eval()
prediction = my_gnn(x, edges)

One limitation of this implementation is that it enforces using message passing (i.e., the GNN architecture) for validation (since usually we use model.eval() before validation and self.training is automatically set to be False). However, this is totally acceptable in most scenarios since validation is not the bottleneck of training efficiency and using message passing in validation often improves performance.

One equivalent way to implement PMLP is to drop all edges but leave only self-loops in training such that message passing operation will not affect node representations. This is an extreme case of DropEdge. Please refer to their paper for more information.

# version C: just drop all edges (but leave self-loops alone)
class My_GNN(nn.Module):
    ...
    def forward(self, x, edges):
        ...
        x = self.feed_forward_layer(x) 
        x = self.message_passing_layer(x, edges)  
        ...
        return x

my_gnn = My_GNN()

# create a graph with only self-loops
self_loops = torch.nonzero(torch.eye(node_numbers)).t()

# in the training loop
my_gnn.train()
for epoch in range(args.epochs):
    prediction = my_gnn(x, self_loops)

# for inference
my_gnn.eval()
prediction = my_gnn(x, edges)

This implementation makes PMLP very easy to adapt since one does not need to modify anything on the original GNN class. But please note that this implementation is LESS efficient than other PMLP versions as we have observed in practice, presumably because it still relies on the messaga passing layer. Please also be careful that it is possible that for some specific message passing implementations, the corresponding transition matrix for self-loops is not an indentity matrix, and then this version would not be exactly equivalent to PMLP.

Another equivalent way to implement PMLP is to define two models, i.e., MLP model and GNN model. We first pretrain the MLP model, save the state_dict(), then load it to the GNN model, and finally use it directly for inference or do whatever we want on top of it.

# version D: Load Pretrained MLP
class My_GNN(nn.Module):
    ...
    def forward(self, x, edges):
        ...
        x = self.feed_forward_layer(x) 
        x = self.message_passing_layer(x, edges)  
        ...
        return x

class MLP(nn.Module):
    ...
    def forward(self, x, edges):
        ...
        x = self.feed_forward_layer(x) 
        ...
        return x

model_mlp = MLP()
model_gnn = My_GNN()

# in the training loop
model_mlp.train()
for epoch in range(args.epochs):
    prediction = model_mlp(x)
    ...
torch.save(model_mlp.state_dict(), PATH)

# for inference
model_gnn.load_state_dict(PATH)
model_gnn.eval()
prediction = model_gnn(x, edges)

This version is more complicated than others in terms of implementation but could be suitable for empowering pre-trained MLP (or any other backbone models) by incorporating message passing / graph convolution in inference and some other scenarios.

This section summarizes frequently asked questions and will continously update.

Many existing GNNs consist of parameterized message passing layers. For those cases, one can consider set a two-stage training. For the first stage of training, one can remove message passing layers such that the model is equivalent to MLP or other backbone NN architectures. For the second stage of training, one can add those layers (with trainable parameters) back and fine-tune the model using a few epochs. Here is a simple example code where this procedure is implemented in an end-to-end manner:

# in the training loop
my_gnn.train()
for epoch in range(1000):
    prediction = my_gnn(x, edge, use_conv = False if epoch < 900 else True)

# for inference
my_gnn.eval()
prediction = my_gnn(x, edge, use_conv = True)

Such a solution was inspired by a concurrent work MLPInit. Please refer to their paper for more information.

Indeed, there exist GNNs whose layers are hard to be disentangled into MP layers and FF layers. In such cases, one can consider using our Version C (Just Drop All Edges) which replaces the original adjacency matrix with an indentity matrix, and is equivalent to dropping all edges (except self-loops) in the graph. But note again that this version might not be as efficient as other versions. One can optimize the code accordingly in such cases.

PMLP can be extended to transductive learning by flexibly choosing various approaches to incorporate unlabeled nodes in training such as the one described in Q1. Particularly, we found that generating pseudo labels for unlabeled nodes (in a non-differiatable way) as data augmentation can further improve the performance of PMLP. The extent to which performance is improved depends on the specific task and the approach used for label generation. Here ia an example code for reference

if args.trans == True:
    out_prob = F.softmax(out, dim=1)
    true_label = F.one_hot(dataset.label, dataset.label.max() + 1).squeeze(1).to(torch.float)
    train_mask = torch.zeros(out.shape[0]).bool().to(device)
    train_mask[train_idx] = True

    pseudo_label = torch.empty_like(true_label)
    pseudo_label[train_mask] = true_label[train_mask]
    pseudo_label[~train_mask] = out_prob[~train_mask].detach()
    pseudo_label[~train_mask] = gcn_conv(pseudo_label, dataset.graph['edge_index'])[~train_mask]
    pseudo_label[~train_mask] = pseudo_label[~train_mask] / pseudo_label[~train_mask].sum(dim=-1, keepdim=True)
    
    loss = nn.KLDivLoss(reduction='batchmean')(F.log_softmax(out, dim=1), pseudo_label)

The main focus of our paper is on node classification. Though we have not tested PMLP on other GNN-related tasks, we encourage you to explore its potential in other tasks. Feel free to try it out and see how it performs.

It is worth noting that these GNN-related problems are orthogonal to PMLP framework. Whether PMLP can solve these issues depends on the design of its GNN counterpart. In our paper's Appendix, we have provided some discussions on these issues. While we have not conducted an in-depth investigation into these GNN-related problems, we encourage you to leverage PMLP as an analytical tool or explore its potential in your research.

Step 1. Install the required package according to requirements.txt.

Step 2. Specify your own data path in parse.py and download the datasets.

Step 3. To run the training and evaluation, one can use the following scripts (for version A).

Step 4. Results will be saved in a folder named results

# GCN: use message passing in training, validation and testing
python main.py --dataset cora --method pmlp_gcn --protocol semi --lr 0.1 --weight_decay 0.01 --dropout 0.5 --num_layers 2 --hidden_channels 64 --induc --device 0 --conv_tr --conv_va --conv_te 

# PMLP_GCN: use message passing only in testing
python main.py --dataset cora --method pmlp_gcn --protocol semi --lr 0.1 --weight_decay 0.01 --dropout 0.5 --num_layers 2 --hidden_channels 64 --induc --device 0 --conv_te 

# MLP: not using message passing
python main.py --dataset cora --method pmlp_gcn --protocol semi --lr 0.1 --weight_decay 0.01 --dropout 0.5 --num_layers 2 --hidden_channels 64 --induc --device 0

--induc and --trans are used to specify inductive or transductive learning settings.