In layer l of JointLK, the Joint Reasoning Module receives the query token representations Ql−1={qil−1}i=1M∈RM×D and the GNN-updated knowledge graph entity representations X~l={x~jl}j=1∣V∣∈R∣V∣×D, where M is the number of query tokens, ∣V∣ is the current number of graph nodes, and D is the hidden feature dimension.
First, an affinity matrix Sl∈RM×∣V∣ is constructed between each token i and node j:
Sijl=WS⊤[qil−1;x~jl;qil−1∘x~jl]
where WS∈R3D is a learnable parameter vector, [;] denotes vector concatenation, and ∘ denotes element-wise multiplication.
Row-wise and column-wise softmax operations derive the KG-to-LM attention maps Sql∈RM×∣V∣ on query tokens and LM-to-KG attention maps Sxl∈RM×∣V∣ on entity nodes:
Sqil=softmax(Si,:l)
Sxjl=softmax((S:,jl)⊤)
The attended representations q^ij and x^ij are computed via matrix multiplication ⊗:
q^ij=qil−1⊗Sqil,x^ij=x~jl⊗Sxjl
The cross-attended features are fused with the original representations and projected back to dimension D using trainable matrices WQ,WX∈RD×4D:
qil=WQ[qil−1;x^ij;qil−1∘x^ij;qil−1∘q^ij]
xˉjl=WX[x~jl;q^ij;x~jl∘q^ij;x~jl∘x^ij]
The updated query representation Ql={qil}i=1M is fed to the next layer's joint reasoning module, and the updated graph representation Xˉl={xˉjl}j=1∣V∣ is sent to the dynamic pruning module.