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

2025-02-10

AI Summary

Large language models (LLMs) and knowledge graphs (KGs) suffer from a fundamental granularity mismatch—KGs operate at the entity level, whereas LLMs generate text token-by-token (often at the subword level). Method: We propose the first multi-step *k*-token head prediction architecture, introducing a *k*-step token prediction mechanism atop the LLM to enable end-to-end entity-level generation. Crucially, we integrate KG-aware contrastive loss directly into the LLM’s autoregressive decoding process, bridging token-level generation and entity-level reasoning. Our approach combines entity-level contrastive learning, KG embedding alignment, and frozen fine-tuning. Results: On multi-task KG reasoning benchmarks, our method significantly outperforms state-of-the-art methods—including those fusing textual and multimodal inputs—achieving substantial gains in entity recognition and linking accuracy while reducing inference latency by 37%.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion
• Natural Language Processing: Generation

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

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.

Problem

• Address granularity mismatch in knowledge graphs
• Integrate KG knowledge into LLMs
• Enable entity-level prediction and contrastive loss

Innovation

• Integrates KG into LLM
• Utilizes multiple head layers
• Enables entity-level predictions