QA-GNN: Reasoning with Language Models and Knowledge Graphs for Question Answering

In plain words

Language models hold a lot of knowledge from their training text, but they are often weak at precise logic, such as noticing that a question says "not". A knowledge graph stores facts as linked entities, which suits that kind of reasoning, but the graph is large and much of it is off-topic for any one question.

QA-GNN combines the two for multiple-choice questions. For each candidate answer it takes the entities mentioned in the question and the answer, pulls in the entities that link them, and asks a language model how relevant each of those entities is to the question. It then adds the question and answer as one extra node in the graph. A graph neural network passes messages along the edges, so the question node and the entities update each other until the model has a score for that answer.

The authors report better accuracy than the comparable models that keep the language model and the graph separate, and they show that the attention patterns can be read as a small reasoning graph.

TL;DR

• Joins a pretrained LM (RoBERTa-large, or SapBERT for medicine) with a retrieved knowledge-graph subgraph §3, §4.4.
• Adds one QA-context node to the subgraph, linked to the topic entities, then scores every answer choice with a GNN §3.1.
• Relevance scores from the LM weight each KG node; the GNN uses them in attention and messages §3.2, §3.3.
• Author-reported: CommonsenseQA in-house test 73.41 (RoBERTa-large 68.69, MHGRN 71.11) tab. 3.
• The released code differs from the paper in several places that change the numbers. See Paper vs code.

Problem

Answering a question from knowledge needs two sources: facts an LM absorbed from text, and an explicit graph such as ConceptNet or Freebase §1. LMs cover a lot but handle structured reasoning, such as negation, poorly. Graphs suit structured reasoning but may be incomplete and noisy §1. Earlier LM+KG methods take the topic entities and their few-hop neighbours, which pulls in many nodes irrelevant to the question, and they run a language model and a GNN on separate inputs without mutual updates §1, §3.2. The paper states that with ConceptNet, 3-hop neighbourhoods give more than 400 nodes on average §3.2.

Inputs are a question q and one answer choice a, concatenated as the QA context [q; a] §3. The task is multiple-choice: the model picks one choice per question from a set §3.4.

Method

QA-GNN works in four stages §3:

1. Encode the QA context with the LM to get z_LM, the vector at the [CLS] token §2.
2. Retrieve a KG subgraph of the topic entities and their 2-hop neighbours, then keep the top 200 nodes by relevance §4.2.
3. Score each KG node's relevance to the QA context with the LM §3.2 eq. 1.
4. Build the working graph: add the QA-context node z, connect it to question-entity nodes (relation r_{z,q}) and answer-entity nodes (relation r_{z,a}) §3.1. Run L = 5 GNN layers of attention-weighted message passing §3.3 eq. 2–eq. 8. Predict with an MLP over z_LM, the final state of z, and a pooled graph vector §3.4.

Node types are Z (context), Q (question entity), A (answer entity) and O (other) §3.1. The context node is initialised with z_LM and KG nodes with entity embeddings §3.1, §4.2.

Relevance score §3.2 eq. 1:

\rho_v = f_{\text{head}}\big(f_{\text{enc}}([\text{text}(z); \text{text}(v)])\big) eq. 1

The attention weight between a source s and a target t §3.3 eq. 8:

\alpha_{st} = \frac{\exp(\gamma_{st})}{\sum_{t' \in \mathcal{N}_s \cup \{s\}} \exp(\gamma_{st'})}, \quad \gamma_{st} = \frac{q_s^\top k_t}{\sqrt{D}} eq. 8

The pseudocode below follows the paper's equations. Where the released code differs, the difference is listed under Paper vs code.

# QA-GNN: logit for one (question, answer choice) pair  [§3]
# Input:  qa_text (str); kg_names, kg_types (N,); edges (E, 3) src/dst/rel
# Output: scalar logit

def qa_gnn_logit(qa_text, kg_names, kg_types, edges):
    # Encode the QA context with the LM: z_LM  [§2] [§3.1]
    z_lm = encoder(qa_text)                              # (D_lm,)

    # Relevance of each KG node to the QA context  [§3.2] [eq. 1]
    rho = relevance_scores(qa_text, kg_names)            # (N,)

    # Node 0 is the context node (type 3, score 0), then the KG nodes  [§3.1]
    types = concat([3], kg_types)                        # (N+1,)
    rho = concat([0.0], rho)                             # (N+1,)
    h = initial_states(z_lm, kg_names)                   # (N+1, D)

    # L = 5 GNN layers, width D = 200  [§4.3]
    for layer in range(L):
        h = gnn_layer(h, rho, types, edges)              # [§3.3]

    # Pool the KG nodes; read out the context node  [§3.4]
    g = pool(h[1:])                                      # not specified
    return mlp(concat([z_lm, h[0], g]))                  # (1,)

One layer of the working-graph GNN, following the paper's equations:

# One message-passing layer on the working graph  [§3.3] [eq. 2]
# h: (N, D) node states; rho: (N,) relevance; types: (N,) in {0,1,2,3}
def gnn_layer(h, rho, types, edges):
    src, dst, rel = edges
    # Type and score embeddings, each of width D/2  [§3.3] [eq. 3] [eq. 5]
    u = type_linear(one_hot(types))                      # (N, D/2)
    s = score_mlp(rho)                                   # (N, D/2)
    # Relation embedding of each edge src -> dst  [eq. 3]
    r = relation_mlp(concat([one_hot(rel),
                             one_hot(types[src]),
                             one_hot(types[dst])]))      # (E, D)
    # Message from src to dst: src state, src type, relation  [eq. 4]
    msg = message_linear(concat([h[src], u[src], r]))    # (E, D)
    # Attention query from src, key from dst  [eq. 6] [eq. 7]
    query = query_linear(concat([h[src], u[src], s[src]]))  # (E, D)
    key = key_linear(concat([h[dst], u[dst], s[dst], r]))   # (E, D)
    # Scaled dot product, softmax over each source's neighbours  [eq. 8]
    gamma = sum(query * key, axis=1) / sqrt(D)           # (E,)
    alpha = softmax_grouped_by(src, gamma)               # (E,)
    # Attention-weighted sum of messages into each target  [eq. 2]
    total = scatter_sum(alpha[:, None] * msg, dst, N)    # (N, D)
    # 2-layer MLP with batch norm, then the paper's residual  [eq. 2]
    return mlp_batchnorm(total) + h                      # (N, D)

Relevance scores and the normalisation the code applies:

# Relevance of each KG node to the QA context  [§3.2] [eq. 1]
# Paper: LM probability of the node text given the QA context.
# Code: summed token loss of the unmasked sentence (see Paper vs code).
def relevance_scores(qa_text, node_names):
    # QA context alone: the reference  [code: utils/graph.py#L281-L313]
    losses = [sentence_loss(qa_text)]
    for name in node_names:
        # Sentence = QA context followed by the node name  [§3.2]
        losses.append(sentence_loss(qa_text + " " + name + "."))
    losses = array(losses)                               # (N+1,)
    # Centre on the context node  [code: modeling/modeling_qagnn.py#L159-L167]
    centered = losses - losses[0]                        # context -> 0
    # Divide by the mean absolute value over real nodes  [code: same lines]
    mean_abs = sum(abs(centered[1:])) / (len(node_names) + 1)
    return centered[1:] / (mean_abs + 1e-05)             # (N,)

Hyperparameters

Values marked "code" come from the released scripts or defaults, not the paper.

Results

Author-reported. Setup: ConceptNet for CommonsenseQA and OpenBookQA, a self-built UMLS and DrugBank graph for MedQA-USMLE; RoBERTa-large for the commonsense tasks, SapBERT for MedQA; baselines use the same LM §4.1–§4.4. The official CommonsenseQA test is hidden, so the in-house split of prior work is used for controlled comparison §4.1.

Other author-reported numbers: on CommonsenseQA the gain over RoBERTa is +4.7 and over MHGRN +2.3 §4.5. In a pilot on logical queries over FB15k, QA-GNN reaches Hit@3 40 against 15 for an LM alone §3.6.

Paper vs code

Code at commit 53eb88998e3365b19239a77ba4038042ca0a449a, MIT license.

• Relevance score. The paper defines ρ_v as the LM probability of the node text given the QA context §3.2 eq. 1. The code's preprocessing takes the sentence "QA context node name." (lowercased), sums the masked-LM token losses over every token, with no tokens masked, and stores the negated sum graph.py. The model then uses the difference from the context node's loss, divided by the mean absolute difference over the graph modeling_qagnn.py. Scores are computed once, in preprocessing, by a frozen RoBERTa-large in eval mode graph.py, not by the model being trained.
• Score embedding. The paper uses an MLP R → R^{D/2} §3.3. The code uses a fixed sinusoidal basis sin(1.1^j · ρ) for j = 0 … D/2−1, followed by a linear layer and GELU modeling_qagnn.py.
• Message input. The paper's f_m takes h_s, u_s and r_st, width 2.5D §3.3 eq. 4. The code's message layer also takes the source's score embedding, width 3D modeling_qagnn.py.
• Residual. Eq. 2 adds h^(l) to each layer's output §3.3. The code's GAT layer returns only the MLP output, with no residual modeling_qagnn.py. The input enters only at the end, through a residual on the model's output modeling_qagnn.py. Each layer also applies GELU and dropout modeling_qagnn.py.
• Attention scaling and heads. The paper uses one head, scaled by √D §3.3 eq. 8. The code uses 4 heads (head_count=4) modeling_qagnn.py, each scaled by √(D/4) modeling_qagnn.py.
• Out-degree rescaling. The code multiplies each attention weight by the number of edges leaving its source node. The paper does not mention this modeling_qagnn.py.
• Self-loops. The paper's sum includes s = t eq. 2 but does not say which relation embedding a self-loop gets. The code adds a dedicated self-loop edge type and uses its own embedding modeling_qagnn.py.
• Context node score. The paper does not define ρ for the QA-context node. The code sets it to zero after centring on the context node modeling_qagnn.py.
• Graph pooling. The paper says "pooling" without specifying it §3.4. The code uses multi-head attention pooling with the LM vector as query over KG nodes, excluding the context node modeling_qagnn.py.
• Output head. The paper describes an MLP over three vectors §3.4. With fc_layer_num=0, the default, the code uses one linear layer with no LayerNorm qagnn.py.
• Extra nodes. The paper says the subgraph holds all nodes on k-hop paths between topic entities §2. The code adds nodes that neighbour at least two topic entities, which is a 2-hop bridge rule graph.py.
• Truncation order. The code places topic entities first, then extra nodes sorted by relevance, and keeps the first 199 KG nodes plus the context node graph.py. Topic entities are kept unless there are more than 199.
• Graph pruning. The code uses a pruned ConceptNet that drops hascontext edges and filtered nodes conceptnet.py. The paper's counts in §4.2 are not stated for the pruned graph.
• Entity embeddings. The paper does not say whether they are trained. The code freezes them by default qagnn.py.
• Encoder freezing. The paper does not mention a freeze schedule. The code freezes the encoder for the first four epochs by default qagnn.py.
• Question text. The QA context in preprocessing is the stem plus the answer text, with a trailing period graph.py. This matches [q; a] §3.

No discrepancies were checked for the MedQA biomedical preprocessing, which lives in a notebook under utils_biomed/. Those numbers are reported from the paper only.

Pitfalls

• Relevance score is not what the paper describes. An implementer following eq. 1 literally will not reproduce the code's scores. The code's scores are negated unmasked sentence losses over the whole "QA context + node" string. Decide which one you are reproducing before comparing numbers §3.2 graph.py.
• Scores are fixed. They come from a frozen RoBERTa computed once during preprocessing. They do not change during training, so the LM is not learning relevance graph.py.
• Normalisation matters. Centring on the context node and dividing by the mean absolute value (with 1e-5 added) sets the scale of every score feature. Without it the feature range depends on the question modeling_qagnn.py.
• Out-degree rescaling changes the aggregation. In the code each source's attention weights sum to its out-degree rather than 1, so messages from high-degree nodes are amplified modeling_qagnn.py.
• Unstated details that the code fixes: the encoder's first four frozen epochs, frozen entity embeddings, the single linear output layer, the 2 pooling heads, and the 38 or 34 relation types qagnn.py.
• OpenBookQA training is unstable. Best dev accuracy varies with the seed and with library versions. The authors suggest trying several seeds or starting from a released checkpoint README.md. The scripts use seed 0 only.
• Ablation and subgroup tables have no standard deviations. Only the main tables carry ± values. The negation subset size is not stated tab. 10, tab. 11.
• Ablation numbers are in the text, not the table. The ablation table (tab. 8) is an image in the source. Use the numbers quoted in §4.6.1.
• Table and figure numbers are inferred from LaTeX order. The source has no compiled PDF in this repository. The numbers used here assume source order: Figure 1 task, Figure 2 overview, Figure 3 relevance scoring, Figure 4 interpretability, Figure 5 structured reasoning; Table 1 complexity, Table 2 logical query, Tables 3–7 main results, Table 8 ablation, Table 9 structured-reasoning examples, Tables 10–11 subgroup results. Check against the arXiv PDF before publishing.
• Costs. The model is about 360M parameters, and the paper reports about 20 hours on two RTX 2080 Ti for CSQA §4.3, §4.5. Per-layer time is O(|V|²L) on a dense graph and O(|V|LΔ) on a sparse one, where Δ is the maximum degree tab. 1.
• Leaderboard numbers are on the hidden test. The 76.1 (CSQA) and 82.8 (OBQA) rows come from the official test, which the authors score by submission. Only the in-house rows can be reproduced from the released splits §4.1 tab. 4, tab. 6.

Lineage

• Builds on: GAT for the attention framework §3.3; MHGRN for the KG preprocessing, entity embeddings and 2-hop retrieval §4.2; KagNet for the subgraph retrieval and for the LM+KG framing §1, §4.4.
• Compared against: RN, RGCN, GconAttn, KagNet and MHGRN under the same LM §4.4; T5 and UnifiedQA on the official leaderboards §4.5.
• Superseded by: not stated in this wiki.