BertGCN: Transductive Text Classification by Combining GCN and BERT

In plain words

Text classification assigns a label to each document. Large pretrained language models do this well, and graph methods do it by linking related documents, including unlabelled ones. Each approach has a gap. Graph methods start from weak representations of the documents, and pretrained models do not use the links between documents.

BertGCN combines the two. It builds a graph in which documents and words are the nodes, and it starts each document node from a pretrained language model's representation of that document. A graph network then passes information along the links, so a document's label estimate is informed by its neighbours, labelled or not. The language model and the graph network are trained together.

The point is that the model keeps the strength of a large pretrained reader and adds the signal that comes from the whole collection of documents.

TL;DR

BertGCN is a text classifier that puts BERT-style document embeddings into a word-document graph and runs a two-layer GCN over it §3.1. BERT and the GCN are trained jointly, with a memory bank that keeps every document embedding up to date §3.3. The GCN prediction is mixed with a BERT-only prediction using a weight lambda §3.2. The authors report state-of-the-art test accuracy on five benchmarks, averaged over 10 runs tab. 1. Code is available, but the repository has no license file.

Problem

Transductive learning uses unlabelled test documents during training. Graph neural networks have been applied to it, for example TextGCN, but those models initialise the nodes with one-hot identity vectors and do not use large-scale pretraining §1 §2. The paper argues that the benefit of pretraining for transductive text classification was unclear §1.

Combining a pretrained encoder with a graph over the whole corpus is hard in practice. Full-batch GCN training needs every node's representation at once, which does not fit in memory when the encoder is BERT §3.3. The paper's setting assumes a fixed corpus with a known split into training and test documents, and it evaluates on that fixed split §4.1 app. A.

Method

BertGCN has three parts: a heterogeneous text graph, a GCN with BERT-initialised document nodes, and a training scheme with a memory bank.

Graph §3.1. The graph follows TextGCN and contains document and word nodes §3.1. Word-word edges carry positive pointwise mutual information (PPMI), document-word edges carry TF-IDF weights, self-loops have weight 1, and all other pairs have weight 0 eq. 1.

A_{i,j} = \begin{cases} \operatorname{PPMI}(i,j) & i, j \text{ are words and } i \neq j \\ \operatorname{TF\text{-}IDF}(i,j) & i \text{ is a document, } j \text{ is a word} \\ 1 & i = j \\ 0 & \text{otherwise} \end{cases} eq. 1

Node features §3.1. Document rows of the feature matrix are the document embeddings from a BERT-style model (BERT or RoBERTa). Word rows are zero. The document embeddings are the [CLS] outputs, as stated in §4.1 [code: train_bert_gcn.py#L143].

X = \begin{pmatrix} X_{\text{doc}} \\ 0 \end{pmatrix}_{(n_{\text{doc}} + n_{\text{word}}) \times d} eq. 2

GCN §3.1. The features pass through a GCN over the whole graph, which mixes messages between labelled and unlabelled documents. Each layer applies a normalised adjacency matrix, a weight matrix and an activation ρ. The paper writes this for layer i, starting from L⁽⁰⁾ = X eq. 3. The output of the GCN is passed to a softmax eq. 4. The loss is cross entropy over labelled document nodes §3.1.

L^{(i)} = \rho\left(\tilde{A}\, L^{(i-1)} W^{(i)}\right), \qquad L^{(0)} = X eq. 3

Z_{\text{GCN}} = \operatorname{softmax}\left(g(X, A)\right) eq. 4

Auxiliary BERT prediction and interpolation §3.2. A second prediction uses a softmax layer directly on the document embeddings eq. 5. The final prediction is a linear interpolation of the two, with weight lambda eq. 6. Lambda = 1 gives the full BertGCN model, and lambda = 0 gives BERT alone §3.2. The authors say the auxiliary head makes the multi-layer GCN input better regulated and speeds up convergence §3.2.

Z_{\text{BERT}} = \operatorname{softmax}(W X) eq. 5

Z = \lambda Z_{\text{GCN}} + (1 - \lambda) Z_{\text{BERT}} eq. 6

Memory bank §3.3. BERT cannot process the whole graph at once, so the method keeps a memory M with one row per document and uses it as GCN input. At the start of each epoch, every document embedding is recomputed with the current BERT and stored in M. In each iteration, a minibatch of labelled and unlabelled documents is drawn. Their embeddings are recomputed with the current BERT and written into M. The GCN then runs on the updated M. During back-propagation, M is treated as constant except for the rows of the minibatch §3.3. Because the stored embeddings come from different BERT states within an epoch, the BERT learning rate is kept small. BERT is also initialised from a copy fine-tuned on the target dataset before training §3.3.

Training loop. The steps below combine §3.1 to §3.3 and the optimiser setting in §4.1. The paper has no algorithm box; this is our condensed version.

# Build the graph once: word and document nodes  [§3.1] [eq. 1]
graph = build_text_graph(documents, vocabulary)
# Start BERT from a copy fine-tuned on the target data  [§3.3]
bert = fine_tune(pretrained_bert, labelled_documents)
# Two-layer GCN; hidden size 200 comes from code, not the paper
gcn = two_layer_gcn(in_dim=d, hidden_dim=200, num_classes=C)
aux = linear(in_dim=d, num_classes=C)     # auxiliary BERT head  [§3.2]
lam = 0.7                                 # mixing weight  [§4.3] [fig. 1]

for epoch in range(num_epochs):
    # Refresh the memory with every document, using the current BERT  [§3.3]
    memory = bert.embed(all_documents)    # (n_doc, d)
    for batch in shuffled_minibatches(all_documents):  # [§3.3]
        # Recompute the batch rows with gradients and write them back  [§3.3]
        batch_embed = bert.embed(batch.documents)      # (B, d)
        memory[batch.index] = batch_embed
        # GCN over the graph with the updated memory  [§3.1] [eq. 3] [eq. 4]
        gcn_probs = softmax(gcn(graph, memory))        # (n_nodes, C)
        # Auxiliary prediction from the batch embeddings  [§3.2] [eq. 5]
        bert_probs = softmax(aux(batch_embed))         # (B, C)
        # Interpolate the two predictions  [§3.2] [eq. 6]
        probs = lam * gcn_probs[batch.index] + (1 - lam) * bert_probs
        # Cross entropy only on labelled documents in the batch  [§3.1]
        labelled = batch.is_labelled
        loss = cross_entropy(probs[labelled], batch.labels[labelled])
        # Update BERT, the auxiliary head and the GCN  [§4.1]
        optimizer_step(loss, bert, aux, gcn)

Hyperparameters

Only values the paper states are given in the "Value" column for the headline model. Everything marked "code" was read from the repository at the commit above and is not stated in the paper.

Results

Author-reported. Setup: five benchmarks (20NG, R8, R52, Ohsumed, MR), BERT-base or RoBERTa-base, two-layer GCN, mean test accuracy over 10 runs; baselines are TextGCN, SGC, BERT and RoBERTa, plus GAT variants §4.1 tab. 1. Hardware and training time are not stated in the main text.

Test accuracy, tab. 1:

Ablation on the 20NG development set, RoBERTaGCN tab. 2. The development set is sampled from the original training set, so these numbers are not comparable with the test numbers above fig. 1 note §4.3:

Effect of lambda on 20NG development accuracy for RoBERTaGCN fig. 1: accuracy rises with larger lambda; the best value is 0.7, slightly above lambda = 1 (GCN only) §4.3.

Pattern in the text: pretrained models alone beat the GCN-only baselines on the short-document datasets R52 and MR, while the graph gives the largest gains on 20NG and Ohsumed, which have long documents §4.2. The authors attribute this to the word-document statistics of the graph §4.2.

Paper vs code

Repository: https://github.com/ZeroRin/BertGCN, the link given in the paper's abstract footnote. The commit is 43babaf3ab0a535d3bf432cdc632e6b0f4aa0134. The repository has no LICENSE file, so code_license is recorded as none.

• Default encoder. Paper: BertGCN uses BERT-base §4.1. Code: the default is roberta-base train_bert_gcn.py. Running the defaults reproduces RoBERTaGCN, not BertGCN.
• Learning rate of the auxiliary head. Paper: 1e-5 for the fine-tuned BERT module §4.1. Code: the same optimizer group also covers the classifier head train_bert_gcn.py.
• Adjacency normalisation. Paper: "normalized adjacency matrix" with no formula §3.1. Code: symmetric normalisation of A + I train_bert_gcn.py utils.py.
• Document-word weight. Paper: TF-IDF §3.1 eq. 1. Code: raw term count times log(N / df), not normalised build_graph.py.
• Word-word weight. Paper: PPMI §3.1. Code matches: only positive PMI values are kept build_graph.py.
• Activation. Paper: ρ, not named §3.1 eq. 3. Code: ELU on the hidden layer, no activation on the output layer models.py torch_gcn.py.
• Interpolation. Paper and code agree on mixing the two probabilities before the loss §3.2 eq. 6. The code takes the log of the mix with a 1e-10 offset models.py.
• Memory bank refresh. Paper: current BERT, with recomputation per epoch and per minibatch §3.3. Code: the per-epoch refresh runs BERT in eval mode with dropout off train_bert_gcn.py train_bert_gcn.py, while the per-batch update runs BERT in train mode models.py. The paper does not state the mode.
• Loss and validation. Paper: cross entropy on labelled nodes §3.1. Code: validation labels are excluded from the loss and only choose the checkpoint with the best validation accuracy train_bert_gcn.py train_bert_gcn.py. The paper does not say which epoch the reported test accuracy comes from.
• Data split. The paper does not describe a validation split. Code holds out the last 10% of the official training documents after a shuffle with no fixed seed build_graph.py build_graph.py.
• Fine-tuning stage. Paper: fine-tune before training, no settings §3.3. Code: separate script with lr 1e-4, batch 128 and 60 epochs finetune_bert.py.
• Source bundle. The LaTeX bundle contains unused drafts (ablation.tex, experiments.tex, method.tex) with different ablation numbers. This card follows the main file transductive.tex only.

Pitfalls

• Encoder default. The headline numbers use BERT-base, but the default encoder in code is RoBERTa-base. Pass --bert_init bert-base-uncased for the BertGCN row train_bert_gcn.py.
• Long inputs are truncated. BERT sees at most 128 tokens per document, while the graph uses the full text train_bert_gcn.py train_bert_gcn.py. The paper credits the graph for the gains on long documents §4.2 and does not mention truncation.
• Stale and inconsistent memory. Stored embeddings come from different BERT states within an epoch. The paper's remedy is a small BERT learning rate and fine-tuned initialisation §3.3 §4.3. Without the small rate, training did not converge in the reported ablation tab. 2.
• Dropout mode mismatch. Batch embeddings use dropout and the per-epoch memory does not models.py train_bert_gcn.py.
• Lowercased, cleaned text. BERT receives the TextGCN-cleaned text, which is lowercased build_graph.py utils.py. This matters for cased encoders such as RoBERTa.
• Random split with no seed. Splits and the whole pipeline have no fixed seed, so results depend on the run build_graph.py. The paper reports 10 runs without seeds.
• Unstated hyperparameters. Optimiser, schedule, epochs, batch size, hidden size, dropout, activation and maximum length are only in code (see the table). An implementer following the paper alone has to choose them.
• Lambda transfer. Lambda = 0.7 was chosen on the 20NG development set for RoBERTaGCN §4.3 fig. 1. The code uses it for every dataset by default train_bert_gcn.py.
• Development numbers. Table 2 reports development accuracy, not test accuracy tab. 2.
• GAT variants. GAT hyperparameters are not in the paper; the code uses 8 heads and a hidden size of 200 per head by default train_bert_gcn.py models.py.
• Code license. The repository has no license, so the code may not be reused or quoted under our rules. This card cites the code only and copies no code.
• Cost. Each epoch runs BERT over every document, labelled or not, and the memory bank holds one embedding per document. Hardware and runtime are not stated.

Lineage

• Builds on: TextGCN, the graph construction and preprocessing protocol §3.1 §4.1 [arXiv:1809.05679]; GCN torch_gcn.py [arXiv:1609.02907]; BERT §4.1 [arXiv:1810.04805]; RoBERTa §4.1 [arXiv:1907.11692]. The memory bank is adapted from contrastive learning §3.3.
• Compared against: SGC [arXiv:1902.07153], TextGCN, BERT, RoBERTa and GAT variants tab. 1.
• Superseded by: not known from content in this wiki.