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

arXiv: https://arxiv.org/abs/2502.06257 · HTML: https://mlsanity.com/paper/2502.06257/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 40 references that are in the map.
- First posted on arXiv 10 Feb 2025.

## Builds on (most cited first, 22 of 22)

- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space](https://mlsanity.com/paper/1902.10197/) (2019)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [Convolutional 2D Knowledge Graph Embeddings](https://mlsanity.com/paper/1707.01476/) (2017)
- [A Survey on In-context Learning](https://mlsanity.com/paper/2301.00234/) (2022)
- [Unifying Large Language Models and Knowledge Graphs: A Roadmap](https://mlsanity.com/paper/2306.08302/) (2023)
- [Composition-based Multi-Relational Graph Convolutional Networks](https://mlsanity.com/paper/1911.03082/) (2019)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [KG-BERT: BERT for Knowledge Graph Completion](https://mlsanity.com/paper/1909.03193/) (2019)
- [The Learnability of In-Context Learning](https://mlsanity.com/paper/2303.07895/) (2023)
- [MMKG: Multi-Modal Knowledge Graphs](https://mlsanity.com/paper/1903.05485/) (2019)
- [Learning to Exploit Long-term Relational Dependencies in Knowledge Graphs](https://mlsanity.com/paper/1905.04914/) (2019)
- [Large Language Models and Knowledge Graphs: Opportunities and Challenges](https://mlsanity.com/paper/2308.06374/) (2023)
- [Mitigating Large Language Model Hallucinations via Autonomous Knowledge Graph-based Retrofitting](https://mlsanity.com/paper/2311.13314/) (2023)
- [Image-embodied Knowledge Representation Learning](https://mlsanity.com/paper/1609.07028/) (2016)
- [HittER: Hierarchical Transformers for Knowledge Graph Embeddings](https://mlsanity.com/paper/2008.12813/) (2020)
- [MEAformer: Multi-modal Entity Alignment Transformer for Meta Modality Hybrid](https://mlsanity.com/paper/2212.14454/) (2022)
- [NativE: Multi-modal Knowledge Graph Completion in the Wild](https://mlsanity.com/paper/2406.17605/) (2024)
- [KGLM: Integrating Knowledge Graph Structure in Language Models for Link Prediction](https://mlsanity.com/paper/2211.02744/) (2022)
- [Unleashing the Power of Transformer for Graphs](https://mlsanity.com/paper/2202.10581/) (2022)

## Built on by (most cited first, 2 of 2)

- [LLMAtKGE: Large Language Models as Explainable Attackers against Knowledge Graph Embeddings](https://mlsanity.com/paper/2510.11584/) (2025)
- [Every Little Helps: Building Knowledge Graph Foundation Model with Fine-grained Transferable Multi-modal Tokens](https://mlsanity.com/paper/2602.15896/) (2026)

## Abstract (arXiv)

Recent advancements in large language models (LLMs) have significantly improved various natural language processing (NLP) tasks. Typically, LLMs are trained to predict the next token, aligning well with many NLP tasks. However, in knowledge graph (KG) scenarios, entities are the fundamental units and identifying an entity requires at least several tokens. This leads to a granularity mismatch between KGs and natural languages. To address this issue, we propose K-ON, which integrates KG knowledge into the LLM by employing multiple head layers for next k-step prediction. K-ON can not only generate entity-level results in one step, but also enables contrastive loss against entities, which is the most powerful tool in KG representation learning. Experimental results show that K-ON outperforms state-of-the-art methods that incorporate text and even the other modalities.