GraphCL: Graph Contrastive Learning Framework with Augmentations

Introduction

Table of Contents

👨‍🏫 Method

Figure 1: GraphCL Framework

Given a Graph $\large \mathcal{G}$ we generate two views $\large \hat{\mathcal{G}_i}, \hat{\mathcal{G}_j}$ by performing augmentations. The authors selectively learn these augmentations based on graph domains.
These two views $\large \hat{\mathcal{G}_i}, \hat{\mathcal{G}_j}$ are then passed through a graph encoder $\large f(\cdot)$ leading to representations $\large h_i, h_j$. These graph encoders can be any architecture.
These representations are then passed through a projection head $\large g(\cdot)$ a simple MLP network which generates two views $\large z_i, z_j$.
We then apply a contrastive objective $\large \mathcal{L} (\cdot)$ between the two views, the objective in this case is the normalized temperature-scaled cross entropy loss.

Node Dropping: From any given graph $\large \mathcal{G}$, we randomly "drop" some nodes along with their edges.
Edge Perturbation: This involves perturbing the edges in $\large \mathcal{G}$ through randomly adding or dropping a certain ratio of edges.
Attribute masking: Attribute masking prompts the model to recover masked node attributes using their context information, i.e., the remaining attributes.
Subgraph Generation: This involves creating a subgraph from the original graph by performing random walks.

Code

class GraphCL(torch.nn.Module):
    ... 
    def train_step(
        self,
        augmented_views: List[torch.Tensor],
    ) -> torch.Tensor:
        """
        Perform a single training step.

        Args:
            augmented_views (List[torch.Tensor]): Views generated by performing augmentations

        Returns:
            float: Loss.
        """
        # Generate Graph Views

        ## Generating representations : intermediate_reps: List[torch.Tensor]
        intermediate_reps = self.encode(augmented_views)

        ## Generate views : reps: List[torch.Tensor]
        reps = self.projection_head(intermediate_reps)

	## Calculate Loss : loss: torch.Tensor
	loss = self.contrastive_loss(reps)

        return loss

Summary