Technical Summary: Query-Aware Learnable Graph Pooling Tokens as Prompt for Large Language Models

1. Problem Statement

The paper addresses the challenge of integrating graph-structured data with Large Language Models (LLMs) for Text-Attributed Graph Question Answering (Graph QA) tasks. While LLMs possess strong reasoning capabilities, effectively incorporating complex graph structures remains difficult. Existing approaches generally fall into two categories, both of which face significant limitations:

• Node-Level Projection: This method passes all node embeddings generated by a Graph Neural Network (GNN) to the LLM as individual tokens. While this preserves fine-grained structural information, it suffers from severe scalability issues. As graphs grow exponentially in size, the prompt length exceeds the context window limits of LLMs, or the computational cost becomes prohibitive.
• Graph-Level Projection: This method compresses all node information into a single vector (e.g., via mean pooling) before passing it to the LLM. While this solves the scalability issue, it inevitably leads to information loss, as complex relational contexts are condensed into a single representation, making it difficult for the LLM to access specific details.

Furthermore, existing methods often employ Late Query Fusion, where the query and graph are encoded independently and combined only after the graph representation is fully constructed. This approach may fail to effectively utilize query-specific context during the critical graph encoding phase.

2. Methodology

The authors propose a novel framework named Learnable Graph Pooling Token (LGPT), integrated with an Early Query Fusion strategy. The method operates without fine-tuning the LLM itself, focusing instead on optimizing the graph encoder and prompt generation.

2.1 Early Query Fusion

Instead of encoding the graph and query separately, the proposed method integrates query context before constructing the graph representation.

• The query $q$ is embedded into the graph embedding space as a virtual query node ($n_q$) using a text encoder.
• This virtual node connects to all nodes in the graph, allowing message passing via a GNN ($GNN_{query}$) to generate query-aware node embeddings.
• A subsequent GNN ($GNN_{graph}$) encodes the relational information of this modified graph structure. This ensures that the graph representation is inherently aware of the query context, rather than just being a static structure.

2.2 Learnable Graph Pooling Tokens (LGPT)

To bridge the gap between node-level and graph-level projections, the authors introduce $n$ learnable parameters (tokens) that act as a bridge between the graph and the LLM.

• Mechanism: The graph is processed by a pooling GNN ($GNN_{pool}$) where all original nodes are fully connected to $n$ learnable tokens. Through message passing, these tokens aggregate information from the graph.
• Output: Instead of a single vector or $k$ node vectors (where $k$ is the number of nodes), the method outputs $n$ tokens ($\{g'_1, \dots, g'_n\}$).
• Projection: These tokens are passed through a projection layer (MLP) to match the LLM's embedding dimension and are concatenated with the textual graph prompt and the query embedding.
• Conceptual Basis: The approach conceptually combines soft clustering (where nodes can belong to multiple clusters, similar to hierarchical pooling) and self-attention mechanisms (similar to Graph Transformers). By using multiple tokens rather than a single vector, it mitigates information loss while maintaining a fixed, manageable prompt length.

2.3 Training Objective

The LLM and the text embedding layer are kept frozen. The training objective is to optimize the learnable parameters $\psi$ of the graph encoder components ($GNN_{query}$, $GNN_{graph}$, $GNN_{pool}$, LGPT, and the projection layer) to minimize the Cross-Entropy Loss between the predicted answer and the ground truth.

3. Key Contributions

The paper outlines two primary contributions:

1. Learnable Graph Pooling Token (LGPT): A novel pooling mechanism that represents graph information using a fixed number of learnable tokens. This balances the trade-off between the scalability of graph-level projection and the information richness of node-level projection.
2. Early Query Fusion: An approach that integrates query context into the graph embedding process prior to graph encoding. The authors argue this is more effective than late fusion methods where query and graph are combined post-encoding.

4. Experimental Results

The method was evaluated on the GraphQA benchmark, comprising three datasets: ExplaGraphs (commonsense reasoning), SceneGraphs (spatial understanding), and WebQSP (knowledge graph QA).

• Performance without LLM Training: The proposed method achieved a 4.13% average performance improvement over the strong baseline G-Retriever (which uses graph-level projection with 1 token) across all datasets. Specifically, it improved by 5.77% on ExplaGraphs, 3.92% on SceneGraphs, and 2.38% on WebQSP.
• Ablation Studies:
  • Early Query Fusion: Contributed an average improvement of 2.88% over the baseline.
  • LGPT: Contributed an average improvement of 3.87% over the baseline.
  • Combination: Using both techniques together yielded the highest gains, confirming their complementary nature.
• Comparison with Late Fusion: Experiments comparing Early Fusion against Late Fusion (using cross-modality pooling) showed that Early Fusion consistently outperformed Late Fusion. In some cases, Late Fusion actually degraded performance compared to not fusing at all, suggesting that combining fully processed embeddings can introduce noise.
• Token Count Analysis: The authors found that using 8 LGPTs yielded optimal performance. Increasing the number to 32 did not improve results and, in some cases, led to performance drops due to overfitting or information redundancy.
• LLM Training (LoRA): When the LLM was also trained using Low-Rank Adaptation (LoRA), the proposed method still outperformed baselines, achieving an additional 11.48% improvement over the LoRA-trained G-Retriever.

5. Significance and Claims

The paper claims that the proposed approach effectively addresses the "information bottleneck" in graph-augmented LLMs. By introducing learnable tokens, the method avoids the computational prohibitions of processing every node while preventing the information loss associated with compressing a graph into a single vector.

The authors emphasize that their method is particularly significant because it achieves these gains without requiring the training of the large language model itself, making it a scalable and efficient solution for handling complex textual-attributed graph data. The integration of query context early in the pipeline is presented as a critical factor in generating more effective graph embeddings that are tailored to the specific question being asked.