GreaseLM: Graph REASoning Enhanced Language Models for Question Answering

In plain words

Answering a question often needs two things: the wording of the question and general knowledge about how things relate. A pretrained language model is good with wording but does not reliably store how concepts connect. A knowledge graph stores those connections explicitly, as a network of linked concepts.

GreaseLM runs a language model and a graph network side by side over the same question. Several times while they process it, they exchange information. The exchange goes through two small shared "bridge" vectors: one stands for the text and one stands for the graph. The text can therefore be read in light of the linked facts, and the facts can be read in light of the wording.

The model was tested on three multiple-choice benchmarks, two of common sense and one of medical knowledge. On each it beat the same language model without the graph, and it also beat earlier systems that combine a language model with a knowledge graph.

TL;DR

GreaseLM is a question-answering model that runs a pretrained language model and a graph neural network over a subgraph of a knowledge graph retrieved for each question and candidate answer. Three to six fusion layers exchange information between the two streams through one interaction token and one interaction node. Author-reported gains over QA-GNN are 0.8 points on CommonsenseQA (IHtest), 0.5 on MedQA-USMLE and 2.0 on OpenbookQA. Check the Paper vs code list before reimplementing: the released code departs from the paper's equations in several places.

Problem

Pretrained language models (LMs) do not robustly represent latent relations between concepts, which reasoning needs. Knowledge graphs (KGs) can supply those relations, but it is open how to fuse the KG with the language context so that each can inform the other §1. Earlier LM+KG systems start the graph from a single vector taken from the LM's final layer. That compresses the context before the two modalities interact, which the authors argue limits the cross-relations that can be captured §6, Quantitative Analysis. Task: given a context paragraph c, a question q and candidate answers A, plus a KG, pick the correct answer §4.1.

Method

GreaseLM has a unimodal LM stack followed by M fusion layers. Each fusion layer adds one LM block, one graph-attention layer, and a modality interaction MLP that mixes only the interaction token and node §4.2 §4.3 fig. 1.

1. Input and KG retrieval §4.1 app. B.1

1. Concatenate c, q and the candidate a with separator tokens, and tokenize them into w_1 ... w_T §4.1.
2. Prepend an interaction token w_int to the sequence §4.1.
3. Retrieve a subgraph of the KG with the procedure of QA-GNN app. B.1: link entities in the text; add bridge entities that lie on a 2-hop path between two linked entities; score each node by its relevance to the QA context with a pretrained LM; keep the top 200 nodes; keep the edges between kept nodes app. B.1.
4. Give each node a type: 0 question entity, 1 answer entity, 2 other node, 3 the context node that stands for the interaction node [code: modeling/modeling_greaselm.py#L189-L191]. Types 0 and 1 come from the linking step §4.1.
5. Connect the interaction node e_int to the linked nodes only, not to bridge nodes §4.1. Connecting it to all nodes lowers the dev score from 78.5 to 77.6 app. C.1.

2. Unimodal LM layers, l = 1 .. N §4.2 eq. 1

The first N transformer blocks encode the text on their own. The LM is initialised from the pretrained checkpoint §4.2 §5.1.

3. Fusion layers, l = 1 .. M §4.3

Each layer runs the three steps below, in this order.

• LM block: one more transformer block over the text states eq. 2.
• GNN layer: one round of relation-aware attention over the KG nodes eq. 3 to eq. 10.
• Modality interaction: concatenate the interaction token and interaction node states, pass them through a two-layer MLP called MInt, and split the output back into the two parts eq. 11. All other tokens and nodes skip this step, but they read the bottleneck vectors in the next layer §4.3.

4. Answer score and training §4.3 §5.1

The score of candidate a is an MLP applied to three vectors: the interaction token after the last layer, the interaction node after the last layer, and an attention-pooled summary g of the KG node states, with the interaction token as the query. Train with cross-entropy over the candidates and predict the argmax §4.3.

The GNN layer in the paper's form:

# Relation-aware graph attention over the retrieved subgraph  [§4.3]
# node_states: (J, D_node); edges: (source s, target j, relation)
def gnn_layer(node_states, edges, node_types):
    # Keep the input states for the residual term  [eq. 4]
    previous_node_states = node_states
    # Relation embedding from the edge type and both node types  [eq. 5]
    relations = relation_mlp(edge_relation, node_types[s], node_types[j])
    # Message from neighbour s to node j  [eq. 6]
    messages = message_linear(node_states[s], node_types[s], relations)
    # Query comes from the source s, key from the target j  [eq. 7-8]
    queries = query_linear(node_states[s], node_types[s])
    keys = key_linear(node_states[j], node_types[j], relations)
    # Scaled dot-product score per edge  [eq. 9]
    scores = dot(queries, keys) / sqrt(D)
    # Softmax over the neighbours of j, self-loop included  [eq. 10]
    weights = softmax_over_neighbours_of(j, scores)
    # Weighted sum of messages, a 2-layer MLP, plus the residual  [eq. 4]
    aggregated = sum_over_edges_into(j, weights * messages)
    node_states = mlp_node(aggregated) + previous_node_states
    return node_states

One fusion layer and the full forward pass:

# One fusion layer: LM block, GNN layer, then modality interaction  [§4.3]
# token_states: (T + 1, d_lm); node_states: (J + 1, D_node)
# Index 0 of each is the interaction token or interaction node
def fusion_layer(token_states, node_states, edges, node_types):
    # Language branch: one more transformer block  [eq. 2]
    token_states = lm_block(token_states)
    # Graph branch: one round of message passing  [eq. 3-10]
    node_states = gnn_layer(node_states, edges, node_types)
    # Modality interaction: mix only the two bottleneck vectors  [eq. 11]
    joint = concat(token_states[0], node_states[0])
    joint = mint_mlp(joint)
    token_states[0], node_states[0] = split(joint, d_lm)
    return token_states, node_states

# Full forward pass for one (context, question, answer) input  [§4]
# Output: one score for this candidate answer
def greaselm_score(token_ids, node_ids, edges, node_types):
    # Prepend the interaction token to the text  [§4.1]
    token_states = embed_tokens(prepend_interaction(token_ids))
    # KG node vectors are pretrained embeddings  [§5.1] [app. B.2]
    node_states = kg_embedding(node_ids)
    # Interaction node: random in the paper, zero in the code  [§4.3]
    node_states = prepend_interaction_node(node_states)
    # Unimodal LM layers, l = 1..N  [§4.2] [eq. 1]
    for layer in lm_layers_before_fusion:
        token_states = layer(token_states)
    # Fusion layers, l = 1..M  [§4.3]
    for step in range(M):
        token_states, node_states = fusion_layer(
            token_states, node_states, edges, node_types)
    # Answer score from interaction token, node and pooled graph  [§4.3]
    graph_vec = attention_pool(node_states[1:], query=token_states[0])
    return mlp(concat(token_states[0], node_states[0], graph_vec))

The symbols D, d_lm, N and M are not given as numbers in the main text; the values per dataset are in the Hyperparameters table. The paper's sqrt(D) is not defined in the active text §4.3. The code uses the per-head dimension for this scale modeling_gnn.py.

Hyperparameters

Paper values are from the appendix table tab. 7 app. B.3. Where the released scripts differ, the script value is listed in the Notes column.

Datasets and graphs §5.1:

• CommonsenseQA: 12,102 questions, 5-way. In-house split of Lin et al. 2019; the official test is not reported §5.1.
• OpenbookQA: 5,957 questions, 4-way, official splits §5.1.
• MedQA-USMLE: 12,723 questions, 4-way, original splits §5.1.
• ConceptNet (CSQA, OBQA): 799,273 nodes and 2,487,810 edges §5.1.
• Medical KG (MedQA): UMLS disease database and DrugBank, 9,958 nodes and 44,561 edges §5.1.
• For about 5% of MedQA examples no entity is retrieved. A dummy node with a zero vector stands in for the graph there app. B.2.

Results

Author-reported. Hardware is not stated in the paper's main text. The LM+KG baselines use the same seed LM as GreaseLM §5.2 tab. 2 tab. 3.

CommonsenseQA, in-house split, accuracy in % tab. 2:

OpenbookQA test accuracy, same seed LM tab. 3: AristoRoBERTa (no KG) 78.4; RGCN 74.6; GconAttn 71.8; RN 75.4; MHGRN 80.6; QA-GNN 82.8; GreaseLM 84.8. No standard deviations are given for OpenbookQA.

OpenbookQA public leaderboard comparison tab. 4. The paper says GreaseLM has the third highest score, and that among models of comparable size it is the highest §6:

* The table footnote says UnifiedQA (11B) and T5 (3B) are 30x and 8x larger than GreaseLM.

MedQA-USMLE accuracy in % tab. 6 (split not stated in the table):

Question analysis on CommonsenseQA IHdev tab. 5 §6. The columns count prepositional phrases (a proxy for explicit constraints), negation terms and hedge terms:

On questions with no prepositional phrase, the paper says QA-GNN and GreaseLM perform comparably §6. The qualitative example is a node-attention trace from the interaction node; the paper reads its intermediate-layer rise on the "bug" node as resembling the hedge word "unlikely" §6, fig. 2.

Ablations on CommonsenseQA IHdev tab. 8 app. C.1:

LM initialisation tab. 9 tab. 10 app. C.2:

The paper's own summary of gains: on CommonsenseQA test, 5.5 points against the fine-tuned LMs and 0.9 against the LM+KG baselines; on OpenbookQA, 6.4 against the raw LM and 2.0 against QA-GNN §6. The IHtest table values give 0.8 against QA-GNN, not 0.9; see Pitfalls.

Paper vs code

Code: github.com/snap-stanford/GreaseLM at commit 6ccbc964152eeae6a040f3bbe39116c3baa75616, MIT license. The repository is under the snap-stanford organisation, and its LICENSE file names co-author Xikun Zhang. Paths below are relative to the repository root.

Discrepancies with the paper's equations or text:

• Interaction node initialisation. Paper: "initialized randomly" §4.3. Code: the position-0 vector is set to zero before the GNN [code: modeling/modeling_greaselm.py#L201-L203], so the interaction node starts at zero. Its type and score embeddings are added at each GNN layer [code: modeling/modeling_greaselm.py#L408].
• Node relevance scores. The paper uses LM relevance scores to keep the top 200 nodes app. B.1. The model's forward pass overwrites the node score input with zeros modeling_greaselm.py, so the score embedding is constant. The scores are used only in preprocessing. The paper's score-aware attention text is not in the active source.
• Attention normalisation. Paper eq. 10 normalises over the neighbours of the target node j. Code takes the softmax over edges grouped by the source index modeling_gnn.py and then multiplies each weight by the source node's out-degree modeling_gnn.py. Neither the grouping by source nor the degree factor appears in the paper.
• Residual connection. Paper eq. 4 adds the previous node state. The code's GNN layer returns the MLP output with no residual term [code: modeling/modeling_gnn.py#L91-L92]. The next step applies only activation and dropout modeling_greaselm.py. The output MLP contains a BatchNorm layer modeling_gnn.py, which the paper does not show.
• Attention heads. The table lists "2 attention heads in GNN" tab. 7. The GNN attention layer is built with 4 heads, its default [code: modeling/modeling_gnn.py#L36] modeling_greaselm.py. The value 2 is the argument of the final pooling over KG nodes [code: greaselm.py#L579].
• Query for the graph pooling. Paper: g is attention-pooled with h_int as the query §4.3. Code: the query is the LM's pooler applied to the first token of the last layer modeling_greaselm.py. In the Hugging Face implementation, that pooler is a dense layer with tanh.
• Node input features. Each GNN layer receives the node state concatenated with a learned node-type embedding and a score embedding [code: modeling/modeling_greaselm.py#L408] modeling_gnn.py. The paper uses type embeddings in eqs. 5 to 8 but does not describe this concatenation.
• OpenbookQA configuration. The table says M = 6 and MInt hidden units of 200 tab. 7. The OpenbookQA script uses k = 5 and ie_dim = 400 [code: run_greaselm.sh#L17] run_greaselm.sh. The paper does not say which setting produced 84.8.
• Entity embeddings. The pretrained KG embedding table is frozen by default; only the projection to 200 dimensions is trained [code: greaselm.py#L75-L78] layers.py. The paper does not say whether the table is frozen.
• Classifier head. Paper: "an MLP" over the three vectors, with no layer count given §4.3. Code default: fc_layer_num = 0, which is a single linear layer after dropout modeling_greaselm.py layers.py.

Confirmed as written (no discrepancy found): the single MInt MLP shared across fusion layers modeling_greaselm.py; the MInt as two linear layers with one hidden layer layers.py; the LM frozen for the first epochs greaselm.py; RAdam with a constant learning rate tab. 7; the 1e-3 learning rate for non-LM parameters in the scripts run_greaselm.sh.

The subgraph retrieval steps (linking, bridge nodes, top-200 pruning) live in preprocess_utils/. I did not check them line by line, so the description in app. B.1 is the paper's and the code is unverified for this step.

Pitfalls

• Cost grows with the number of candidates. Each candidate answer gets its own KG subgraph and its own pass through the LM [code: modeling/modeling_greaselm.py#L74-L82], with up to 200 nodes per graph. The paper gives no run-time or memory figures.
• Relation count is not in the paper. The code sets 38 relation types for ConceptNet and 34 for the medical KG greaselm.py. It adds 2 more when the interaction node connects to all nodes greaselm.py. Size the relation embedding from these numbers.
• Gradient accumulation. The loss of each mini-batch is divided by the full batch size, so the sum over mini-batches equals the batch mean [code: greaselm.py#L334-L343]. A mini-batch size of 8 or 2 (scripts) does not change the effective batch size of 128.
• Model selection is not stated in the paper. The code reports the test accuracy at the epoch with the best dev accuracy; a tie goes to the later epoch greaselm.py. Training stops after 10 epochs without a new best dev score, counted from the unfreeze epoch greaselm.py.
• Seeds and variance. The number of runs behind the ± values is not stated. The scripts use one seed (5) run_greaselm.sh.
• Compute is not stated in the paper's main text.
• Weight decay is not in the paper. The code default is 1e-2. It applies to all parameters except biases and LayerNorm weights [code: utils/parser_utils.py#L87] greaselm.py [code: greaselm.py#L247-L250].
• Text and table disagree on one number. The text gives +0.9 over existing LM+KG models on CommonsenseQA test. The table gives 74.2 for GreaseLM and 73.4 for QA-GNN, a gap of 0.8 tab. 2 §6. Use the table.
• OpenbookQA facts. The baseline AristoRoBERTa takes an extra corpus of scientific facts §5.2. The paper does not say whether GreaseLM receives that corpus, and the OpenbookQA converter in the code does not refer to it [code: preprocess_utils/convert_obqa.py].
• LM frozen for the first epochs. The LM parameters are frozen for 4 epochs on CommonsenseQA and OpenbookQA and for 0 epochs on MedQA tab. 7 [code: greaselm.py#L319-L325].
• Learning-rate default trap. The parser's default for OpenbookQA's non-LM learning rate is 3e-4 greaselm.py. The paper and the scripts use 1e-3, so pass 1e-3 explicitly.
• Sequence length. MedQA uses 512 tokens, the other two 100 tab. 7.
• Ablations use one split. All ablations are on CommonsenseQA IHdev only tab. 8.

Lineage

• Builds on: QA-GNN, which shares five of GreaseLM's seven authors. GreaseLM reuses QA-GNN's subgraph retrieval app. B.1. Node embeddings follow MHGRN app. B.2.
• Compared against: RoBERTa-Large, RN, RGCN, GconAttn, KagNet, MHGRN, QA-GNN §5.2. On the OpenbookQA leaderboard, T5 and UnifiedQA tab. 4.
• Superseded by: not stated in this card.