Graph-KV: Breaking Sequence via Injecting Structural Biases into Large Language Models

TL;DR

Graph-KV introduces a method to incorporate structural biases into large language models by using a graph-structured attention mechanism, improving performance on tasks involving complex data structures like graphs and long documents.

Contribution

The paper proposes Graph-KV, a novel approach that injects structural inductive biases into LLMs using graph-structured attention, enabling better handling of non-sequential data dependencies.

Findings

Outperforms baseline models on seven RAG benchmarks.

Effectively reduces positional bias in large language models.

Enhances reasoning and understanding in graph-structured tasks.

Abstract

Modern large language models (LLMs) are inherently auto-regressive, requiring input to be serialized into flat sequences regardless of their structural dependencies. This serialization hinders the model's ability to leverage structural inductive biases, especially in tasks such as retrieval-augmented generation (RAG) and reasoning on data with native graph structures, where inter-segment dependencies are crucial. We introduce Graph-KV with the potential to overcome this limitation. Graph-KV leverages the KV-cache of text segments as condensed representations and governs their interaction through structural inductive biases. In this framework, 'target' segments selectively attend only to the KV-caches of their designated 'source' segments, rather than all preceding segments in a serialized sequence. This approach induces a graph-structured block mask, sparsifying attention and enabling a message-passing-like step within the LLM. Furthermore, strategically allocated positional encodings for source and target segments reduce positional bias and context window consumption. We evaluate Graph-KV across three scenarios: (1) seven RAG benchmarks spanning direct inference, multi-hop reasoning, and long-document understanding; (2) Arxiv-QA, a novel academic paper QA task with full-text scientific papers structured as citation ego-graphs; and (3) paper topic classification within a citation network. By effectively reducing positional bias and harnessing structural inductive biases, Graph-KV substantially outperforms baselines, including standard costly sequential encoding, across various settings. Code and the Graph-KV data are publicly available.