A Context-Aware Citation Recommendation Model with BERT and Graph Convolutional Networks

In plain words

When you write a paper, you often need to attach a reference at a spot in a sentence where you know a citation belongs, but you cannot remember which paper to cite. This work tries to suggest that paper automatically. You mark the spot with a placeholder, and the system ranks the papers that could fill it.

The system looks at two things. The first is the words next to the spot, which a language model reads. The second is the map of which papers cite which other papers, which a graph model turns into a numeric description of each paper. The two descriptions are joined and fed to a classifier that scores every candidate paper.

The authors also build a benchmark dataset for this task, because earlier datasets lacked either the surrounding text or the paper metadata. Their model scores higher than the earlier model they compare against, and they report the numbers below.

TL;DR

• Task: given the text around a citation placeholder, rank the papers that could be cited there §1.
• Model: BERT encodes the context; a GCN, pretrained as a variational graph autoencoder on the citation graph, gives a per-paper vector; the two are concatenated and classified with a softmax over cited papers §3.1.
• Data: a new dataset, FullTextPeerRead, built from PeerRead; the AAN-derived dataset is not released §2, Table 1.
• Result (author-reported): higher MAP, MRR and Recall@K than CACR on AAN and FullTextPeerRead tab. 2.
• Caution: several numbers, split details and metric definitions disagree between text, tables and code (see Pitfalls).

Problem

Context-aware citation recommendation picks a paper for a placeholder from the sentences around it, optionally with paper metadata §1. The authors say the field lacked a benchmark with both context and metadata. AAN did not provide preprocessed sentences and metadata together, DBLP gave only bibliographic data, and CiteseerX gave context without metadata §1. The paper also notes that earlier systems were compared on different datasets §1.

Method

Overview. A context encoder (BERT) and a citation encoder (GCN) each produce a vector. The citation encoder is pretrained as a variational graph autoencoder on the citation graph §3.1, §3.2. The two vectors are concatenated and passed to a softmax output over cited papers, trained with cross-entropy §3.1.

Citation encoder. The GCN uses the identity matrix as node input, so it is featureless, and the normalised adjacency à is built from the diagonal degree matrix D §3.2.1 eq. 2. The two-layer GCN is written as a function of the adjacency and the identity input eq. 1. The paper uses VGAE for unsupervised training: the latent Z is inferred as a Gaussian per node from the GCN outputs §3.2.2 eq. 4, and the adjacency is reconstructed from inner products of latents through a sigmoid §3.2.2 eq. 5. The training objective is a variational lower bound: reconstruction minus a KL term eq. 3.

Context encoder. The text is the left and right context of a citation placeholder §2.2, §4.2.1. BERT's first-token output (768 dimensions) is the text representation §4.2.3. The GCN embedding of the citing (target) paper is appended to it §3.1. The paper describes a feedforward network before the softmax §3.1; see Paper vs code.

Training. Cross-entropy over all cited papers as classes §3.1. BERT parameters, learning rate, batch size and epochs are in Hyperparameters.

Data split and filter. Cited papers are kept when they have at least five citations (for the main table) §4.3.1. The authors split by year, with earlier papers for training and later for testing §4.2.1.

# Context-aware citation scoring  [§3] [§3.1]
# Input:  left, right: text around the placeholder (strings)
#         target_id: id of the citing paper
# Output: one score per cited paper (softmax over all cited papers)

def score_candidates(left, right, target_id):
    # Keep the last n characters before and the first n after the
    # placeholder; n is 50 or 100 in the paper  [§4.2.1]
    # [code: utils.py#L20-L25]
    left_part = left[-n:]
    right_part = right[:n]
    # BERT encodes the pair; the first-token output is the text vector
    # [§3.1] [code: modeling.py#L224-L229]
    text_vector = bert_first_token(left_part, right_part)   # (768,)
    # Citation-graph embedding of the citing paper  [§3.2]
    # [code: utils.py#L116-L123]
    graph_vector = graph_embedding[target_id]                # (768,)
    # Concatenate the two vectors  [§3.1]
    # [code: modeling.py#L251-L252]
    features = concatenate(text_vector, graph_vector)        # (1536,)
    # Dropout on the concatenation, keep probability 0.9 (training only)
    # [code: run_classifier.py#L541-L544]
    features = dropout(features, keep_prob=0.9)
    # Linear layer over all cited papers, then softmax  [§3.1]
    # [code: run_classifier.py#L546-L548]
    logits = linear(features)                                # (num_cited,)
    return softmax(logits)

# Graph pretraining: VGAE over the citation graph  [§3.2] [eq. 1]
# Input:  adjacency A (N x N, papers as nodes), no text features
# Output: one embedding per paper (N, hidden2)

def pretrain_graph(adjacency, hidden1, hidden2, epochs, lr):
    # Symmetric degree normalisation of A  [eq. 2]
    # [code: GCN/code/preprocessing.py#L25-L31]
    a_tilde = normalise(adjacency)
    # Node input is the identity matrix, so A~ X W0 reduces to A~ W0
    # [§3.2.1] [code: GCN/code/gcn_train.py#L52]
    for epoch in range(epochs):    # 200 epochs  [§4.2.3]
        # First GCN layer with ReLU, width hidden1  [eq. 1]
        # [code: GCN/code/model.py#L92-L98]
        hidden = relu(a_tilde @ W0)                          # (N, hidden1)
        # Two heads give the Gaussian mean and log-std  [§3.2.2] [eq. 4]
        # [code: GCN/code/model.py#L100-L112]
        mean = a_tilde @ hidden @ W_mean                     # (N, hidden2)
        log_std = a_tilde @ hidden @ W_std                   # (N, hidden2)
        # Sample latent z  [§3.2.2] [code: GCN/code/model.py#L114]
        z = mean + noise * exp(log_std)                      # (N, hidden2)
        # Reconstruct A from inner products  [§3.2.2] [eq. 5]
        # Training loss: reconstruction minus KL  [eq. 3]
        # [code: GCN/code/optimizer.py#L25-L32]
        loss = reconstruction_loss(z, adjacency) - kl(mean, log_std)
        adam_step(loss, lr)
    # Export: the code stores the sampled z for all papers
    # [code: GCN/code/gcn_train.py#L173]
    return z

# Evaluation on test contexts  [§4.2.2]
# Input:  scores (contexts, num_cited), true_sets (correct papers per context)
# Output: Recall@K, MRR, MAP

def evaluate(scores, true_sets, top_ks):
    # Recall@K: share of correct papers found in the top-K list, averaged
    # over contexts; K in {5, 10, 30, 50, 80} in code
    # [code: run_classifier.py#L893-L897] [code: utils.py#L167-L191]
    recall = {k: mean_recall_at_k(scores, true_sets, k) for k in top_ks}
    # MRR: reciprocal rank of the first correct paper, averaged
    # [code: utils.py#L237-L255]
    mrr = mean_reciprocal_rank(scores, true_sets)
    # MAP: average precision per context, averaged; code uses the full
    # ranking, not K = 30  [§4.2.2] [code: utils.py#L194-L201]
    map_value = mean_average_precision(scores, true_sets)
    return recall, mrr, map_value

Hyperparameters

Values are as the paper states them. Where the code differs, the code value is given in Notes and detailed in Paper vs code.

Results

Author-reported. Setup: AAN and FullTextPeerRead, BERT-base and a two-layer GCN, compared against CACR §4.1, §4.2. Table 2 uses citation frequency over five and 50 pair context sentences tab. 2. Table 3 uses frequency over one and 200 pair context sentences, and only AAN tab. 3.

Table 2 tab. 2. Frequency over five citations, 50 pair context sentences. Bold values are the authors' own marks.

The table has no Recall@100 column, though the text lists K = 100 §4.2.2.

Table 3 tab. 3. AAN only; frequency over one citation; 200 pair context sentences; the CACR row is labelled "CACR metric result" in the caption.

In Table 3 CACR is higher at Recall@30, @50 and @80 tab. 3, as the text notes for Recall@10 and beyond §4.3.1.

Table 5 tab. 5. Effect of citation frequency; the metric is not named in the table; see Pitfalls.

Qualitative (Table 4) tab. 4: top-3 recommendations for two example contexts, labelled correct or not. Not reproduced here.

Figures. fig. 2 architecture; fig. 3 VGAE architecture; fig. 4 BERT, GCN and BERT-left effects; fig. 5 performance against context length, which the text says is less affected once the context reaches about 100 §4.3.4. fig. 1 shows a manuscript with a placeholder and is an illustration only.

Paper vs code

Code: https://github.com/TeamLab/bert-gcn-for-paper-citation at commit 2770789826e72bc7fdb52319cfc12297ca07a5cd. The repository belongs to the TeamLab organisation, which the paper lists as the authors' lab (affiliation TEAMLAB, Gachon University). It has no LICENSE file and GitHub reports no license, so no code is quoted here, only cited. GitHub search found only a fork of this repository under another user and a master's-thesis fork; neither was used. We did not run the code.

• Embedding export is a sample, not the mean. The paper's inference step is described as the learned representation Z §3.2.2. The code exports model.z, which is the mean plus noise scaled by exp(log-std) gcn_train.py model.py. Whether the paper's reported numbers use the sample is not stated.
• Graph pretraining uses no text. The paper says each encoder is "pre-trained with context data" §3.1. The code pretrains the GCN only on the citation graph with identity features gcn_train.py.
• GCN width. The paper sets the first hidden width to the number of documents §4.2.3. The code default is 7071 config.py, and the command files use 7071 for AAN and 4837 for PeerRead pretrain.txt pretrain.txt. The paper gives 6,500 (text) and 7,073 (Table 1) for AAN, and 4,898 for FullTextPeerRead §4.2.1 tab. 1.
• Graph loss. Eq. 3 is the plain variational bound eq. 3. The code uses a positive-weighted cross-entropy with a normalising constant, and scales the KL term by 0.5 / N optimizer.py gcn_train.py.
• GCN dropout 0.5 during training is in the code only gcn_train.py.
• Classifier head. The paper describes a feedforward network before the softmax §3.1. The code applies dropout to the 1536-dimensional concatenation and a single linear layer, with no hidden layer run_classifier.py modeling.py.
• Graph embedding table trainability. For AAN the table is frozen; for FullTextPeerRead it is trainable modeling.py. The paper does not say.
• Maximum sequence length. The paper states 128 §4.2.3; the code default is 128 run_classifier.py, but the command files pass 50 for every run main_classifier.txt. The same flag also sets the character slice length run_classifier.py utils.py.
• Split year. The code's default is 2013 run_classifier.py. The command files do not pass a year, so FullTextPeerRead is also split at 2013 unless it is overridden. The paper says 2017 for FullTextPeerRead test §4.2.1.
• MAP definition. The paper says MAP is measured at K = 30 §4.2.2. The code computes average precision over the full ranking of all classes utils.py run_classifier.py.
• Recall K list. The text lists K = 100 §4.2.2; the code lists 5, 10, 30, 50, 80 run_classifier.py.
• Optimizer. Weight decay, betas and epsilon match the paper §4.2.3 optimization.py. Warmup, linear decay and gradient clipping are in the code only optimization.py optimization.py optimization.py.
• GCN learning-rate flag. --gcn_lr is declared type=int, so a command-line value such as 0.01 would not parse; the default 0.01 is used as a float config.py.
• Left-only variants. The code has no option to drop the right context run_classifier.py. The paper does not say how the Left rows were produced.
• Data. The README links a PeerRead CSV; the AAN-derived data is not released README.md and Table 1 says the same tab. 1.

Pitfalls

• Headline ratio. The abstract claims "more than 28%" improvement in MAP and Recall@K §abstract. The tables give about 2.1x CACR's MAP on AAN in Table 2 (0.6189 / 0.2893) and about 1.5x in Table 3 (0.4516 / 0.301). The paper does not show how 28% is computed. Use the table values.
• "Three times" vs tables. The text says "approximately three times" §4.3.1; Table 2 gives about 2.1x MAP on AAN and about 2.7x on FullTextPeerRead.
• Context length unit. The paper gives 50 and 100 without a unit §4.2.1. The code cuts characters, not words utils.py. Token count after BERT tokenisation is capped at the max sequence length, so the two numbers can interact.
• One flag, two meanings. --max_seq_length sets both the BERT token cap and the context character slice run_classifier.py utils.py. The command files use 50, so the context is at most 50 characters each side.
• Label space. The label encoder and the frequency filter are fit on all years before the split utils.py utils.py utils.py. Cited papers whose citations all fall in the test years therefore get a class with no training examples (inferred from the code order).
• Identical contexts are merged. Test rows with the same left and right strings form one query whose relevant set is the union of their citations utils.py. The paper does not describe this.
• Probability threshold. Recall counts only candidates with probability above 1e-10 run_classifier.py utils.py.
• Table 5 metric. The table names no metric. At frequency 5 its values equal the Recall@5 values in Table 2 (0.6736 and 0.6593). This is an inference from the numbers, not something the paper states.
• CACR numbers. The text says CACR was reproduced from its released code §4.3.1, but also that the comparison uses CACR's reported performance, which assumes frequency one §4.3.1. Table 3's caption says "CACR metric result". Which numbers are reproduced and which are copied is unclear.
• Left-context gap. The text says the left-only model is about 0.03 lower §4.3.3. Table 2 gives 0.022 (AAN MAP) and 0.030 (FullTextPeerRead MAP).
• Split counts. The AAN split numbers in the text §4.2.1 do not match the totals in Table 1. The AAN test-set sentence gives the same count of 973 twice for different test sets, so it cannot be used as written. Use the code's split rule.
• Data rebuild is manual. Extraction from arXiv Vanity was followed by manual removal of noise, and remaining placeholder text is a leak §2.2. Rebuilding the dataset needs that manual step.
• Costs the abstract hides. Each test query scores all cited papers, so inference cost grows with the number of classes. The label space includes every cited paper that passes the frequency filter utils.py.

Lineage

• Builds on: BERT (the paper cites it as its pretrained text encoder §1), Kipf and Welling's graph convolutional network §1, §3.2.1, and VGAE §3.2.2.
• Compared against: CACR, the earlier context-aware model that uses an LSTM encoder and paper metadata §3.1, §4.1.
• Superseded by: not stated in the content available here.