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

Lingbing Guo, Yichi Zhang, Zhongpu Bo, Zhuo Chen, Mengshu Sun +3 more
2/10/2025
cs.CLcs.AI

No official code found. Best implementation: zjukg/K-ON (11 stars, MIT, last commit 2025-02-22).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

zjukg/K-ONCommunity73%
MIT

K-ON: Stacking Knowledge On the Head Layer of Large Language Model, AAAI 2025 Oral

111Shell, PythonFeb 22, 20251 years agoMIT

Similar Papers

Cite this paper

@article{guo2025kon,
  title  = {K-ON: Stacking Knowledge On the Head Layer of Large Language Model},
  author = {Lingbing Guo and Yichi Zhang and Zhongpu Bo and Zhuo Chen and Mengshu Sun and Zhiqiang Zhang and Wen Zhang and Huajun Chen},
  year   = {2025},
  eprint = {2502.06257},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2502.06257v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2502.06257.svg)](https://codeofpaper.com/paper/2502.06257)

Discussion