PaperlayerView exact revision on arXiv

Source-linked AI summary

AGE: Adaptive-masking for Graph Embedding in Graph Retrieval-Augmented Generation

Bao Long Nguyen Huu, Atsushi Hashimoto

arXiv:2607.00052v1cs.IRcs.AI

TL;DR

GraphRAG can struggle to represent graph structure for frozen LLMs because graph-based and text-based latent features are misaligned. AGE addresses this with Transformer-based adaptive masking and self-supervised graph embedding, and it improves non-parametric GraphRAG methods across benchmark tasks. Its scope remains limited by fixed sampling, GraphRAG-only evaluation, small-scale models, and focus on representation rather than direct graph-learning tasks.

  • Problem

    GraphRAG needs better graph representations for frozen LLMs because graph-based and text-based latent features can be misaligned.

  • Method

    AGE uses Transformer-based mask-based self-supervised learning with a learnable sampler that predicts auxiliary nodes instead of masking key nodes.

  • Results

    AGE consistently improved G-Retriever and AMAR with frozen LLMs and achieved state-of-the-art results on three other benchmarks.

  • Takeaways & Limitations

    Adaptive masking provides a graph-embedding approach compatible with non-parametric retrieval, open LLMs, and frozen-LLM GraphRAG settings.

  • Takeaways & Limitations

    AGE uses a fixed sampling rate, has been evaluated only on GraphRAG tasks and primarily small-scale models, and does not directly target graph-learning tasks.

Abstract

from arXiv · show

GraphRAG is an extension of retrieval-augmented generation (RAG) that supports large language models (LLMs) by referring to graph-structured data as external knowledge. While this technique ideally captures intricate relationships, it often struggles with graph representations for LLMs, particularly for frozen LLMs, due to the misalignment between graph-based and text-based latent features. We tackle this issue by introducing the {\it Adaptive-masking for Graph Embedding (AGE)}. AGE employs a Transformer in a mask-based self-supervised learning (SSL) approach. We designed the architecture similar to text embedding encoders, addressing the latent feature misalignment. In contrast to natural language texts, graphs are concise representations, and there exist {\it key nodes} that hold dominant contextual information, which are challenging to predict from their surroundings. Masking such key nodes leads to inefficiency in the SSL process. Therefore, AGE focuses on predicting nodes apart from key nodes, utilizing a learnable node sampler. Our experimental results indicate that AGE significantly improves approaches using non-parametric search component in GraphQA tasks, achieving superior accuracy across four benchmark datasets with distinct characteristics.

1 Introduction

GraphRAG uses graph structure to improve access to complex relationships, but practical systems with non-parametric retrieval and frozen LLMs still face representation challenges. AGE addresses these challenges through adaptive masking in graph embedding self-supervision and reports improvements across GraphRAG benchmarks.

  • GraphRAG represents entities and relationships explicitly, supporting data integration, search accuracy, inference, and hallucination reduction.
  • Non-parametric retrievers are efficient and low-cost but may retrieve redundant or missing critical nodes without explicit structural constraints.
  • AGE uses a Transformer-style mask-based self-supervised process with a reinforcement-learning node sampler that selectively masks nodes rather than masking randomly.
  • AGE represents retrieved subgraphs through key-node and auxiliary-node embeddings learned with RL-guided mask-based self-supervised learning.
  • AGE improves GraphRAG methods using non-parametric retrieval and achieves state-of-the-art results on three additional benchmarks while using non-parametric retrieval and open LLMs.

2 Related Work

Prior work represents graphs for LLMs through textualization or graph embeddings, while self-supervised methods learn graph representations through contrastive or reconstruction objectives. AGE builds on these directions by improving graph embeddings with JEPA-based semantic prediction and practical prompt-tuning integration.

  • Graph Representation for LLMs: Graph representation for LLMs commonly uses textualization or graph embeddings, with G-Retriever adding graph embeddings as tunable prompts alongside textualized representations.
  • Graph Representation for LLMs: Textualization includes node sequences and reasoning paths, but studies report negative conclusions about interpreting text-encoded graphs with current LLMs.
  • Graph Representation for LLMs: The proposed training integrates self-supervised learning with prompt tuning in a one-stage process rather than prioritizing a two-stage procedure for each LLM input space.
  • Self-Supervised Learning: Contrastive joint-embedding methods learn node representations from positive and negative samples, whereas generative methods reconstruct masked input portions.
  • Self-Supervised Learning: JEPA avoids reconstructing pixel- or token-level details and instead targets more semantic self-supervised representations.

3 Preliminaries

GraphRAG seeks a relevant subgraph and text-modal knowledge for answering queries, while mask-based self-supervision and reinforcement learning provide the framework for learning graph representations. JEPA predicts masked representations without reconstructing unnecessary input details, and REINFORCE learns masking decisions from reconstruction rewards.

  • For a query on a textual graph, GraphRAG seeks an optimal relevant subgraph and text-modal knowledge to guide LLM generation.
  • The output sequence is modeled conditionally on the query, relevant text-modal knowledge, and optimal subgraph information.
  • JEPA uses an encoder, predictor, and target encoder to learn representations while avoiding reconstruction of unnecessary input details.
  • The JEPA loss applies only to predictions of masked input features, comparing predicted and stop-gradient target representations.
  • REINFORCE learns a masking policy over nodes, using reconstruction quality as the reward for each masking action.
  • The policy-gradient objective sums log-policy gradients over nodes and weights them by the reward associated with each masking action.

4 Approach

AGE converts retrieved subgraphs into graph embeddings through a graph encoder, adaptive node sampling, and a concept encoder-decoder trained against a target encoder. Separate objectives optimize downstream representations, reconstruction, and node sampling, with the decoder output used for inference.

  • Inference: At inference, AGE connects h_out to the downstream LLM instead of h_target, and aggregates h_out before projecting it to the LLM input dimension.The graph-structure-based aggregator uses the original subgraph edges before mean pooling and projection.
  • Architecture: AGE encodes the retrieved subgraph into node features, samples key nodes, reconstructs auxiliary-node representations, and aggregates the output for an LLM.Its modules include a graph encoder, node sampler, concept encoder-decoder, target encoder, and graph-structure-based aggregator.
  • Adaptive masking: The node sampler estimates node probabilities with attention, a linear layer, and softmax, then samples N_key = ⌈ρN⌉ nodes as key nodes.The remaining nodes form the auxiliary set used for prediction.
  • Representation learning: The concept encoder-decoder receives key-node features and predicts h_target, which the target encoder computes from all input nodes.The target encoder operates in a semantic space, while the decoder output h_out is trained to match the target representation.
  • Optimization: The target loss averages the L2 distance between predicted and stop-gradient target embeddings over auxiliary nodes, implementing knowledge distillation for reconstruction.The target loss is optimized for the graph encoder and concept encoder-decoder parameters.
  • Optimization: The sampling loss back-propagates only to the node sampler and increases a node’s sampling probability when its auxiliary prediction error is larger.This pushes nodes that are difficult to predict from their surroundings into the key-node set.

5 Experiments

AGE consistently improves frozen-LLM GraphRAG performance across benchmark settings and benefits from JEPA-based prediction with a learnable node sampler. Experiments also examine sampling-rate sensitivity and embedding organization.

  • Evaluation Setup: The evaluation uses accuracy for ExplaGraphs and SceneGraphs, and Hit@1 for WebQSP and CWQ.Experiments use frozen open-source Llama backbones and set the sampling rate to ρ = 0.3.
  • Main Results: AGE consistently improved G-Retriever and AMAR across frozen-LLM graph-embedding settings, with or without PEFT.The comparison covers ExplaGraphs, SceneGraphs, and WebQSP under multiple backbone LLMs.
  • Main Results: 26.72 percentage points was the largest reported gain for Llama3.2-1B with AGE over G-Retriever on ExplaGraphs without PEFT.The smallest reported gain was 2.02 points for Llama3.2-3B on WebQSP.
  • Ablation Study: 26.72% improvement was achieved by JEPA with the learnable node sampler over the baseline on Llama3.2 1B ExplaGraphs.The corresponding generative architecture with a random mask achieved 9.37%, while random-mask AGE achieved 15.46%.
  • Sampling Rate: ρ = 0.3 produced the best ExplaGraphs performance for both Llama3.2 settings: 81.4% for 1B and 92.6% for 3B.On WebQSP, the same rate was best for Llama3.2-1B at 62.5% and second best for Llama3.2-3B at 72.2%, behind 73.5% at ρ = 0.35.
  • Embedding Analysis: AGE reorganizes node-wise embeddings while preserving graph-text clustering structure in the graph encoder.Nonisolated key nodes show lower auxiliary-node prediction errors, indicating that relations between key nodes support prediction.

6 Limitation and Conclusion

The paper reports several scope and evaluation limitations for AGE. These include fixed sampling, limited task and model coverage, and an emphasis on KGQA rather than general graph learning.

  • Limitations: AGE uses a fixed sampling rate despite variation in key-node density across graphs.Adaptive sampling-rate selection remains an identified area for improvement.
  • Limitations: AGE was evaluated only on GraphRAG tasks, although the authors state that it may apply to other modalities.Broader applicability was not tested in the reported experiments.
  • Limitations: AGE primarily targets small-scale models, leaving effectiveness for large-scale LLMs unexplored because of computational constraints.The paper does not establish performance for larger LLMs.
  • Scope: The method represents retrieved structured data for LLMs rather than directly addressing node classification or link prediction.Without theoretical guarantees for node and link integration, the current scope is mainly KGQA.

Appendix AGE: Adaptive-masking for Graph Embedding in

The paper presents AGE, or Adaptive-masking for Graph Embedding, for representing graph-structured data in GraphRAG. The listed authors are Bao Long Nguyen Huu and Atsushi Hashimoto.

  • Paper Identity: AGE stands for Adaptive-masking for Graph Embedding and is proposed for GraphRAG graph representation.The title identifies the method and its application context.
  • Authors: The paper is authored by Bao Long Nguyen Huu and Atsushi Hashimoto.The affiliations listed are OMRON Corporation and OMRON SINIC X Corporation.

1 Proof of Concept

The proof of concept frames retrieved subgraph representation as a bottleneck for frozen LLMs and motivates AGE’s RL-guided mask-based self-supervised embedding. AGE is designed to preserve structural and semantic information while reducing diluted representations that hinder reasoning.

  • Problem definition: Static retrieval can produce corrupted subgraphs containing redundant nodes or missing critical elements.Popular retrieval systems select top-k nodes by cosine similarity, but non-optimized retrieval may yield incomplete or redundant subgraphs.
  • AGE embedding module: The embedding module is trained to transform a retrieved subgraph representation into an embedding that guides an LLM toward expected answers.The pipeline combines graph representation, contextualized reasoning, and LLM generation over retrieved text and subgraph knowledge.
  • Problem definition: The generation objective decomposes LLM output into knowledge recall followed by contextualized reasoning conditioned on the recalled knowledge.The formulation models response generation as interleaved domain knowledge and reasoning steps.
  • Problem definition: Weighted-sum aggregation without explicit structural constraints may dilute node representations and make relationships harder for frozen LLMs to capture.The paper links diluted subgraph representations to weakened knowledge recall and reasoning in frozen LLMs.
  • AGE embedding module: AGE uses RL-guided masking and reconstruction to learn structural-dependency representations while selectively preserving critical nodes.The sampler estimates which nodes are critical for structural and semantic information, then guides mask-based self-supervised learning.
  • Problem definition: When retrieved representations are diluted, the loss shifts toward knowledge identification rather than reasoning, creating inefficient knowledge use.The paper contrasts this with explicit subgraph representations, where the loss shifts toward integration and supports application of retrieved graph knowledge during fine-tuning.

2 Additional Experimental Details

Additional experiments evaluate training configurations, alternative graph-structure component arrangements, and implementation settings. The combined graph encoder and graph-structure-based aggregator performs best in the reported WebQSP comparison.

  • Training settings: The experiments use AdamW with learning rate 1e−4, weight decay 0.05, cosine decay after warm-up, and early stopping with patience 3.Experiments run on either 2 NVIDIA 2080Ti-11G or 2 NVIDIA A100-80G GPUs.
  • Training settings: AMAR comparisons retain its training settings, including dataset-specific retrieved-data counts, soft-prompt lengths, beam-search numbers, and LoRA parameters.WebQSP and CWQ use different retrieval and prompting configurations; Llama2 uses LoRA learning rate 5e−5 and scaling factor 32.
  • Architecture arrangement: The architecture investigation compares graph encoder only, graph-structure-based aggregator only, and their combination.The comparison is presented as an investigation of core component arrangement.
  • Architecture arrangement: 73.46% Hit@1 is achieved by the combined architecture on WebQSP, versus 71.12% with Graph Encoder and 72.44% with graph-structure-based aggregator.The combined graph encoder and graph-structure-based aggregator is the best-performing arrangement in this comparison.

The Choice of GNN on AGE G-Retriever

The study compares GCN, GAT, and Graph Transformer as graph encoders and graph-structure-based aggregators on WebQSP and ExplaGraphs. Graph Transformer performs best on ExplaGraphs and is selected for all experiments.

  • WebQSP: On WebQSP, GCN, GAT, and Graph Transformer achieve Hit@1 scores of 56.75, 61.42, and 62.53, respectively.The comparison evaluates the three GNN architectures under the same reported dataset setting.
  • ExplaGraphs: On ExplaGraphs, Graph Transformer reaches 0.8501 accuracy, ahead of GCN at 0.8321 and GAT at 0.8212.The reported ordering motivates choosing Graph Transformer across datasets.
  • Architecture selection: Graph Transformer is chosen for all experiments to maintain performance across datasets with different properties and demands.The paper presents architecture selection as important for dataset-specific performance.

The design of AGE with Generation Architecture

The generation architecture complements input nodes during inference by training an encoder-decoder to reconstruct randomly masked nodes. The appendix also analyzes the number of graph-embedding layers on WebQSP.

  • Generation architecture: The generation architecture randomly masks input nodes at a 70% ratio during encoder-decoder training.The encoder embeds unmasked nodes, while the decoder reconstructs the masked nodes.
  • Generation architecture: Its purpose is to complement input nodes and enhance the graph-structure-based aggregator’s embedding during inference.The design is presented as an auxiliary generative architecture for AGE.
  • Layer analysis: Table B.2 analyzes the number of GNNge layers with LLaMA 3.2 3B on WebQSP.GE denotes Graph Embedding in the table caption.

Analysis on the Layer Number of the Graph Encoder GNNge on AGE

AGE maintains or improves performance without requiring deeper graph encoders, while its adaptive sampling and embedding design support efficient training and retrieval across settings. The analyses also show robustness to retrieval volume and stronger comparisons with non-parametric baselines.

  • Graph encoder depth: AGE achieves superior performance while maintaining the baseline training time, because its additional modules substitute for deeper GNN layers.G-Retriever benefits from deeper GNNge, whereas AGE performs better with fewer layers.
  • Graph encoder depth: AGE outperforms G-Retriever by approximately 10 points with Llama 3.2 1B and 6 points with Llama 3.2 3B, even against a 20-layer GNN baseline.The 20-layer G-Retriever is configured to have computational cost similar to AGE with two layers.
  • Key-node sampling: RL-based node samplers achieve high performance at lower sampling rates than static PageRank and Degree Centrality strategies.Static strategies tend to require higher sampling rates, whereas RL-selected key nodes more effectively support the concept encoder-decoder.
  • Retrieval volume: AGE remains stable across retrieval volumes, scoring 86.8 with 5 retrievals, 87.0 with 10 retrievals, and 86.5 with 100 retrievals.The reported stability is contrasted with limited AMAR improvement when recalled information is insufficient and noise when large volumes are retrieved.
  • Retriever comparison: AGE with a non-parametric retriever reaches 80.3% Hit@1 on WebQSP, exceeding GraphToken at 57.1% and G-Retriever at 70.2%.AGE is reported as approaching DualR with a parametric retriever, which achieves 82.8%.
Loading 2607.00052v1…