Document Classification with Graph Neural Networks

Summary: You can construct a document graph, then apply a GNN to learn document embeddings that will be fed into a softmax layer to produce a probabilistic distribution over the set of class labels.

Typical pipeline for document classification with GNNs

Step 1 — Graph Construction: Usual approach is to construct a single heterogeneous graph for the whole corpus containing word nodes and document nodes, and connect the nodes with edges based on word co-occurrence and document-word relations. For short text and non text sources you can enrich the semantics with additional information such as topics and named entities.

Step 2 — Graph Representation Learning: Encode graph structure in a way that it can be easily exploited by machine learning models. You can use various GNN models like Graph CNNs, TensorGCN, GGNN (gated graph neural network) etc.

Step 3 — Node Embedding Initialization: Mapping nodes to low-dimensional embeddings. You can use pre-trained word embeddings.

Further Reading:

1. https://arxiv.org/abs/1809.05679
2. https://arxiv.org/pdf/2106.06090.pdf
3. https://www-cs.stanford.edu/people/jure/pubs/graphrepresentation-ieee17.pdf
4. http://snap.stanford.edu/proj/embeddings-www/files/nrltutorial-part1-embeddings.pdf