Research map / Large language models / Retrieval and knowledge

K-ON: Stacking Knowledge on the Head Layer of Large Language Model

AAAI 2025 · Large language models › Retrieval and knowledge · approach: Knowledge-augmented question answering · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes K-ON, which adds K LoRA-adapted head layers to Llama-2-chat-7B so that entities are scored in one step via entity-level contrastive loss, plus head trajectory tuning to keep the original LLM head's token distributions.

Problem

Knowledge graph completion with LLMs suffers from a granularity mismatch: entities need multiple tokens, so token-level next-token prediction cannot score entities directly or support entity-level contrastive learning efficiently.

Approach

Attach K multi-step head layers (head MLPs, a small causal conditional Transformer, and per-step LoRA score layers) to the LLM so all candidate entities get joint probabilities from K parallel token predictions, trained with an entity-level contrastive loss and a KL-based head trajectory tuning objective that aligns with the original LLM head.

Key results

• On DB15K, K-ON achieves MRR 38.10 and Hits@1 30.13, versus 32.68 MRR for the best prior baseline MMRNS.
• On MKGW, K-ON achieves MRR 36.64 and Hits@10 48.26, outperforming MMRNS (35.03 MRR, 47.47 Hits@10).
• Removing the entity-level contrastive loss drops DB15K MRR from 38.10 to 14.09, the largest ablation drop.

Measured effects reported in the paper

Surprising findings

• Multiplicative joint probability aggregation performs much worse than addition (MRR 23.24 vs 38.10), which the authors attribute to vanishing gradients.
• Performance saturates for K>=8 while trainable parameters grow significantly with K.

Limitations

• Does not support arbitrarily large K; entity names longer than K tokens are truncated.
• Lacks support for multimodal (vision) input in the LLM.
• Multiplicative aggregation of token probabilities underperforms, and the approach requires more GPU resources than conventional KG methods.

Experimental setup

Datasets: DB15K MKGW

Benchmarks: DB15K MKGW

Metrics: MRR Hits@1 Hits@3 Hits@10

Baselines: TransE DistMult RotatE IKRL TransAE KG-Bert MMKRL OTKGE MMRNS KGLM QEB VISTA

Builds on: Medusa MultiToken Llama 2 LoRA KG-Llama

Closest papers by meaning

• MKGL: Mastery of a Three-Word Language (NeurIPS 2024 · similarity 0.91)
• KG-FIT: Knowledge Graph Fine-Tuning Upon Open-World Knowledge (NeurIPS 2024 · similarity 0.89)
• A Systematic Exploration of Knowledge Graph Alignment with Large Language Models in Retrieval Augmented Generation (AAAI 2025 · similarity 0.88)
• KGGen: Extracting Knowledge Graphs from Plain Text with Language Models (NeurIPS 2025 · similarity 0.88)
• Think-on-Graph: Deep and Responsible Reasoning of Large Language Model on Knowledge Graph (ICLR 2024 · similarity 0.88)
• Select and Augment: Enhanced Dense Retrieval Knowledge Graph Augmentation (Abstract Reprint) (AAAI 2024 · similarity 0.88)