Improving Code Representation Learning via Multi-view Contrastive Graph Pooling for Abstract Syntax Tree

As the field of code intelligence continues to grow, Code representation learning has emerged as a research hot spot. Given that code structure can be naturally represented as graphs, Graph Neural Networks (GNNs) have proven highly effective for learning graph representations of source code. Pooling, as an essential operation for GNN-based models, is limited in its ability to leverage the rich hierarchical information presented in tree-like graph, especially Abstract Syntax Trees. In order to learn the graph representation of code more effectively, we propose a novel pooling method called TreePool. TreePool directly splits tree-like graphs using depth filtering based on the tree structure to form a sequence of pooled graphs sorted by descending size of subgraphs. Then local-local contrastive learning between these neighboring subgraphs is conducted to preserve the information of the graph before pooling. Through TreePool, multiple views of representation are learned and fused to obtain the final code graph representation. We conduct TreePool on a supervised framework and experimental results demonstrate that the average improvements on two real-world datasets in terms of accuracy are 1.1% and 3.3%. It also exhibits excellent performance in an unsupervised framework. Our results show that TreePool can effectively learn meaningful Abstract Syntax Tree representation of code and exhibit good performance in code classification tasks.