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

1 Introduction

• GraphRAG addresses RAG’s difficulty capturing structured relationships, while AGE targets practical settings with frozen LLMs and non-parametric retrieval that can omit critical nodes or include redundant ones.Graph embeddings are intended to bridge retrieved graph data and the LLM input space.Source: §1 · pp. 1–2
• AGE uses reinforcement-learning-guided masking to distinguish key from auxiliary nodes, reporting effectiveness over random masking and state-of-the-art performance on three other benchmarks.It targets GraphRAG with a non-parametric retriever and open LLMs.Source: §1 · p. 3

2.1 Graph Representation for LLMs

• AGE improves graph embeddings in G-Retriever and uses one-stage SSL integrated with prompt tuning rather than two-stage SSL followed by prompt alignment.Prior work also represents graphs through textualized node sequences, but studies report limitations in LLM interpretation of text-encoded graphs.Source: §2.1 · pp. 3–4

2.2 Self-Supervised Learning

• AGE applies JEPA’s abstract, semantic-level prediction to GraphRAG and adds a node sampler to select key knowledge and link it with related data.Unlike generative methods that reconstruct masked input details, JEPA avoids reconstructing pixel- or token-level detail.Source: §2.2 · p. 4

3 Preliminaries

• GraphQA seeks a query-relevant subgraph and text-modal knowledge to condition the LLM’s answer generation.The preliminaries define the optimal subgraph relative to a textual graph and query.Source: §3 · p. 4
• JEPA uses an encoder, predictor, and target encoder with stop-gradient, and applies its loss only to predictions of masked input.It focuses on abstract representations rather than reconstructing unnecessary input details.Source: §3 · p. 5
• REINFORCE trains a node-level policy to choose whether to mask each retrieved node, maximizing reconstruction quality with masked-node SSL as its environment.The policy estimates the probability of masking or not masking a node.Source: §3 · p. 5

4 Approach

• The approach applies SentenceBert to indexed knowledge, retrieves with static k-nearest neighbors, constructs a Prize-Collecting Steiner Tree subgraph, and permits arbitrary inference-time LLMs.The section focuses on the embedding step.Source: §4 · p. 5

4.1 Text Embedding of Query and Text Graph

• AGE concatenates the textualized retrieved subgraph with the query and embeds the combined text using a pretrained function for the frozen LLM.The resulting text embedding is an output token sequence.Source: §4.1 · p. 5

4.2 Adaptive-masking for Graph Embedding

• AGE encodes the retrieved graph, samples key nodes, and uses a concept encoder-decoder to predict target-encoder representations for auxiliary nodes before aggregating the output into an LLM token.The target encoder is trained with downstream tasks, while the graph-structure-based aggregator projects the output using the subgraph’s edges and a projector maps it to the LLM input dimension.Source: §4.2 · pp. 6–8

4.3 Optimization of Adaptive-masking for Graph Embedding

• AGE uses separate prompt-tuning, auxiliary-node reconstruction, and sampler losses to optimize different modules without overlap or loss-weight adjustment.The sampler loss uses reconstruction error to favor selecting nodes that are difficult to predict from their surroundings as key nodes.Source: §4.3 · pp. 8–9

4.4 Analysis of Learning Objectives

• The learning-objective analysis separates latent identification from representation learning and argues that a learned latent space restructured for the frozen LLM can improve knowledge recall and indirectly enhance reasoning.The paper says this separation provides efficient learning signals, faster convergence, and improved graph representations.Source: §4.4 · p. 10

5 Experiments

• Across ExplaGraphs, SceneGraphs, WebQSP, and CWQ, the experiments use accuracy for the first two datasets and Hit@1 for the latter two.The datasets cover commonsense reasoning, visual question answering, and question answering over large knowledge graphs.Source: §5 · pp. 10–11

5.1 Main Results

• AGE consistently improves G-Retriever and AMAR across backbones in frozen-LLM graph-embedding settings; its largest reported gain is 26.72 percentage points on ExplaGraphs, and its smallest is 2.02 points on WebQSP.AGE AMAR outperforms ReKnoS on WebQSP and CWQ; versus Paths-over-Graph, Plan-on-Graph, and DoG, it underperforms on WebQSP but outperforms them on CWQ.Source: §5.1 · p. 11, Table 1 · p. 12

5.2 Ablation study

• With Llama3.2 1B on ExplaGraphs, GA with random masking improves 9.37% over baseline, random masking improves 15.46%, and JEPA with the learnable sampler improves 26.72%.The comparison supports JEPA over GA and indicates that the learnable sampler adds improvement over random masking.Source: §5.2 · p. 12, Table 2 · p. 13
• A sampling rate of ρ = 0.3 gives the best ExplaGraphs scores for both Llama3.2 1B and 3B, and the best WebQSP score for 1B.On WebQSP with Llama3.2 3B, ρ = 0.3 is second best at 72.2%, behind 73.5% at ρ = 0.35; the paper uses ρ = 0.3 in its experiments.Source: §5.2 · p. 13, Figure 3 · p. 13, Figure 4 · p. 13
• Equal loss weighting (1:1:1) consistently performs best across datasets and model sizes, while deviating from equal weights degrades performance without consistent benefit.The losses optimize disjoint parameters, so the authors conclude explicit weighting is unnecessary by design.Source: §5.2 · p. 13, Table R2 · p. 15
• Qualitative results show AGE reorganizes node-wise embeddings, and nonisolated key nodes have lower auxiliary-node prediction errors, suggesting their relations support prediction.Examples include inferring “missionaries” and “Christians” from the key node “saving souls.”Source: §5.2 · pp. 13–14, Figure 5 · p. 14

Conclusion

• AGE consistently improved GraphRAG performance while maintaining computational cost, and its node sampler identified key nodes in the ablation study.The method introduces JEPA to enhance graph-structure embeddings for downstream reasoning tasks.Source: §6 · p. 15

Limitations

• AGE used a fixed sampling rate despite variation in key-node density, and was evaluated only on GraphRAG tasks although it may apply to other modalities.The authors identify adaptive sampling and broader evaluation as future directions.Source: §6 · p. 15
• AGE primarily targets small-scale models; its effectiveness for large-scale LLMs remains unexplored because of computational constraints.Source: §6 · p. 15
• AGE represents retrieved structured data for LLMs rather than directly addressing graph-learning tasks, leaving its current scope mainly to KGQA without theoretical guarantees for node and link integration.Source: §6 · p. 15

• The formulation decomposes responses into recalled knowledge and contextualized reasoning, with a static single-turn retriever providing text knowledge and a subgraph for generation.The graph embedding module represents the retrieved subgraph to guide the LLM toward expected answers.Source: §1 · p. 21
• Static retrieval often selects top-k graph elements by cosine similarity, but retrieved subgraphs may contain redundant nodes or lack suggestive or critical elements.Source: §1 · pp. 21–22
• Without explicit structural constraints, attention-weighted GNN or Transformer aggregation can dilute node representations, making it harder for frozen LLMs to capture relationships.Source: §1 · p. 22
• AGE uses reinforcement learning to guide mask-based self-supervision, reconstructing masked node features from sampled key-node encodings to impose structural constraints on representations.The sampler’s binary action masks or keeps each node; masked features serve as decoder input and key-node features as encoder input.Source: §1 · pp. 22–23
• Explicit subgraph representations shift fine-tuning emphasis toward knowledge integration and reasoning, whereas diluted representations prioritize identifying graph patterns over applying them in reasoning.For frozen LLMs, weak knowledge recall can cascade into flawed reasoning and degraded response quality.Source: §1 · pp. 23–24

2.1 Implementation Settings (AGE G-Retriever)

• AGE G-Retriever uses AdamW, a 1e −4 initial learning rate, 0.05 weight decay, cosine decay after warm-up, and early stopping with patience 3.The Graph Transformer uses 2 layers, 4 attention heads, and hidden dimension 1024, following. LLMs are Llama3.2 1B/3B and Llama3.1 8B; LoRA scaling is 16, input length is capped at 512, and generation at 32 new tokens. Subgraph top-k is 3 for SceneGraphs, nodes 3 and edges 5 for WebQSP, and 0 for ExplaGraphs to retrieve its full graph.Source: §2.1 · p. 24

2.2 Implementation Settings (AGE AMAR)

• AGE AMAR retains AMAR [82] training settings, with dataset-specific retrieved-data counts, soft-prompt lengths, beam-search counts, and 256 maximum new tokens.WebQSP uses 100 retrieved items, prompt length 7, and beam search 8; CWQ uses 4, 16, and 15, respectively. Llama2 [73] uses LoRA learning rate 5e −5 and scaling factor 32.Source: §2.2 · p. 24

2.3 The choice of AGE architecture

• 73.46% Hit@1 on WebQSP for the combined graph encoder and graph-structure-based aggregator exceeds 71.12% and 72.44% for either component alone.Figure B.1 depicts the three tested arrangements.Source: Figure B.1 · p. 25, §2.3 · p. 25

The Choice of GNN on AGE G-Retriever

• WebQSP Hit@1 is 62.53 for Graph Transformer, 61.42 for GAT, and 56.75 for GCN; ExplaGraphs accuracy is 0.8501, 0.8212, and 0.8321 for Graph Transformer, GAT, and GCN, respectively.Source: The Choice of GNN on AGE G-Retriever · pp. 25–26
• AGE uses Graph Transformer for all experiments to maintain performance across datasets, given the importance of choosing an architecture suited to dataset properties.Source: The Choice of GNN on AGE G-Retriever · p. 26

The design of AGE with Generation Architecture

• GA complements input nodes to enhance the graph-structure-based aggregator’s embeddings during inference.At inference, random masking is disabled; all input nodes pass through the encoder-decoder, which reconstructs the input for the aggregator.Source: The design of AGE with Generation… · p. 26
• During training, GA randomly masks 70% of input nodes, embeds unmasked nodes, and trains the decoder to reconstruct masked nodes through the target loss.Source: The design of AGE with Generation… · p. 26
• Prompt-tuning loss trains the graph-structure-based aggregator and MLP to refer to edges E∗, while the MLP aligns aggregated embeddings with the LLM input dimension.Source: The design of AGE with Generation… · p. 26

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

• Unlike G-Retriever, which benefits from deeper GNNge layers, AGE performs better with fewer layers and achieves superior performance while maintaining baseline training time.AGE’s additional inference modules effectively substitute for reduced GNN layers.Source: Analysis on the Layer Number of the… · p. 26

Comparison with Deeper GNNs

• AGE outperforms 20-layer G-Retriever by approximately 10 points with Llama 3.2 1B and 6 points with Llama 3.2 3B, despite similar computational cost to 2-layer AGE.Using 20 G-Retriever layers substantially improves its performance, but AGE still achieves the stated margins.Source: Comparison with Deeper GNNs · p. 27

The choice of concept decoder and node sampler architecture.

• On ExplaGraphs and WebQSP, a one-block concept decoder performs best; increasing depth from 1 to 4 blocks raises parameters from 81M to 106M but drops WebQSP from 62.5 to 57.4.The study therefore selects one transformer block for the concept decoder.Source: The choice of concept decoder and node… · p. 27
• Adding transformer blocks to the node sampler yields only marginal WebQSP performance gains and requires more memory.The authors choose a single transformer block with hidden dimension 1024 for subsequent experiments to balance computational and performance efficiency.Source: The choice of concept decoder and node… · p. 27
• The selected node sampler uses one transformer block with hidden dimension 1024 in subsequent experiments.This configuration was chosen to maintain computational and performance efficiency.Source: The choice of concept decoder and node… · p. 27

The study on key nodes sampling

• Static PageRank and Degree Centrality sampling tend to perform better at higher sampling rates on ExplaGraphs and WebQSP, suggesting their selected key nodes are less impactful.These static strategies require a larger set of key nodes to support graph embedding and guide the LLM toward desired answers.Source: The study on key nodes sampling · pp. 28–29, Figure B.3 · p. 28, Figure B.4 · p. 28
• RL-based node samplers achieve high performance at lower sampling rates, indicating their selected key nodes support the Concept Encoder-Decoder more effectively and can guide the LLM with fewer nodes.The reported outcome is that the LLM can produce expected answers with fewer key nodes in the guidance process.Source: The study on key nodes sampling · pp. 28–29, Figure B.3 · p. 28, Figure B.4 · p. 28

The study on transferability

• AGE supports cross-dataset transfer: when trained on WebQSP, it improves generation on ExplaGraphs with Llama 3.2 3B and outperforms GRAG’s transferability.The passage says training on a large dataset can enhance generation on a smaller one.Source: The study on transferability · p. 29

The study on number of retrieval

• AGE scored 86.8 at 5 retrievals and 87.0 at 10; at 100 retrievals, Hit1@ was 86.5, with minimal performance variation.The passage attributes the gains at smaller retrieval counts to AGE mapping retrieved data to useful embeddings, and says large retrieval volumes introduce noise for LLMs.Source: Figure B.5 · p. 29, The study on number of retrieval · p. 30

Impart of AGE with LoRA Finetuning performance

• AGE’s mask-based SSL reconstruction objective imposes structural constraints in embedding space, producing faster convergence and lower initial-stage training loss than G-Retriever.The comparison is reported across ExplaGraphs and WebQSP training-loss trends.Source: Impart of AGE with LoRA Finetuning… · p. 30, Figure B.6 · p. 30, Figure B.7 · p. 31

The choice of RL method for Node Sampler

• On ExplaGraphs, Gumbel-Softmax and REINFORCE outperform the Straight-Through Estimator; on WebQSP, REINFORCE leads slightly and was selected.The comparison concerns three node-sampler RL methods.Source: The choice of RL method for Node Sampler · pp. 30–31, Table B.6 · p. 31
• Using hout as LLM input at inference outperforms htarget on both datasets; connecting htarget during training to hout during inference avoids internalizing the learning-inference mismatch.The passage says this design lets the student model surpass the teacher and supports a more robust representation.Source: The choice of RL method for Node Sampler · p. 31, Table B.7 · p. 32

The stability of Target encoder

• Identical concept- and target-encoder parameters help prevent distribution shifts; normalization improves stability, normalization alone scores highest, and EMA adds further enhancements.Table B.8 reports MSE performs better with L1; normalization is intended to keep activation distributions consistent and reduce internal covariate shift.Source: The stability of Target encoder · pp. 31–32, Table B.8 · p. 32

The landscape of existing KGQA methods

• KGQA methods include GNN-based NSM, ReaRev, and G-Retriever on dense subgraphs, while ToG [66] retrieves facts hop-by-hop and RoG [50] maps generated relation paths to the KG.Frequent LLM calls in LLM-based retrieval and reasoning increase training and inference costs.Source: The landscape of existing KGQA methods · p. 32, Figure B.8 · p. 33
• AGE enhances GNN graph embeddings with RL-inspired supervision in SSL, improving non-parametric retrieval to performance comparable to LLM-based retrievers.This approach is intended to improve LLM reasoning without relying on frequent LLM calls for retrieval.Source: The landscape of existing KGQA methods · p. 32

2.4 Additional Qualitative Evaluation

• The qualitative examples show that auxiliary-node embeddings are easier to predict from key nodes than key-node embeddings are from auxiliary nodes.The passage illustrates this asymmetry with examples from ExplaGraphs and notes similar patterns in additional ExplaGraphs and WebQSP visualizations.Source: §2.4 · p. 33, Figure B.10 · p. 37, Figure B.11 · p. 38
• In a WebQSP failure case, the true answer node 98175 had high auxiliary-node target loss and was omitted by the LLM; adjusting the trainable sampling rate was suggested as a possible solution.The case used AGE trained with LLaMA 3.2 1B at a sampling rate of 0.3.Source: §2.4 · pp. 33–34, Figure B.9 · p. 34

Training Computational Resources on AGE G-Retriever

• AGE with Llama3.1 8B reaches Hit@1 78.25 in 6.4 minutes per epoch, while its LoRA-tuned setting reaches Hit@1 80.34 in 7.3 minutes per epoch.The passage also reports that AGE with Llama3.2 3B outperforms G-Retriever with Llama2 7B without longer training time.Source: Training Computational Resources on… · p. 35, Table B.9 · p. 35

Inference Computational Resources on AGE G-Retriever

• AGE Llama 3.2 3B offers balanced Hit@1 performance at 87.6 tokens per second, while Llama 3.2 1B runs at 148.5 tokens per second with lower Hit@1.The passage attributes the 1B model’s greater speed to its smaller parameter count and notes that higher-parameter models provide superior accuracy.Source: Inference Computational Resources on… · pp. 35–36, Table B.10 · p. 35

Comparison with trainable retriever methods

• AGE (Llama3.1 8B-LoRA, non-parametric) reaches 80.3% Hit@1 on WebQSP, near DualR (ChatGPT, parametric) at 82.8%, and above GraphToken at 57.1% and G-Retriever at 70.2%.The comparison reports that AGE closes the performance gap without the added complexity and training overhead associated with parametric retrievers.Source: Comparison with trainable retriever… · pp. 36–37, Table B.11 · p. 36