SPECTER: Document-level Representation Learning using Citation-informed Transformers

In plain words

Search and recommendation for research papers need a numeric "fingerprint" for each paper, so that papers on similar topics get similar fingerprints. A plain language model trained to read text is not good at this, because it learns to predict words, not which papers belong together.

SPECTER gets its training signal from citations. Papers that cite each other are treated as related, and papers that are not are treated as unrelated. The model is pushed so that a paper's fingerprint lands closer to a paper it cites than to a paper it does not. Papers that are cited by a cited paper are used as the harder "unrelated" examples, because they are topically close.

After this training, the model needs only the title and abstract of a paper to produce its fingerprint. No citation data is needed at use time, so it also works for new papers that nobody has cited yet. The fingerprints can be compared with a simple distance and used as features in other tasks, without any further training.

TL;DR

• SPECTER starts from SciBERT and continues training it on a triplet margin loss over citation triples §2.3 eq. 2.
• Input is the title and abstract only; the paper embedding is one vector per paper, used without fine-tuning §2.2 §2.5.
• Author-reported: average 80.0 over 12 SciDocs metrics, against 76.9 for SGC, the next-best baseline tab. 1 §5.
• Relevant if you need fixed document embeddings for scientific papers and can pretrain or use an existing checkpoint; the citation-based objective is the part that matters.

Problem

Off-the-shelf pretrained language models (including SciBERT) produce poor document-level representations for document tasks, because their training objective predicts words from nearby text and never sees relatedness between documents §1. Graph-based methods do use relatedness but cannot embed unseen papers §6 §7. The paper assumes that the title and abstract are a sufficient summary of a paper §2.2 and that citations are a usable proxy for topical relatedness §2.3. It also introduces SciDocs, a benchmark of seven tasks to evaluate such embeddings §3.

Method

Overview. SPECTER is a Transformer encoder (SciBERT initialisation) whose output for one paper is a single dense vector fig. 1. Training uses triples of papers (query, positive, negative). The positive is a paper the query cites; the negative is a paper the query does not cite §2.3. Training minimises a triplet margin loss with L2 distance and margin 1 eq. 2 §2.3.

Step 1: input encoding §2.2. Concatenate the title and the abstract, separated by a [SEP] token, with a [CLS] token in front, then tokenise with WordPieces §2.2. The paper says a single concatenated input is used, since encoding title and abstract separately did worse §2.2, footnote.

Step 2: paper vector §2.2 eq. 1. The paper defines the vector as the final hidden state of the [CLS] token eq. 1. The released AllenNLP code instead averages token representations; see Paper vs code.

Step 3: triplet loss §2.3 eq. 2. For a triple, compute the L2 distance from the query vector to the positive vector and to the negative vector. The loss is zero if the negative is farther than the positive by at least the margin m = 1. The paper reports that cosine distance underperformed L2 §2.3, footnote.

Step 4: negative selection §2.4 §4. Two kinds of negatives are used:

• Hard negatives: papers that the query does not cite but that a cited paper of the query does cite (a citation of a citation) §2.4.
• Easy negatives: papers drawn at random from the corpus §2.4.

For each query the paper uses up to 5 triples: 2 hard and 3 easy negatives, with positives sampled from the query's direct citations §4. The paper reports that removing hard negatives lowers the average from 80.0 to 78.4 tab. 2 §6.

Step 5: training §4. Initialise from SciBERT, train all parameters on the triplet objective, using Adam with the BERT-style settings described in §4. Train for 2 epochs with an effective batch size of 32 (micro-batch 4 with gradient accumulation) on a single 12 GB GPU.

Step 6: inference §2.5. Only the title and abstract are needed. The output vector is the paper embedding, and the paper compares embeddings with L2 distance (or cosine for the recommendation feature, which uses L2-normalised embeddings §3.4).

# SPECTER training step  [§2.3] [eq. 2]
# Input:  query, positive, negative: title and abstract text
# Output: scalar loss

def embed(paper):
    # Concatenate title and abstract with [SEP] and encode  [§2.2]
    tokens = wordpieces(paper.title + " [SEP] " + paper.abstract)
    # Read one vector per paper from the encoder output  [§2.2] [eq. 1]
    states = scibert(tokens)  # (N, 768)
    return pool(states)  # (768,); pooling is discussed under Pitfalls

def training_step(query, positive, negative):
    # Embed the query, the positive and the negative paper  [§2.2]
    query_vector = embed(query)
    positive_vector = embed(positive)
    negative_vector = embed(negative)
    # L2 distances to the positive and to the negative  [§2.3]
    positive_distance = l2(query_vector, positive_vector)
    negative_distance = l2(query_vector, negative_vector)
    # Triplet margin loss with margin m = 1  [§2.3] [eq. 2]
    return max(positive_distance - negative_distance + 1.0, 0.0)

# Build the training triples for one query paper  [§2.4] [§4]
# Input:  cites (papers the query cites), cites_of_cites (candidate hard
#         negatives), corpus (all papers)
# Output: up to 5 (query, positive, negative) triples

def triples_for(query):
    # Positives are direct citations of the query  [§4]
    positives = sample(cites[query])
    # Two hard negatives: cited by a cited paper, not cited by the query  [§4]
    hard = sample(cites_of_cites[query] - cites[query], 2)
    # Three easy negatives: random papers  [§4]
    easy = sample(corpus - cites[query], 3)
    return [(query, positive, negative) for negative in hard + easy]

The second block is a simplification: it draws the positive once per query and does not show the released code's count-based selection (see Paper vs code).

Hyperparameters

Baseline settings (for reproducing comparisons, app. A):

• Doc2Vec: 300-dimensional vectors, 27 configurations from window {5, 10, 15}, sample {0, 1e-6, 1e-5}, epochs {50, 100, 200}, min_count 3, negative sampling 5 app. A.
• fastText: 300-dimensional, trained on about 3.1B tokens of scientific text, minCount 25 app. A.
• SIF: first principal component removed; the weighting parameter a chosen from [1e-5, 1e-3] on a log grid app. A.
• Citeomatic: 75-dimensional app. A. SGC: k from 1 to 8 with SIF features app. A.
• SciBERT baseline: pooled [CLS] of the last layer app. A.

Results

Author-reported. Setup: SciDocs with seven tasks and 12 metrics, embeddings used without fine-tuning, baselines extensively tuned by the authors §3 §4 tab. 1. The table below reproduces the paper's Table 1 in full. "Avg." is the mean of the 12 metrics.

Source: tab. 1. ^ Propensity-adjusted (the paper's hat notation) §3.4. MAP and nDCG are ranking metrics over L2 distances of raw embeddings, except the recommendation task, which uses the embedding cosine as one feature in a ranker §3 §4.

Headline numbers §5:

• Average 80.0 against 76.9 (SGC) and 3.1 points above the next-best baseline §5 tab. 1.
• MeSH F1 86.4 and MAG F1 82.0 (about +2.3 and +1.5 over the best baseline on each) §5.
• Co-view MAP reported as 83.8 in the text §5 (table shows 83.6, see Pitfalls).
• Direct citation nDCG 94.9, slightly above Citeomatic (94.1) and below SGC (96.2) tab. 1.
• Online A/B test on 4,113 clicks: the SPECTER-based recommender raised click-through by 46.5 percent over a textual-similarity production system §5.

Ablations §6 tab. 2. Averages over task groups: CLS (classification), USR (user activity), CITE (citation), REC (recommendation).

Comparison with task-specific fine-tuning §6 tab. 3. Fine-tuning SciBERT on each user-activity signal, or on all of them jointly, scored below the fixed SPECTER embedding (average 80.0 against 76.0 to 78.0). Columns are the same as tab. 2:

Source: tab. 3.

Visualisation fig. 2. A t-SNE projection of SPECTER and SciBERT embeddings, labelled by MAG topic, shows more compact topic clusters for SPECTER. Clustering homogeneity and completeness are 0.41 and 0.72 for SPECTER, against 0.19 and 0.63 for SciBERT §6. The figure itself is not reproduced here (see Pitfalls on images).

Paper vs code

Official repository: https://github.com/allenai/specter (linked by the paper in a footnote to the abstract and in §4), commit f7afbc9cff79d487f6529a4b5802f10ed666df94, Apache-2.0 LICENSE. The paper and code differ in the following places.

• Pooling. Paper uses the [CLS] final hidden state §2.2 eq. 1. The AllenNLP model passes the BERT token embeddings to a title encoder of type boe (bag of embeddings) simple.jsonnet model.py. The README's errata says the original model mixes all BERT layers and uses average pooling, not [CLS] README.md.
• Pooling in the Hugging Face release. The Lightning training code uses BERT's pooler output, model(...)[1], to compute the training loss train.py train.py. The Hugging Face embedding script takes the last-layer [CLS] vector embed_papers_hf.py. These are different vectors, so training and inference pooling do not match in that code.
• Learning rate. Paper: 2e-5 §4. Config: 3e-5 simple.jsonnet. The Lightning script defaults to 2e-5 train.py.
• Scheduler step count. Paper: "number of train steps equal to training instances" §4. Config: steps per epoch set to instances divided by 32 simple.jsonnet.
• Hard-negative ratio. Paper: 2 hard and 3 easy per query §4. Code default: ratio 0.3 and 5 samples, giving ceil(0.3 x 5) = 2 hard create_training_files.py triplet_sampling_parallel.py. Same counts here, by a different route.
• Hard-negative definition. The code works on weighted relation counts from data.json (direct citation 5, citation of citation 1 in the README). A hard negative is a candidate whose count is below the best candidate's count by at least a margin of 0.5 times the count spread triplet_sampling_parallel.py triplet_sampling_parallel.py create_training_files.py. The paper describes only the citation-of-citation rule §2.4.
• Title normalisation. The code appends a . to the title if it does not end with one, then adds [SEP] before the abstract create_training_files.py create_training_files.py. The paper mentions [SEP] only §2.2.
• Unused layers. The model defines a feed-forward layer and a layer norm on the title output, but forward does not apply them model.py model.py. The embedding-layer norm is applied to venue only model.py model.py.
• Dropout. The config sets dropout 0.25 on the pooled title vector model.py simple.jsonnet. The paper does not mention dropout §4.
• Gradient clipping and weight decay. Config: clipping 1 simple.jsonnet and weight decay 0.01 simple.jsonnet. Neither is in the paper §4.
• Casing. The AllenNLP config lowercases text simple.jsonnet. The Lightning training code loads the cased SciBERT tokenizer train.py.
• Results. The authors' Hugging Face model reports average 80.5 with different per-task values from Table 1; for example, direct citation MAP 92.0 (paper 88.3) README.md. The README says the Hugging Face model is slightly higher on average than the paper, and that the original implementation is needed to reproduce the exact numbers README.md. Embeddings from the public API come from a different SPECTER version and cannot be mixed with the repository's embeddings README.md.

Pitfalls

• Pooling is the main ambiguity. Paper says [CLS], the AllenNLP code averages tokens, and the Hugging Face training code uses the pooler output. Use the pooling that matches the checkpoint you load, and do not mix checkpoints. If you retrain, pick one pooling and keep it for training and inference.
• Max sequence length is not stated in the paper. The data pipeline truncates title and abstract to 512 tokens create_training_files.py; the shell script defaults to 256 run-exp-simple.sh; the Hugging Face inference script pads to 512 embed_papers_hf.py. Long abstracts are truncated at whichever limit you use, so the embedding of the same paper changes with it.
• Scheduler total steps are ambiguous. "Number of train steps equal to training instances" (about 684K) and the config's instances/32 per epoch differ by a factor of 32 in total steps §4 simple.jsonnet. Implement the config's reading and mark it.
• Learning rate differs by 50 percent between paper and config. 2e-5 §4 vs 3e-5 simple.jsonnet.
• Co-view number in the paper is inconsistent. Text gives 83.8 MAP §5; Table 1 shows 83.6 tab. 1. The text's +2.7 gain over Citeomatic uses 83.8.
• Citation links are needed for training only. Training needs the citation graph for the training papers. Inference needs only the title and abstract of the input paper §2.5.
• Authors and venue hurt or do not help. Adding authors lowers the average to 67.3 and venue leaves it at 79.9 tab. 2 §6. The paper says venue normalisation was not done, so venue results may be affected by noise §6, footnote.
• Expected cost. One epoch took about 1 to 2 days on a single 12 GB GPU §4. Two epochs are required.
• Results are author-reported. Baselines were tuned on a validation set, but the authors say their own method was only "minimally tuned" §4. The SciDocs data and metrics come from the same authors.
• Figures. Figures 1 and 2 are not reproduced in this card; the images are not allowed on the page yet. Figure 1 shows the training overview fig. 1. Figure 2 shows the t-SNE embeddings fig. 2.
• Cosine vs L2. The paper finds L2 better than normalised cosine for the triplet loss §2.3. The recommendation feature uses cosine on L2-normalised vectors, which is equivalent to L2 there §3.4 (footnote).

Lineage

• Builds on: SciBERT, the scientific-text BERT that is the initialisation §2.2 §4. Not compiled in this wiki.
• Builds on: BERT, the base architecture §2.2. Not compiled in this wiki.
• Compared against: Citeomatic (citation-trained paper embeddings, 75-dimensional) app. A; SGC (graph convolution over citations; cannot embed unseen papers) §5 app. A; Sentence-BERT (triplet-trained sentence embeddings) app. A; Doc2Vec, fastText-sum, SIF, ELMo tab. 1.
• Introduces: SciDocs, seven document-level tasks with release of the data §3. Its own GitHub repository is linked from the paper §1 §3.
• Superseded by: not known from content in this wiki.