1 Introduction↩︎

Large Language Models (LLMs) such as GPT [1], [2], Claude [3], Gemini [4], Qwen [5], and LLaMA [6] have significantly advanced natural language understanding and generation capabilities. Retriever-Augmented Generation (RAG) [7]–[9] integrates query-relevant information into the generation process, enabling LLMs to access and utilize domain-specific knowledge beyond their pretraining corpus. However, although RAG enhances LLMs with external data, it may struggle to capture essential structured relationships, reducing search precision and reasoning effectiveness [10], [11].

Graph Retriever Augment Generation (GraphRAG) [13], [14] is a technology that uses graphs to overcome the limitations of RAG. Graph data, represented by nodes (entities) and edges (relationships), clearly presents complex relationships. This provides several benefits, such as facilitating data integration [15], improving search accuracy [16], [17], enhancing inference capabilities [18], [19], and reducing hallucinations [12]. By capturing sub-graphs, the broader context and interconnections within the graph structure can be captured, enabling comprehensive information to be accessed for LLMs enhance the performance in domain-specific tasks.
This study investigates GraphRAG methods that operate within practical computational costs. Fine-tuning LLMs can enhance GraphRAG performance, yet it is resource-intensive. Instead, previous methods often focused on the retrieval module as it is a key factor for GraphRAG performance. Trainable retrievers, such as LLM-based retrievers [20], [21] realize a higher retrieval accuracy. However, this strategy still requires significant computational overhead. Non-parametric retrievers [12], [22] are efficient and low-cost but may contain redundant or missing critical nodes, leading to the lack of explicit structural constraints. To maintain practicality, we base our method on non-parametric retrievers with frozen LLM, and improve performance of structural representation by updating the graph embedding module. Embeddings play a crucial role in bridging the gap between retrieved graph data and the LLM input space. Multiple methods [12], [23] use graph embedding together with textualized graph representation (Fig. 1), indicating that the graph encoder should embed relationships between elements, rather than their individual contents. What is the optimal strategy to achieve such encoding for a frozen LLM? We considered two factors: similarity of the embedding space to the LLM’s text encoder and its relationship embedding capability. Since the LLM’s text encoder uses mask-based SSL [24]–[26], which learns to embed relationships between elements into the embedding space by optimize the reconstruction of masked elements. These factors aim to embed relationships between nodes within the retrieved subgraph into the embedding space by adapting the mask-based SSL with minimal modification.
To realize this intention, we propose a novel embedding strategy, the Adaptive-masking for Graph Embedding (AGE). The architecture of AGE is designed to imitate the general self-supervised text embedding process, while incorporating the Joint-Embedding Predictive Architecture (JEPA) [27], which improves representations in the embedding space by eliminating unnecessary detail reconstruction. Although generative SSL shows promise, the quality of reconstruction relies on the discriminability of input nodes [28], [29]. Random masking fails on non-discriminative nodes, leading to poor representations [30], [31]. To avoid this, the only major modification is adding a node sampler trained via reinforcement learning (RL) to selectively mask nodes, replacing traditional random node masking with an adaptive approach. The motivation for this approach stems from the fact that, graphs are concise logical structures with minimal redundancy. Hence, some nodes are crucial for maintaining graph integrity; we refer to them as key nodes. Our RL-based strategy aims to guide SSL to distinguish the representations of key and auxiliary nodes, encouraging LLMs to identify redundant information within the retrieved graph. The contributions of this paper are outlined as follows:

1. We propose AGE, a novel method that represents retrieved subgraphs via key-node and auxiliary-node embedded by RL-guided mask-based SSL.

2. Our study reveals adaptive masking approach’s notable effectiveness over random masking within GraphRAG.

3. AGE uses a non-parametric retriever and open LLMs, while also achieving SOTA on three other benchmarks.

2 Related Work↩︎

2.1 Graph Representation for LLMs↩︎

In the context of representing graphs as input to LLMs, it is necessary to first convert the retrieved graph data into specific formats. We summarize two distinct formats: textualization and graph embeddings. Textualization [32]–[34] is a text-based formalization method designed to characterize and represent graph data. Node sequences are a popular form of textualization[20], [21], [35]. Some methods [20], [21], [36] propose LLM-based retrievers to extract reasoning paths. A node sequence ordered along the path aids LLM’s reasoning. However, many studies report negative conclusions in interpreting text-encoded graphs with concurrent LLMs [37]–[39], suggesting a need for solutions beyond textualization.
The other format, graph embeddings, have recently been adopted in GraphToken [23]. Following this, G-Retriever [12] proposed a retrieval framework for graph embeddings. This occurs when graph embeddings are added as tunable prompts to the LLM in addition to their textualized representations. In this work, we improve the quality of LLM responses on the G-Retriever framework through enhancing the representation of graph embeddings. Some methods [15], [40] build self-alignment and cross-question module among retrieved entities, relations, and subgraph embedding elements. Some methods enhance embeddings through a two-stage training process [41], [42]. The first stage trains the embedding module on SSL alone; in the second stage, prompt tuning aligns the structured relationships embedded for LLM input by the pretrained module. Since each LLM has its own domain embeddings and input spaces, two-stage training process prioritize maximizing performance. Instead, focusing on practical, we propose a one-stage training process SSL that integrates with prompt tuning.

2.2 Self-Supervised Learning↩︎

Many existing self-supervised learning architectures focus on learning representations that effectively capture relationships between input data. Joint-Embedding Architecture (JEA) [43]–[45] has shown considerable promise in advancing SSL methodologies. Joint-Embedding SSL for GNNs, such as GraphCL [46], GCA [47] and JOAO [48], learn node representations by contrasting positive and negative samples. Subsequent studies identified areas for enhancement in JEA [49]–[51], particularly the issue of mapping all inputs to a single constant vector, known as the collapsing problem. Generative Architecture (GA) [52]–[54] focuses on reconstructing masked portions of the input at either the pixel or token level. GraphMAE [55], [56] learns representations by reconstructing masked samples. These methods encourage the model to learn more robust and diverse representations, potentially reducing the risk of representation collapse [49], [50], [57]. Joint-Embedding Predictive Architectur (JEPA) [27] eliminate reconstruction of pixel or token-level details and enhances the semantic level of self-supervised representations [57], [58]. In this work, we first demonstrate JEPA’s effectiveness in the GraphRAG framework. JEPA originates from cognitive neuroscience, suggesting that humans have an ability of top-down schema reasoning, aiding planning, decision-making, and problem-solving on complex tasks [59]–[61]. In GraphRAG, the tokens fed into the LLM’s hidden layer should capture the ability. We implement it as JEPA in the graph embedding module. Cognitive science has revealed that a brain region called the temporal lobe plays a role in bottom-up associative learning, selecting key knowledge and linking it to related data [62]–[64]. These selection processes can also occur with new information. Aiming to reproduce this functionality in GraphRAG, AGE has the novel node sampler module.

3 Preliminaries↩︎

GQA with LLM. For a query \(q\) on a textual graph \(G\), there is an optimal subgraph \(\overline{S^{*}} \in S(G)\) and query relevant text-modal knowledge \(T^{*}\) that guides the LLM to produce expected answers, where \(S(G)\) is the set of all subgraphs of \(G\). The challenge of GraphRAG is to efficiently search for the relevant subgraph \(S^{*}\) and represent it to \(\overline{S^{*}}\) for an LLM \(p_\Phi\) improve generation. The probability distribution of the output sequence \(Y\) is given by: \[p_\Phi(Y \mid [q, G]) = \prod_{i=1}^{n} p_\Phi(y_i \mid y_{<i}, [q, T^{*}, \overline{S^{*}}]),\] where \(y_{<i}\) represents the prefix tokens, and \([q, \overline{S^{*}}]\) indicates the concatenation of the query, relevant text-modal knowledge and optimal subgraph information, respectively.
Joint-Embedding Predictive Architecture. Mask-based SSL methods such as MAE are well suited to handle corrupted input, as they learn to reconstruct missing or corrupted input parts. JEPA improves upon MAE by eliminating the reconstruction of unnecessary input feature details, focusing instead on learning more abstract representations. JEPA consists of an encoder \(E_\theta(\cdot)\), predictor \(P_{\phi}(\cdot)\) and target encoder \(T_{\theta}(\cdot)\). The stop-gradient operation \(\operatorname{sg}\) is employed to prevent representation collapse in the target encoder \(T_{\theta}(\cdot)\). The predictor generates \(y\) from visible input \(x\) and masked input \(\Delta_x\). The encoder and predictor are trained simultaneously with the objective: \[\min || P_{\phi}\big(\Delta_x, E_{\theta}(x)\big) - \operatorname{sg}\big(T_{\theta}(y)\big) ||_2,\] The loss is applied only to the predictions of the masked input \(\Delta_x\).
Reinforcement Learning. To estimate the key and auxiliary nodes on retrieved graph for mask-based SSL discriminative embedding. We adopt REINFORCE [65], a basic policy gradient method in RL. Let \(\mathcal{D}={(q,Y^{*})}\) denote a corpus of training data, where \(Y^{*}\) is the complete reference label for query \(q\). REINFORCE optimizes a policy \(\pi_{\theta}\) parameterized by \(\theta\), to maximize reconstruction quality \(R_a\) for each masking action \(a\). The policy gradient is given by: \[\nabla_\theta \mathcal{J}(\theta) = \mathbb{E}_{q \sim \mathcal{D}, S^{*} \sim q} \left[ \sum_{v \in S^{*}} \nabla \log \pi_\theta(a | v) \cdot R_a \right]\] where where \(\mathcal{J}(\theta)\) is the expected return. \(\pi_\theta(a | v)\) denotes the probability distribution estimated by policy \(\pi_{\theta}\) for taking action \(a\) (masking or not) on node \(v\). The SSL framework with masked node reconstruction serves as the RL environment.

4 Approach↩︎

Our framework, illustrated in Figure 1, consists of four main steps: input, graph preprocessing, embedding, and inference. We adopt the previous method [12] that applies SentenceBert [25] to indexed knowledge data at (input) step and employ a static k-nearest neighbors [66] retrieval approach combined with Prize-Collecting Steiner Tree [67] subgraph construction during graph preprocessing. For inference, we can use arbitrary LLMs, as usual RAG methods. Therefore, this section focuses on the details of embedding step.

4.1 Text Embedding of Query and Text Graph↩︎

We transform the retrieved subgraph \(S^{*}\) into a textual format, following [12] as in the first two steps. The converted text is then concatenated with the input query \(q\). The concatenated texts are embedded into \(h_{\text{text}}\) using a pretrained function for the frozen LLM, \(\text{TextEmbedding}\), where \([;]\) denotes concatenation and \(L\) is the output token sequence length, as follows: \[\begin{gather} h_{\text{text}}=\text{TextEmbedding}(\text{[textualize}(S^{*});q]) \in \mathbb{R}^{L\times d_{l}}, \end{gather}\]

4.2 Adaptive-masking for Graph Embedding↩︎

AGE comprises a node sampler, concept encoder-decoder, target encoder, and graph-structure-based aggregator modules, as overviewed in Fig. 2. The AGE’s input, \({\boldsymbol{h}}_{\text{in}}\), is encoded from the retrieved graph \(S^{*}\) with a conventional graph encoder. \({\boldsymbol{h}}_{\text{in}}\) is then passed to node sampler and target encoder. The node sampler categorizes the nodes into key nodes and the remaining auxiliary nodes. The selected key nodes are fed to concept encoder-decoder. The output \({\boldsymbol{h}}_{\text{out}}\) is trained to predict \({\boldsymbol{h}}_{\text{target}}\), the output of target encoder, forming a JEPA. The embedding \({\boldsymbol{h}}_{\text{out}}\) is then aggregated into a token to be fed to the LLM. The rest of this subsection explains each module in Fig. 2 individually.
Graph Encoder prepares input for AGE. The retrieved subgraph \(S^{*}=(V^{*},E^{*})\) consists of query-relevant nodes \(V^{*}\) and edges \(E^{*}\). \({\boldsymbol{h}}_{\text{in}}\) is obtained from \(S^{*}\) by the graph encoder \(\text{GNN}_{\text{GE}}\) as follows: \[\begin{gather} {\boldsymbol{h}}_{\text{in}}=\text{GNN}_{\text{GE}}(S^{*};{\boldsymbol{\theta}}_{\text{GE}})\in \mathbb{R}^{N\times d_{g}}, \end{gather}\] where \({\boldsymbol{\theta}}_{\text{GE}}\) is parameter of \(\text{GNN}_{\text{GE}}\), \(d_{g}\) is dimension of each output node feature, and \(N=|V^{*}|\).
Node sampler estimates key nodes using \({\boldsymbol{h}}_{\text{in}}\) from the graph encoder for adapting masks on auxiliary nodes. The node sampler processes \({\boldsymbol{h}}_{\text{in}}\) through a Multi-Head Attention (MHA) network, a linear layer, and a softmax activation, to obtain nodes probability scores \({\boldsymbol{p}}_{\text{NS}}\) as follows: \[\begin{gather} {\boldsymbol{z}}_{\text{NS}}=\text{MHA}_{\text{NS}}({\boldsymbol{h}}_{\text{in}};{\boldsymbol{\theta}}_{\text{NS}}^{\text{MHA}}) \in \mathbb{R}^{N \times d_{g}}, \\ {\boldsymbol{p}}_{\text{NS}}=\text{Softmax}(\text{Linear}({\boldsymbol{z}}_{\text{NS}};{\boldsymbol{\theta}}_{\text{NS}}^{\text{Linear}})) \in [0,1]^{N\times 1}, \label{eq:node95sampler} \end{gather}\tag{1}\] where \({\boldsymbol{\theta}}_{\text{NS}}=\{{\boldsymbol{\theta}}_{\text{NS}}^{\text{MHA}},{\boldsymbol{\theta}}_{\text{NS}}^{\text{Linear}}\}\) is the parameter of this module. We sample \(N_{\text{key}}\) nodes based on a categorical distribution defined by \({\boldsymbol{p}}_{\text{NS}}\), where we decide \(N_{\text{key}}\) by the sampling rate \(\rho\) as \(N_{\text{key}}=\lceil \rho N\rceil\). Hereafter, we denote the sampled key nodes as \(I_{\text{key}}\) and auxiliary nodes as \(I_{\text{aux}}(=V^{*}\backslash I_{\text{key}})\). Based on \(I_{\text{key}}\), we extract key node features \({\boldsymbol{h}}_{\text{key}}\in \mathbb{R}^{k \times d_{g}}\) from \({\boldsymbol{h}}_{\text{in}}\) and input it to our concept encoder-decoder module.
Concept Encoder-Decoder consists of a concept encoder \(\text{MHA}_{\text{CE}}\) and a concept decoder \(\text{MHA}_{\text{CD}}\). \(\text{MHA}_{\text{CE}}\) encode the input \({\boldsymbol{h}}_{\text{key}}\) into the latent representation \({\boldsymbol{z}}_{\text{key}}\) as follows: \[\begin{gather} {\boldsymbol{z}}_{\text{key}}=\text{MHA}_{\text{CE}}({\boldsymbol{h}}_{\text{key}}+ \text{PE}({\boldsymbol{h}}_{\text{key}});{\boldsymbol{\theta}}_{\text{CE}}) \in \mathbb{R}^{k \times d_{g}}, \; \end{gather}\] where \(\text{PE}(\cdot)\) represents positional encoding as defined by [68], and \({\boldsymbol{\theta}}_{\text{CE}}\) is the parameter of \(\text{MHA}_{\text{CE}}\). \({\boldsymbol{z}}_{\text{key}}\) is combined with \({\boldsymbol{z}}_{\text{aux}}\), placeholder vectors for unsampled auxiliary nodes with values copied from \({\boldsymbol{h}}_{\text{in}}\), as in [69]. Let \({\boldsymbol{z}}\in \mathbb{R}^{N\times d_g}\) be the combined node features, which maintains \({\boldsymbol{h}}_{\text{in}}\)’s original node position. \(\text{MHA}_{\text{CD}}\) decodes \({\boldsymbol{h}}_{\text{out}}\) with the positional encoding \(\text{PE}\) as follows: \[\begin{gather} {\boldsymbol{h}}_{\text{out}}=\text{MHA}_{\text{CD}}({\boldsymbol{z}}+\text{PE}({\boldsymbol{z}});{\boldsymbol{\theta}}_{\text{CD}}) \in \mathbb{R}^{N \times d_{g}}, \end{gather}\] where \({\boldsymbol{\theta}}_{\text{CD}}\) is the parameter of \(\text{MHA}_{\text{CD}}\). \({\boldsymbol{h}}_{\text{out}}\) is the output of AGE, which we train to predict \({\boldsymbol{h}}_{\text{target}}\), an embedding obtained from all nodes through the target encoder.
Target Encoder is applied to obtain a prediction target for the previous module in a semantic space (JEPA), as it works more robustly than in the input space (GA) [57], [58], [70], [71]. The target encoder \(\text{MHA}_{\text{TE}}\) projects \({\boldsymbol{h}}_{\text{in}}\) to a target embedding \({\boldsymbol{h}}_{\text{target}}\) as follows: \[\begin{gather} {\boldsymbol{h}}_{\text{target}}=\text{MHA}_{\text{TE}}({\boldsymbol{h}}_{\text{in}}+\text{PE}({\boldsymbol{h}}_{\text{in}});{\boldsymbol{\theta}}_{\text{TE}}) \in \mathbb{R}^{N \times d_{g}}, \end{gather}\] where \({\boldsymbol{\theta}}_{\text{TE}}\) is the parameter of \(\text{MHA}_{\text{TE}}\). We train \(\text{MHA}_{\text{TE}}\) with downstream tasks, optimizing it to produce embeddings that directly contribute to the task. \(\text{MHA}_{\text{CE}}\) and \(\text{MHA}_{\text{CD}}\) are trained in parallel with \(\text{MHA}_{\text{TE}}\), with the target encoder learning graph representations for LLMs and the concept encoder-decoder learn to exploit key-node representations for mimic the target encoder’s auxiliary node representations. Inferring masked auxiliary nodes condenses relational concepts between key and auxiliary nodes into \({\boldsymbol{z}}_{\text{key}}\) (and thus \({\boldsymbol{h}}_{\text{out}}\) in the downstream). This JEPA-derived mechanism would be synergetic with GraphRAG as the previous studies suffers from embedding graph’s structured relationships efficiently [37].
Graph-structure-based Aggregator \(\text{GNN}_{\text{GSA}}\) projects \({\boldsymbol{h}}_{\text{out}}\) to a single token \(\bar{{\boldsymbol{h}}}_{g}\). As the graph encoder module, we followed [12] for this module. It aggregates \({\boldsymbol{h}}_{\text{out}}\) referring \(E^{*}\), the edge connections of the original subgraph \(S^{*}\). Where \(\text{POOL}\) is mean pooling and \(d_{l}\) is the input dimension of the target layer. The projector \(\text{MLP}_{\text{Proj}}\) adjusts the aggregated embeddings to fit the LLM’s input dimension. \[\begin{gather} {\boldsymbol{h}}_{g} = \text{POOL}(\text{GNN}_{\text{GSA}}({\boldsymbol{h}}_{\text{out}};{\boldsymbol{\theta}}_{\text{GSA}})) \in \mathbb{R}^{d_{g}}, \; \; \; \; \bar{{\boldsymbol{h}}}_{g} = \text{MLP}_{\text{Proj}}({\boldsymbol{h}}_{g};{\boldsymbol{\theta}}_{\text{Proj}}) \in \mathbb{R}^{d_{l}}, \end{gather}\]

4.3 Optimization of Adaptive-masking for Graph Embedding↩︎

This subsection describes three loss functions used in our method. In the training phase, we connect \({\boldsymbol{h}}_{\text{target}}\) in the target encoder stream to the LLM and optimize \({\boldsymbol{\theta}}_{\text{TE}}\) with the prompt tuning loss \(L_{\text{PT}}\). During training the target encoder with \(L_{\text{PT}}\), the concept encoder-decoder module is optimized exclusively with \(L_{\text{target}}\), in a JEPA approach. Once entire network has been trained, we connect \({\boldsymbol{h}}_{\text{out}}\) to the downstream LLM rather than \({\boldsymbol{h}}_{\text{target}}\) at the inference phase. One challenge of AGE lies in optimizing \({\boldsymbol{\theta}}_{\text{NS}}\). Optimizing \({\boldsymbol{\theta}}_{\text{NS}}\) using \(L_{\text{target}}\) is difficult due to the non-differentiability of the sampling operation. Therefore, we propose an additional loss function \(L_{\text{NS}}\) for optimizing \({\boldsymbol{\theta}}_{\text{NS}}\).
Prompt tuning loss \(L_{\text{PT}}\) maximizes accuracy of a downstream task. It was originally introduced in [12], and we use the definition as is. We optimize \({\boldsymbol{\theta}}_{\text{TE}}\), \({\boldsymbol{\theta}}_{\text{GSA}}\), and \({\boldsymbol{\theta}}_{\text{Proj}}\) with \(L_{\text{PT}}\). The concrete implementation depends on the benchmark tasks; refer to the original papers for details. Note that we train \({\boldsymbol{\theta}}_{\text{GE}}\) with \(L_{\text{target}}\) rather than \(L_{\text{PT}}\), as the concept encoder-decoder is used in the inference phase and the upstream network should be optimized to that module.
Target loss \(L_{\text{target}}\) optimizes parameters \({\boldsymbol{\theta}}_{\text{GE}}\), \({\boldsymbol{\theta}}_{\text{CE}}\), and \({\boldsymbol{\theta}}_{\text{CD}}\) to maximize embedding reconstruction by minimizing the distance between \({\boldsymbol{h}}_{\text{out}}\) and \({\boldsymbol{h}}_{\text{target}}\) for each auxiliary node indexed by \({I}_{\text{aux}}\) as follows: \[\begin{gather} L_{\text{target}}({\boldsymbol{\theta}}_{\text{GE}},{\boldsymbol{\theta}}_{\text{CE}},{\boldsymbol{\theta}}_{\text{CD}}) = \frac{1}{N_{\text{aux}}} \sum_{i \in I_{\text{aux}}} \parallel {h^i_{\text{out}}} - sg({h^i_{\text{target}}}) \parallel_{2}, \end{gather}\] with \(N_{\text{aux}}\) defined as \(N - N_{\text{key}}\). The objective is to apply knowledge distillation to effectively represent key nodes for reconstructing auxiliary nodes. Furthermore, we apply normalization to enhance stability during the learning process, details are provided in the Appendix Table B.8.
Sampling loss \(L_{\text{NS}}\) optimizes \({\boldsymbol{\theta}}_{\text{NS}}\) using RL-inspired supervision. Regarding the operation as an action, the node sampler as a policy network, and \({\boldsymbol{h}}_{\text{out}}\) as a state, we design \(L_{\text{NS}}\) on \({\boldsymbol{p}}_{\text{NS}}=\{p^1,\ldots,p^N\}\) in Eq. 1 as follows: \[\begin{gather} L_{\text{NS}}({\boldsymbol{\theta}}_{\text{NS}})=-\frac{1}{N_{\text{aux}}}\sum_{i\in I_{\text{aux}}}\left( \log(p^i) \times sg(\parallel h^i_{\text{out}}-h^i_{\text{target}}\parallel_{2})\right). \end{gather}\] The loss is back-propagated only to \({\boldsymbol{\theta}}_{\text{NS}}\). Here, for each node assigned to \(I_{\text{aux}}\), larger \(\parallel h^i_{\text{out}}-h^i_{\text{target}}\parallel_{2}\) increases \(p^i\) more, resulting in pushing such node into \(I_{\text{key}}\). This reflects our aim to classify nodes that are difficult to predict from their surrounding as key nodes. Additional strategies for key node selections and sampling optimization for RL are discussed in the Appendix Table B.6.
Total loss is given as \(L_{\text{PT}} + L_{\text{target}} + L_{\text{NS}}\). Since we designed the optimization process so that each loss optimizes different modules without overlap, our methods does not require weight adjustment between these losses.

4.4 Analysis of Learning Objectives↩︎

To explain the motivation for architecture design and applying a distributed loss to each module, we analyze the learning objective of the \(R{\omega}\) module, which represents expected \(\overline{S^{*}}\) for LLM \(\pi_{\theta}\) on LoRA finetuning as: \[\begin{align} \mathcal{L} &= \underbrace{ -\mathbb{E}_{(S^{*})}\!\left[\log R{\omega}(\overline{S^{*}} \mid S^{*}) \right] }_{\text{Loss of Graph Representation Module}} \\ & \underbrace{ -\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\, \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right] }_{\text{Loss of LLM}} \end{align}\] According to Bayes’ Theorem, given an input \(X\), a target \(Y\), and latent rationales \(Z\), we can sample these latent rationales \(Z\) from the posterior distribution \(P(Z|X, Y)\). This posterior represents the probability of latent \(Z\) given both the input \(X\) and the target \(Y\). To compute the marginal likelihood of obtaining answer \(Y\) given input \(X\), we marginalize over all possible rationales \(Z\): \[\begin{align} P(Y|X) = \sum_{Z \sim P(Z|X,Y)} P(Z, Y|X) =\sum_{Z \sim P(Z|X,Y)} P(Z|X) \cdot P(Y|X, Z) \end{align}\] The equations above show how to compute the marginal likelihood \(P(Y|X)\). Equation (15) with first line makes explicit that \(Z\) is sampled from the posterior distribution \(P(Z|X,Y)\). Second line extend applies the chain rule of probability to decompose \(P(Z, Y|X)\) into two components: \(P(Z|X)\) and \(P(Y|X,Z)\). Following this analysis, we apply it to the learning objective for target representation \(\overline{S^{*}}\) given input \(S^{\*}\), latent \(\mathcal{Z}\) from a posterior \(R{\omega}(\mathcal{Z}|S^{\*}, \overline{S^{\*}})\) that bridges \(S^{\*}\) and \(\overline{S^{\*}}\). The marginal likelihood of \(\overline{S^{\*}}\) given \(S^{\*}\) is: \[\begin{align} R\omega(\overline{S^*}|S^*) &=\sum_{Z\sim R\omega(Z|S^*,\overline{S^*})}R\omega(Z,\overline{S^*}|S^*) \\ & =\sum_{Z\sim R\omega(Z|S^*,\overline{S^*})}R\omega(Z|S^*)\cdot R\omega(\overline{S^*}|S^*,Z) \end{align} \label{eq:}\tag{2}\] Above analysis shows that learning objective Graph Representation implicitly learns to identify the latent \(Z\) and map it to the expected \(\overline{S^{*}}\) for LLM: \[\begin{align} -\mathbb{E}\!\left[\log_{R{\omega}}(\overline{S^{*}} \mid S^{*}) \right] &= \underbrace{ \mathbb{E}\!\left[ \log R{\omega}(\mathcal{Z} \mid S^{*},\overline{S^{*}}) \right] }_{\text{Loss of Latent Identification}} \; \; \\ & \underbrace{ -\mathbb{E}\!\left[ \log R{\omega}(\mathcal{Z} \mid S^{*}) \cdot R{\omega}(\overline{S^{*}} \mid S^{*},\mathcal{Z})) \right] }_{\text{Loss of Representation}}\\ \end{align} \label{ysoahcwg}\tag{3}\] Instead of using a single model for both latent \(\mathcal{Z}\) identification and \(\overline{S^{*}}\) representation learning. We separate the learning into \(Sampler_{\theta}\) for latent identification by minimizing reconstruction loss and \(Encoder_{\theta}\)-\(Decoder_{\theta}\) for representation as: \[\begin{align} -\mathbb{E}\!\left[\log_{R{\omega}}(\overline{S^{*}} \mid S^{*}) \right] &\approx \underbrace{ -\mathbb{E}_{(S^{*},\overline{S^{*}})}\!\left[ V_{key} \in \mathcal{Z} \sim \log\;Sampler_\theta (\mathcal{Z} \mid S^{*},\overline{S^{*}}) \right] }_{\text{Loss of Node Sampling}} \\ &\underbrace{ -\mathbb{E}_{(S^{*})} \!\left[ \log Encoder_{\theta} (\mathcal{Z} \mid V_{key}) \cdot Decoder_{\theta} (\overline{S^{*}} \mid \Delta_{V_{masked}},\mathcal{Z}) \right] }_{\text{Loss of Encoder-Decoder}} \end{align} \label{tblnqdke}\tag{4}\] By separating the learning processes, the target encoder learns the representation directly, while the encoder-decoder learns to reconstruct this representation through Evidence Lower Bound (ELBO) optimization. Specifically, the node sampler learns to extrapolate \(V_{key} \in \mathcal{Z}\), making static sampling illogical. Therefore, our node sampler with encoder-decoder architecture and explicit loss distribution yields efficient learning signals, faster convergence, and improved graph representations. To support our analysis, we provide empirical comparisons of sampling strategies in Appendix Figure B.3, B.4 and analyze the stability of the target encoder teacher module for the encoder-decoder in Appendix Table B.8. In the prompt tuning setting, given \(r\) as the reasoning trajectory, the learning objective for frozen LLMs with graph representation model \(R_{\omega}\) is: \[\begin{align} \mathcal{L} &= \underbrace{ \mathbb{E}\!\left[ \log R_{\omega}(\mathcal{Z} \mid S^{*},\overline{S^{*}}) \right] }_{\text{Loss of Latent Identification}} \; \; \underbrace{ -\mathbb{E}\!\left[ \log R_{\omega}(\mathcal{Z} \mid S^{*}) \cdot R_{\omega}(\overline{S^{*}} \mid S^{*},\mathcal{Z})) \right] }_{\text{Loss of Representation}}\\ & \underbrace{ \underbrace{ -\mathbb{E}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\right] }_{\text{Loss of Knowledge Recalling}} \; \; \underbrace{ -\mathbb{E}\!\left[\log \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right] }_{\text{Loss of Contextualized Reasoning}} }_{\text{Frozen}} \end{align}\] We observe that \(R_\omega\) explicitly learns to identify the latent \(\mathcal{Z}\) from \(S^*\) for LLM-expected \(\overline{S^*}\). During training with frozen LLM parameters, \(R_\omega\) implicitly captures latent identification \(i\) by satisfying LLM’s expectations: \(Retriever(S^* \mid q, T^*) \; R_\omega(\overline{S^*} \mid S^*, \mathcal{Z}) \approx \pi_\theta(\overline{S^*} \mid q, T^*, i)\), yielding \(\mathcal{Z} \subseteq i\). Therefore, \(R_\omega\) able to learns a subspace of the frozen LLM’s complete latent space through this objective. Throughout this, we argue that leveraging a learned latent space \(\mathcal{Z}\), robustly restructured into the LLM-expected representation \(\overline{S^*}\), can directly improve knowledge recall and indirectly enhance reasoning.

5 Experiments↩︎

Datasets and Evaluation Metrics. Following previous work [12], [15], [41], we conduct experiments on ExplaGraphs [72] is a generative commonsense reasoning dataset, SceneGraphs [73] is a visual question answering dataset. And WebQSP [74], ComplexWebQuestions (CWQ) [75] is a large question-answering dataset derived from Web questions, where all queries can be answered using Freebase, a large collaborative knowledge graph database. We use accuracy as the primary metric for ExplaGraphs and SceneGraphs, datasets focusing on reasoning, following [12], [15], [41]. For WebQSP and CWQ, a dataset with extra-large graphs, we use the Hit@1 metric, as in [21].
Implementation Details. We employed the open-source Llama3.2 (1B and 3B) [76], Llama 2 (7b and 13B) [77] and Llama3.1 (8B) [6] as frozen LLM components. Based on the analysis in 5.2, we set the sampling rate \(\rho=0.3\) (see Appendix 2.1 and 2.5 for cost–benefit trade-offs).

5.1 Main Results↩︎

Table ¿tbl:tab:performance? illustrates our main results, comparing the methods in three settings: Frozen LLM with Graph Embedding: Use a graph embedding technique to tune tokens with given prompt. Frozen LLM with Graph Embedding + PEFT: Apply LoRA [78], a PEFT technique, in combination with the graph embedding technique. LLM with LLM-Retriever: Use LLM for retrieval in addition to the one for inference. Any methods train either LLM. These are reference scores with a larger computational cost.

Among Frozen LLM with Graph Embedding settings (with and without PEFT), AGE consistently improved performance of G-Retriever and AMAR regardless of the backbone LLM models. Without PEFT, Llama3.2-1B with AGE showed the most notable gain against G-Retiever: \(26.72\) percent points increase on ExplaGraphs, while the least gain was observed with Llama3.2-3B on WebQSP, which was \(2.02\) points. This might be due to the extra-large size of knowledge graphs in WebQSP datasets, which include textual knowledge absent at pretraining, and non-parameter retriever struggling to provide critical information for representation. AGE maintains consistent superiority against G-Retriever and shows more gains on retrieval from smaller graphs. By employing a cross-question approach enriched with retrieved elements, GRAG improved performance by \(2.8\) points, AMAR achieved an improvement of \(4.2\) points with its baseline. This suggests that improving embedding module is beneficial, it alone may not be enough to boost performance significantly, and relying on a non-parameter retriever could be limiting. Despite that challenge, when integrated with AMAR, AGE continues to achieve further enhancements the performance. Direct comparison with current LLM with LLM Retriever methods on WebQSP and the larger CWQ shows that AGE AMAR, which applies non-parametric retriever-based approaches, outperforms the proprietary LLM-based retriever ReKnoS [79] on both datasets. AGE AMAR underperforms compared with Paths-over-Graph [80], Plan-on-Graph [81] and DoG [82] on WebQSP, but outperforms them on CWQ, suggesting that AGE AMAR is advantageous on larger datasets, showing substantial potential for future work (as reported in the Appendix Table B.11, further including the more baseline methods.)

5.2 Ablation study↩︎

Performance Comparison of Self-Supervised Learning Architectures. Table 1 presents the performance of concept encoder-decoder with some variations and the baseline of G-Retriever, which demonstrates contribution of proposed technique independently. As a SSL variation, we prepared AGE based on generative architecture (GA) against our choice of JEPA. We also compared AGE with random mask against the proposed learnable node sampler. All the results used Llama3.2 1B as its LLM. Using a GA with a random mask, AGE achieves a performance of \(0.6532\), which is a \(9.37\%\) improvement over the baseline. Next, AGE with a random mask, improving performance by \(15.46\%\) over the baseline. Finally, AGE using JEPA with the learnable node sampler achieves a \(26.72\%\) improvement over the baseline. Based on these experiments, we confirmed that JEPA works better than GA as expected, while node sampler further improves performance with a notable margin. Furthermore, we include an additional study on architectural design choices in the Appendix Table B.7, such as reason we choose different input features for LLMs during training and inference in GraphRAG.

Analysis on the Sampling Rate \(\rho\). Figure 3 illustrates how the sampling rate impacts AGE G-Retriever performance on ExplaGraphs and WebQSP, guiding our hyper-parameter setting. Average retrieved nodes are \(18.21\) and \(5.17\) on WebQSP and ExplaGraphs, respectively. A sampling rate \(\rho=0.3\) gives the best performance on ExplaGraphs with both LLM settings (\(81.4\%\) for Llma3.2-1B and \(92.6\%\) for Llama3.2-3B). The same setting also achieves the best performance on WebQSP for Llama3.2-1B (\(62.5\%\)) and the second best for Llam3.2-3B (\(72.2\%)\), compared to the best score of \(73.5\%\) at \(\rho=0.35\). From these observations, we decided to use \(\rho=0.3\) through the experiments.
Other LLM backbones. A Qwen3.5 result (Table 2) confirms the same trend. We will clarify the LoRA setting. AGE consistently outperforms G-Retriever across all scales and tasks, with especially large gains on ExplaGraphs. The improvements remain strong even at the smaller 0.8B scale, indicating efficiency and scalability. Gains on WebQSP further demonstrate the robustness of the proposed method across different tasks.

0pt 1pt 4.5pt

Weighting Loss. The three losses optimize disjoint parameters, so weighting is unnecessary by design; confirmed also in Table 3. Equal weighting (1:1:1) consistently yields the best performance across datasets and model sizes. Deviating from equal weights degrades performance, with no consistent benefit from emphasizing any single loss. This supports that the three losses operate on disjoint parameters, making explicit weighting unnecessary.

0pt 1pt 9pt

Qualitative Evaluation. To analyze node sampling results, we visualized the node sampling results on two samples from the ExplaGraphs test set in Figure 4. Nodes are colored based on clustered text embeddings to track node-wise feature restructuring through graph relationships. The graph encoder process maintains the clustering structure of text graph embeddings in both G-Retriever (first column) and AGE (second column). In contrast, the concept encoder-decoder module (second column) shuffles the colored nodes, indicating a reorganization of the node-wise embeddings.
The last column displays text entities of some key and auxiliary nodes. Our node sampler is designed to sample entities from specific domains as key nodes. As "saving souls" is sampled as a key node, inferring "missionaries" and "Christians" from it seems easier than the reverse. We observe the same tendency with the key node "work with criminals" and the auxiliary nodes "imprison people" and "public defenders." The last column also shows the target loss for each auxiliary node using the color bar, where 1.0 represents the maximum error in the test set, and 0.0 the least. From the color visualization, we observe that non-isolated key nodes achieve lower errors in auxiliary node prediction, suggesting that relations between key nodes support the prediction. We provides additional qualitative results, including failure cases, in Appendix B.9.

6 Limitation and Conclusion↩︎

Although our method delivers substantial improvements in GraphRAG, several limitations remain. First, we used a fixed sampling rate despite variations in key node density between the graphs. Second, we tested AGE only on GraphRAG tasks, even though it is applicable to other modalities. These limitations suggest areas for future improvement and the potential for broader applications. Third, our approach primarily targets small-scale models, and its effectiveness for large-scale LLMs remains unexplored due to computational constraints. Finally, our method focuses on representing retrieved structured data for LLMs rather than directly addressing graph learning tasks (e.g., node classification or link prediction). In the absence of theoretical guarantees on the benefits of node and link integration, our current scope is mainly limited to KGQA scenarios.
We proposed Adaptive-masking for Graph Embedding (AGE) to improve structured graph embeddings and enhance LLM performance on GraphQA tasks. The method introduced JEPA, a self-supervised learning architecture which enhanced the graph-structure embedding for downstream reasoning tasks. Our node sampler demonstrated its effectiveness in the ablation study, successfully identified key nodes within given graphs. The quantitative results confirmed AGE’s consistent performance gain in GraphRAG tasks while maintaining computational cost. We hope this work contributes to structured knowledge representation for intelligent agents and facilitates cross-modal reasoning through structured perceptual representations.

7 Proof of Concept↩︎

In this section, we motivate AGE’s design for representing retrieved subgraphs for higher-order reasoning skills. We assume this requires on the most common static search single-turn retrieval and LLM for academic simplicity.
Problem definition: Given input query \(q\), we generate chain-of-thoughts response \(y = i \oplus r\) (interleaved domain knowledge \(i\) and reasoning steps \(r\)). We use a static retrieval engine that returns \(T^{*}\) for text knowledge and \(S^{*}\) for subgraph. The learning objective \(\min_{\theta} \mathbb{E}[\mathcal{L}]\) optimizes representation \(S^{*}\) to prioritize reasoning \(r\) over knowledge \(i\).
For the chain-of-thoughts response is defined as \(y = i \oplus r\) where \(y\) is the concatenation of knowledge \(i\) and reasoning \(r\) through three discrete generation processes below.

• Graph Knowledge Retrieval: Given the query \(q\) on a textual graph \(G\) and \(S(G)\) is the set of all subgraphs of \(G\). Retrieval system extracts relevant subgraph \(S^{*} \in S(G)\) and text-modal knowledge \(T^{*}\). Popular retrieval systems select top-k elements by cosine similarity, yielding nodes \(V_k^{*}\) and edges \(E_k^{*}\) considered relevant to the query. A non-optimized retriever may yield corrupted subgraphs \(S^{*}=(V_k^{*},E_k^{*})\), as they may be contain redundant or lack suggestive elements.

• Graph Knowledge Representation: Graph embedding module is trained to represent graph that guides the LLM to produce expected answers. Embedding module \(R_\omega(\cdot)\) learn to represent corrupted subgraph \(S^{*}\) to \(\overline{S^{*}}\) for an LLM\(\Phi\) improve generation.

• Contextualized Reasoning: Given \(q\), \(T^{*}\) and \(\overline{S^{*}}\), LLM synthesizes domain knowledge \(t\) by recall their internal parametric knowledge with external inputs, following the conditional distribution \(i \sim \pi_{\theta}(i|q, T^{*},\overline{S^{*}})\). Then LLM generates reasoning steps \(r\) conditioned on \(q\), \(T^{*}\), \(\overline{S^{*}}\) with the recalled internal knowledge \(i\), adhering to the reasoning distribution \(r \sim \pi_{\theta}(r|q, T^{*},\overline{S^{*}}, i)\).

Here we formally analyze and discuss the subgraph representation learning objectives of both vanilla embedding module and AGE embedding module with LLM generation distribution.

• Subgraph embedding module that employ GNN or Transformer: \[\begin{align} \overline{S^{*}} &= \text{R}_\omega^{\text{GNN}}(S^{*}) = \big\{\, \overline{V}_i \,\big\}_{i \in V^{*}} \\ \overline{V}_i &= \sum_{j \in \mathcal{N}(i)} \alpha_{ij}\, W_V V_j, \end{align}\] Here, \(\mathcal{N}(i)\) denotes the local neighborhood of node \(i\) (e.g., \(\mathcal{N}(i)=\{j \mid (i,j)\in E^{*}\}\)). \[\begin{align} \overline{S^{*}} &= \text{R}_{\omega}^{\mathrm{Transformer}}(S^{*}) = \big\{\, \overline{V}_i \,\big\}_{i \in V^{*}} \\ \overline{V}_i &= \sum_{j \in V^{*}} \alpha_{ij}\, W_V V_j, \qquad \alpha_{ij} = \mathrm{softmax}_{j \in V^{*}} \!\Big( \tfrac{(W_Q V_i)^\top (W_K V_j)}{\sqrt{d_k}} \Big), \end{align}\] where \(W_Q,W_K,W_V\) are learnable matrices, \(d_k\) is the key dimension. These formulations use attention-weighted aggregation as node embeddings. Static retrieval from graphs with high-order structural patterns (motifs/role patterns) produces corrupted subgraphs that contain redundant nodes or miss critical elements. Without explicit structural constraints, training weighted-sum aggregation through supervised or semi-supervised learning separates signal from noise, may lead to diluted node representations within the embedding space. In a frozen state, LLMs may fail to capture relationships in the embedding space due to their struggle to handle diluted node representations.

• Subgraph embedding module that employ RL-guide mask-based SSL: \[\begin{align} \overline{S^{*}} = R{\omega}^{\mathrm{AGE}}(S^{*}) = \big\{\, \overline{V}_i \,\big\}_{i \in V^{*}} =Decoder_{\phi}\big(\Delta_{V_{masked}}, Encoder_{\theta}(V_{key})\big) \\ = V_{key} \sim \text{Sampler}_\theta(V_{key} \mid V^{*}) \!\Bigg[ \sum_{j \in V^{*}} \alpha_{ij} W_V \bigg[ \Big[\sum_{j \in V^{key}} \alpha_{ij}\, W_V V_j \Big]; \Delta_{V_{masked}} \bigg] \!\Bigg] \end{align}\] \(\Delta_{V_{masked}}\) denotes the auxiliary masked node features used as decoder input, while \(V_{key}\) denotes the key node features used as encoder input and \([;]\) is the concatenation operation. \[\begin{align} \begin{cases} \Delta_{V_{masked}} = Mask(V_{j}) & \text{if } a_{j} = 1 \\ V_{key} = Visible(V_{j}) & \text{if } a_{j} = 0 \end{cases} \quad \sum_{j \in V^{*}} \text{Sampler}_\theta(a_{j} \mid V_j) , \end{align}\] where \(V_{key}\) represents key node features for encoder input, where \(\text{Mask}(\cdot)\) masks features, \(\text{Visible}(\cdot)\) preserves them, and \(a_j \in \{0,1\}\) denotes the binary RL action (\(1\)=mask, \(0\)=keep) for node \(j\).

  By employing mask-based SSL alongside a reinforcement learning framework, AGE learns structural dependency node representations in the embedding space through reconstruction objectives. The reinforcement learning framework is used to estimate which nodes are critical for preserving graph structure and semantic information. Then, the estimated nodes are applied to guide mask-based SSL to reconstruct that provide structural constraints in the embedding space, enabling LLMs to better separate signals and capture relationships within it.

• The joint generation distribution of LLMs is: \[\begin{align} \pi_{\theta}(y \mid q, T^{*},\overline{S^{*}}) &= \pi_{\theta}(i \oplus r \mid q, T^{*},\overline{S^{*}}) \\ &= \underbrace{\pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})}_{\text{Knowledge Recalling}} \cdot \underbrace{\pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)}_{\text{Contextualized Reasoning}} , \end{align} \label{eq:joint95generation}\tag{5}\]

• The loss function optimizes both knowledge integration and contextualized reasoning: \[\begin{align} \mathcal{L} &= -\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\, \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right] \\ &= -\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\right] \; \; -\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right] \, , \end{align} \label{eq:loss95generation}\tag{6}\]

Through the lens of multi-task learning, we compare above equations and from two perspectives:

• Retrieved Subgraph Representation on LLMs’ generation distribution. Based on equation equation 5 , LLMs decompose response generation into knowledge recall and contextualized reasoning. In the frozen state, diluted subgraph representations trigger a cascade—weak knowledge recall causes flawed reasoning. This limitation creates a bottleneck that degrades response quality.

• Retrieved Subgraph Representation with Parameter-Efficient Fine-Tuning.

  Following previous research [84], we assume the retrieved representation is explicitly, we have  equation 6 as: \[\begin{align} \mathcal{L} &= \underbrace{-\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\right]}_{\boldsymbol{Loss of Integration}}\downarrow \; \; \underbrace{-\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right]}_{\text{Loss of Reasoning}}\uparrow \, , \end{align}\] The arrows indicate the loss function shifts to reasoning. That means the loss term shifts from knowledge identification to integration, that \(\pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\) has already reached "application" levels of retrieved graph knowledge. Therefore, explicit subgraph representation aids knowledge integration beyond mere identification during fine-tuning.

  Conversely, the retrieved subgraph representation is diluted, we have  equation 6 as: \[\begin{align} \mathcal{L} &= \underbrace{-\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\right]}_{\boldsymbol{Loss of Identification}}\uparrow \; \; \underbrace{-\mathbb{E}_{(q, T^{*},\overline{S^{*}})}\!\left[\log \pi_{\theta}(r \mid q, T^{*},\overline{S^{*}}, i)\right]}_{\text{Loss of Reasoning}}\downarrow \, , \end{align}\] The arrows indicate that the loss function prioritizes identification. This reveals inefficient knowledge use: the model identifies patterns directly, treating retrieved subgraphs as training data. This creates a trade-off: instead of learning to apply retrieved subgraphs in reasoning, \(\pi_{\theta}(i \mid q, T^{*},\overline{S^{*}})\) prioritizes identifying graph structures over reasoning. Poor retrieved subgraph representation implicitly hinders reasoning capability development, forcing resources into identification tasks that better representations would cover.

8 Additional Experimental Details↩︎

8.1 Implementation Settings (AGE G-Retriever)↩︎

When integrated with G-Retriever, we consistently use the AdamW [85] optimizer and set the initial learning rate at \(1e-4\), with a weight decay of \(0.05\). Following the baseline work [12], we set learning rate decays with a half-cycle cosine decay after the warm-up period. To avoid overfitting, we implement early stopping with a patience of \(3\) epochs. The experiments used 2 NVIDIA 2080Ti-11G or 2 NVIDIA A100-80G GPUs.
GNN. We use Graph Transformer as the GNN backbone applied in the Graph Encoder and Graph Structure Based Aggregator. Similar to previous approaches [12], our settings for its employ \(2\) layers, each with \(4\) attention heads, and a hidden dimension size of \(1024\).
LLM. We use the open-source Llama3.2 1B, 3B [76], and Llama3.1 8B as the LLM backbone. When LoRA [78] is applied with the LLM, the LoRA scaling factor hyperparameter is set to \(16\). Following previous work [12], we configure the LLM with a maximum input text length of \(512\) and a maximum number of new tokens to generate of \(32\).
Subgraph Construction. We follow previous approaches [12] that select the top \(k\) nodes and edges through subgraph construction by setting \(k\) to \(3\) for SceneGraphs dataset. For WebQSP dataset, \(k = 3\) for nodes and \(k = 5\) for edges. For the ExplaGraphs dataset, the entire graph fits within the LLM’s context window. Thus, setting \(k\) to \(0\) for retrieves the original graph without modification.

8.2 Implementation Settings (AGE AMAR)↩︎

When integrated with AMAR[40], to fairly compare we keep the training settings of AMAR, setting the retrieved data to \(100\) on WebQSP, Soft prompt length to 7, Beam search number to 8, and Max new tokens to 256 for the WebQSP dataset. For the CWQ dataset, we set the retrieved data to 4, Soft prompt length to 16, Beam search number to 15, and Max new tokens to 256. With the Llama2 [77] is trained with LoRA learning rate \(5e-5\) scaling factor hyperparameter is set to \(32\).

8.3 The choice of AGE architecture↩︎

8.3.1 The choice of graph structure extractor architecture on AGE G-Retriever↩︎

Figure 5 shows the architectures with graph encoder only, graph-structure-based aggregator only, and both combined. The best-performing architecture is the combination of graph encoder and graph-structure-based aggregator. It achieves a Hit@1 score of \(73.46\)% on the WebQSP dataset, improving upon \(71.12\)% with Graph Encoder and \(72.44\)% with graph-structure-based aggregator.

8.3.2 The Choice of GNN on AGE G-Retriever↩︎

In Table 4, our investigation extends to existing popular GNNs employed as both graph encoders and graph-structure-based Aggregator, including the Graph Convolutional Network (GCN) [86], Graph Attention Network (GAT) [87] and Graph Transformer [88].
This illustrates the performance comparison of these GNNs on the WebQSP and ExplaGraphs datasets. On the WebQSP dataset, the GCN, GAT, and Graph Transformer achieve Hit@1 scores of \(56.75\), \(61.42\), and \(62.53\), respectively. In the ExplaGraphs dataset, the Graph Transformer achieves the highest accuracy of \(0.8501\), followed by the GCN with an accuracy of \(0.8321\), and the GAT trailing slightly at \(0.8212\).
These findings emphasize the critical role of selecting an appropriate GNN architecture tailored to the unique properties and demands of each dataset. To maintain performance across various datasets, we choose the Graph Transformer for all experiments.

8.3.3 The design of AGE with Generation Architecture↩︎

As shown in Figure 6, we provide more details of AGE with the randomly masked generative architecture (GA) introduced in Section 4.3. Designing AGE with GA aims to complement the input node to enhance the embedding of the graph-structure-based aggregator during the inference stage. To do this, we train the encoder-decoder with input nodes masked by a random mask at a masking ratio of \(70\)%.Then, the encoder is trained to embed unmasked nodes, and the decoder reconstructs masked nodes through the target loss. On the other hand, prompt tuning loss is used to train the graph-structure-based aggregator and MLP for referring to edges \(E^{*}\), and the MLP adjusts the aggregated embeddings to align with the LLM input dimension. During the inference stage, random masking is disabled. All input nodes are fed into the encoder-decoder to reconstruct the input for the graph-structure-based aggregator.

8.3.4 Analysis on the Layer Number of the Graph Encoder \(GNN_{ge}\) on AGE G-Retriever↩︎

Compared to G-Retriever, AGE’s inference path has additional modules. While this might increase processing time, Table 5 indicates otherwise. For G-Retrievers, a deeper \(GNN_{ge}\) performs better. In contrast, AGE performs better with fewer layers, as the added modules effectively substitute for reduced GNN layers. As a result, AGE achieves superior performance while maintaining the training time of the baseline method. We provide further analysis on computational complexity in Appendix.

8.3.5 Comparison with Deeper GNNs↩︎

Table 6 compares the number of GNN layers and the performance of G-Retriever and Adaptive-masking for Graph Embedding models. G-Retriever with 20 layers is prepared as a model whose computational cost (GFLOPs) is similar to AGE with 2 layers.
Applying the G-Retriever with \(20\) layers largely improves performance. However, AGE still outperforms G-Retriever by approximately \(10\) points when using Llama 3.2 1B and \(6\) points when using Llama 3.2 3B, demonstrating AGE’s superior performance.

8.3.6 The choice of concept decoder and node sampler architecture.↩︎

4pt

Table ¿tbl:table:Decoder? illustrates our analysis of model performance across various concept decoder depths. We increased the decoder depth from 1 block to 4 blocks, thereby increasing the parameters from \(81\)M to \(106\)M. Despite this increase, the performance decreased, with scores dropping from \(62.5\) to \(57.4\) on WebQSP. The best performance is achieved with a decoder depth of 1 on both ExplaGraphs and WebQSP. Therefore, we choose a single transformer block to maintain the performance of the concept decoder in this work.
As shown in Table [tbl:table:Maskselector], we investigate different network architectures for the node sampler design. Increasing the number of transformer blocks leads to marginal gains in performance on the WebQSP dataset, although it requires more memory. To maintain computational and performance efficiency, we selected a single transformer block with a hidden dimension of 1024 for the node sampler in all subsequent experiments.

8.3.7 The study on key nodes sampling↩︎

Figure 7 and [fig:keynode2] illustrates how the static strategy node sampler and RL-based node sampler performance in various sampling rate on ExplaGraphs and WebQSP, guiding our node sampler architecture setting.
When using static PageRank [89] and Degree Centrality strategies for node sampling, higher sampling rates tend to better performance. However, this suggests that the key nodes that identified by these static methods are not sufficiently impactful. Lead to the Concept Encoder-Decoder needs a larger set of key nodes to effectively embed the graph, which then helps guide the LLM to produce the desired answers.
In contrast, RL-based node samplers can achieve high performance with lower sampling rates. This indicates that the key nodes chosen by RL-based methods more effectively support the Concept Encoder-Decoder, boosting the quality of the graph embedding. As a result, the LLM can produce the expected answers with fewer key nodes involved in the guidance process.

8.3.8 The study on transferability↩︎

Table 7 show the transferability of AGE when interacted with G-Retriever. AGE support G-Retriever to strong transferability to transfer learned graph embedding encoding capabilities across datasets. When trained on a large dataset, AGE can enhance generation on a smaller dataset using the trained model. Notably, AGE trained on WebQSP on ExplaGraphs with Llama 3.2 3B outperforms transferability of GRAG.

8.3.9 The study on number of retrieval↩︎

As illustrated in Figure8, AMAR was designed to address the challenge posed by excessively long retrieved data inputs and to leverage rich information more effectively. However, when the volume of retrieved data is relatively small, AMAR’s performance shows minimal improvement, indicating that the recalled information is insufficient. In such cases, AGE is capable of mapping retrieved data to useful embeddings, eading to significant performance improvements, with scores of \(86.8\) for 5 retrievals and \(87.0\) for 10 retrievals.
Conversely, when large amounts of data are retrieved, the accompanying noise complicates the ability of LLMs to identify and prioritize the most relevant information. AGE consistently maintains its performance with minimal variation, highlighting the robustness as \(86.5\). Moreover,to fairly compare with AMAR, we choose 100 retrievals with an Hit1@ of \(86.5\).

8.3.10 Impart of AGE with LoRA Finetuning performance↩︎

We conduct extensive experiments to compare the trend of training loss with G-Retriever baselines, as illustrated in Figure9 and 10. AGE provides structural constraints in the embedding space through mask-based SSL with reconstruction objective, enabling LLMs to better separate signals and capture relationships within it, leading to increased convergence rate and lower loss observed in the initial training stage compared with the G-Retriever.

8.3.11 The choice of RL method for Node Sampler↩︎

Table 8 show Both Gumbel-Softmax and REINFORCE demonstrate strong performance on the ExplaGraphs dataset, outperforming the Straight-Through Estimator. On the WebQSP dataset, REINFORCE leads slightly, indicating it may be the most effective method among the three tested for this task. From these observations, we decided to use REINFROCE through the experiments.

8.3.12 The chose different input features for LLMs during training and inference↩︎

Table 9 show AGE with \(h_{out}\) as LLMs input in inference state outperforms AGE with \(h_{target}\) as input on both datasets. This indicates that representing nodes using a concept encoder-decoder \(h_{out}\) is more effective than the target encoder in downstream LLMs input tasks. Based on these observations, we concluded that designing the connection of \(h_{target}\) during training and \(h_{out}\) during inference to the downstream LLM not internalizes the learning-inference mismatch. Instead, it allows the student model has already surpassed the performance of the teacher model, allowing for a more robust representation.

8.3.13 The stability of Target encoder↩︎

Due to differences in node representations during the training inference stage, using identical parameters for both the concept encoder and the target encoder helps prevent distributional shifts. We apply two popular techniques to enhance the stability provided by the target encoder for the concept encoder-decoder during the training stage. EMA weights are defined as an exponential moving average of the encoder weights, and normalization is applied to enhance stability during the learning process. Normalization ensures consistent activation distributions and reduces internal covariate shift. Table 10 shows that MSE performs better with L1, and normalization alone achieved the highest score. These results indicate that normalization consistently improves encoder stability and performance, and adding EMA offers further enhancements.

8.3.14 The landscape of existing KGQA methods↩︎

Figure 11 illustrates the spectrum of current Knowledge Graph Question Answering (KGQA) approaches regarding KG retrieval and reasoning capabilities. Graph Neural Network (GNN)-based methods, including NSM [90], ReaRev [91], and G-Retriever [12], perform reasoning on retrieved dense subgraphs by utilizing the GNN to embed graph structures.
Recent LLM-based methods leverage the power of LLMs for both retrieval and reasoning. ToG [20] uses the LLM to retrieve relevant facts hop-by-hop. RoG [21] uses the LLM to generate plausible relation paths which are then mapped on the KG to retrieve the relevant information. However, the frequent calls to the LLM significantly increase the training and inference costs.
In this work, we improve LLM reasoning by enhancing the graph embedding of the GNN method with RL-inspired supervision integrated into the SSL framework. This improves the performance of the non-parametric retriever to levels comparable to those of LLM-based retrievers.

8.4 Additional Qualitative Evaluation↩︎

We provide additional visualizations in Figures 13, 14 on the ExplaGraphss dataset and Figure 15 on WebQSP dataset.
In the first row of Figure 13, we consider the addition of node text information and its visualization. It is easier to infer the auxiliary node embeddings "Payday loans" and "For the disadvantaged" from an key node embeddings "Provide assistance". Conversely, it is more challenging to infer the key node embeddings "Provide assistance" from the auxiliary node embeddings "help society" and "available".
Similarly, it is easier to infer "Bullying, However they like, Banned" mask node embeddings from a "Expensive clothes, Students" key node embeddings. Conversely, it is more challenging to infer a "Expensive clothes, Students" mask node embeddings from a "Bullying, However they like, Banned" key node embeddings.
In the second row of Figure 13, when considering the addition of node text information and its visualization, it is easier to infer the auxiliary node embeddings "Motivation" and "Students work harder" from the key node embedding "Student loans". Conversely, it is more challenging to infer a "Student loans" auxiliary node embeddings from "Motivation" and "Students work harder" key node embeddings. Similar things are shown in Figure 14 and Figure 15.

Failure Case Analysis: Furthermore, we provide a failure case on the WebQSP dataset where AGE was trained with LLaMA 3.2 1B using a sampling rate of \(0.3\). In this case, the query is "What is the zip code of Seattle, Washington?". Based on the query, the retriever is provided with the node "98175," which is one of the true answers. However, the response of the LLM lacks this node. To analyze this, we visualized the node sampling results from the concept encoder’s input and the concept decoder’s output, as shown in Figure 12. In the first column (top row), similar to the above visualization, the graph encoder process maintains the clustering structure of text graph embeddings to provide input to the concept encoder. Additionally, the output of the concept encoder-decoder module (bottom row) shuffles the colored nodes. In the second column, the target loss of each auxiliary node is represented using a color bar, where 1.0 indicates the maximum error in the test set and 0.0 the minimum error. In this column, the top left side nodes "98194, 98117, 98164,..." include low target loss auxiliary nodes, and the key nodes are the true answers. In parallel, the bottom right side node "98175" is an auxiliary node with a high target loss. This may be why the LLM omits this node in its response, and adjusting the trainable sampling rate could be a solution.

8.5 Discussion on the Complexity↩︎

8.5.1 Training Computational Resources on AGE G-Retriever↩︎

Following the previous G-Retriever [12] method, we utilized the same two A100 GPUs, each with 80GB of memory, and conducted tests on the Llama3-8b, Llama3.1-1B, and Llama3.1-3B on WebQSP datasets. Our experiments had a training batch size of 16 and an evaluation batch size of 32, yielding the following results in Table 11 for training cost and Table 12 for validation speed.
The Table 11 shows the training speed and performance of AGE on the WebQSP dataset. The PEFT setting, without the graph RAG component, takes 18.7 min/epoch through prompt tuning and 19.0 min/epoch when applied with LoRA. Subsequently, the G-Retriever approach via graph RAG reduces graph size and speeds up training time.
By enhancing the embedding module on the graph RAG component, AGE with Llama3.1 8B achieves a higher Hit@1 of \(78.25\) in \(6.4\) minutes per epoch. In the tuned LLM setting, AGE with Llama3.1 8B and LoRA achieves a Hit@1 of \(80.34\) in \(7.3\) minutes per epoch. These results highlight that AGE with Llama3.2 3B outperforms G-Retriever with Llama2 7B, achieving better performance without longer training time.

8.5.2 Inference Computational Resources on AGE G-Retriever↩︎

Table 12 presents the validation speed and performance of various AGE configurations on the WebQSP dataset. Among the AGE models, Llama 3.2 3B model offers a balanced performance with a Hit@1 and an inference speed of \(87.6\) tokens per second. The AGE with Llama 3.2 1B achieves a significantly higher inference speed of \(148.5\) tokens per second while maintaining a lower Hit@1. This increased speed can be attributed to the reduced number of parameters in the 1B model, which allows for faster computation and more efficient processing, albeit at the expense of some accuracy.
These results indicate that while higher parameter models like AGE+Llama 3.1 8B provide superior accuracy, lower parameter models such as AGE+Llama 3.2 1B offer significantly increased processing speeds, supporting diverse application requirements.

8.5.3 Comparison with trainable retriever methods↩︎

AGE, utilizing a non-parametric retriever, achieves accuracy levels comparable to state-of-the-art models that employ trainable parametric retrievers. As shown in Table ¿tbl:tab:trainableretriever?, AGE (Llama3.1 8B-LoRA with a non-parametric retriever) attains a Hit@1 score of \(80.3\%\) on the WebQSP dataset, closely approaching DualR (ChatGPT with a parametric retriever), which achieves \(82.8\%\). This demonstrates that AGE effectively bridges the performance gap between non-parametric and parametric retriever models, achieving high accuracy without the additional complexity and training overhead associated with parametric retrievers. This performance notably surpasses other models employing non-parametric retrievers, such as GraphToken (Llama2-7B) with \(57.1\%\) and G-Retriever (Llama2-7B-LoRA) with \(70.2\%\).
The substantial increase in accuracy demonstrates that AGE enhances reasoning capabilities without relying on trainable parametric retrievers. This positions AGE as a leading approach within non-parametric retriever frameworks, closing the performance gap with models that utilize more complex and resource-intensive trainable retrievers. AGE can be deployed to train and perform inference on two RTX 2080Ti 11GB GPUs or one A100 80GB GPU.

References↩︎