LLM-Aware Relation Tokenizer

• The paper introduces the LLM-Aware Relation Tokenizer, a module that uses LLM-based prompts to efficiently encode complex multi-hop relations in heterogeneous graphs.
• It constructs compact, task-agnostic relational embeddings by pooling same-hop, same-type neighbor tokens and leveraging chain-of-thought reasoning in prompt design.
• Empirical results show state-of-the-art link prediction and node classification performance with reduced LLM inference cost compared to traditional meta-path methods.

The LLM-Aware Relation Tokenizer is a computational module designed to efficiently and effectively encode the multi-hop, multi-type relational structures present in heterogeneous graphs by querying a frozen LLM using specifically engineered prompts. This tokenizer is a key component of the Efficient LLM-Aware (ELLA) framework, addressing the problem of modeling rich and complex semantics among diverse node and relation types in heterogeneous graphs while maintaining computational feasibility. Unlike prior approaches that depend on meta-path heuristics or require heavy supervision, the Relation Tokenizer synthesizes the LLM’s advanced reasoning about graph paths and structure into a small set of compact, task-agnostic relational embeddings for each node, enabling state-of-the-art performance in both link prediction and node classification scenarios with greatly reduced LLM inference cost (Li et al., 22 Nov 2025).

1. Motivation and Conceptual Role

Heterogeneous graphs encode interrelated objects (nodes) of multiple types and relations, typified by bibliographic networks with entities like papers, authors, organizations, and relations such as authorship, citation, and affiliation. Modeling such rich interconnections traditionally requires expert-crafted meta-paths or relation-specific aggregation strategies, both of which scale poorly to large, real-world data and demand extensive supervision.

LLMs encapsulate extensive world knowledge and can perform complex reasoning over text, suggesting new opportunities for data-driven semantic encoding of these structures. However, a naïve approach—where every node or path invokes its own LLM call—incurs prohibitive computational cost. The LLM-Aware Relation Tokenizer addresses these obstacles by constructing per-node, per-hop, per-type “relation tokens” through a small number of carefully pooled and prompted LLM inferences, synthesizing the manifold semantics latent in the graph and its text attributes.

2. Architectural Overview

The Relation Tokenizer acts as an interface between the raw heterogeneous graph data and the downstream Hop-level Relation Graph Transformer. Its operation can be decomposed into three major steps:

• Node Token Extraction:

For each node $s$, if a text-rich attribute $x_s$ exists (such as a paper title), its LLM-derived embedding $u_s = \text{LLM}(x_s)$ is used. For nodes lacking such attributes, $u_s$ is mean-pooled from the LLM embeddings of one-hop neighbors: $$u_t = \operatorname{MeanPool}\{\text{LLM}(x_n) \mid n \in \mathcal{N}_t\}$$.

• Relation Token Construction:
  • All $i$-hop same-type neighbors' tokens are pooled: $$v_{s, \tau_k}^i = \operatorname{MeanPool}\{u_t \mid t \in \mathcal{N}_{s,\tau_k}^i\}$$.
  • A textual prompt $\text{RelationPrompt}^{(i)}(u_s, v_{s, \tau_k}^i)$ is constructed, articulating the connectivity and meta-path patterns:

Given node [PH] and node [PH], there exist these i-hop paths (with proportions p1, p2, ...). Steps: 1. Analyze relations based on path proportions and connection types. 2. Calculate the similarity (0–1) with justification.

• The [frozen] LLM is queried with this prompt to produce the relation token: $$u_{s, \tau_k}^i = \text{LLM}(\text{RelationPrompt}^{(i)}(u_s, v_{s, \tau_k}^i))$$.
  • Feedforward to Downstream Transformer:

For each hop $i$, the set $U_s^i = \{u_{s, \tau_k}^i\}$ passes through a type-mixing Transformer, producing a single hop-level feature $h_s^i$. The sequence $\{h_s^0 \equiv u_s, h_s^1, ..., h_s^K\}$ is then fed into a hop-mixing Transformer, whose output $z_s$ serves as the unified node embedding for subsequent tasks.

3. LLM-Based Encoding of Complex Graph Semantics

The Relation Tokenizer leverages prompt engineering to draw out high-order, multi-type dependencies between nodes:

• Prompt Templates:

Prompts are instantiated using a fixed template, enumerating available node types and relation labels. Meta-paths of specified length $i$ between the source node and pooled neighbors are articulated, including empirical occurrence proportions.

• Chain-of-Thought (CoT) Reasoning:

Prompts embed a two-step instruction that first requests analysis of the relations/path distributions, then similarity computation with contextual justification. This encourages the LLM to execute explicit, interpretable reasoning about the graph structure.

• Call Optimization:

By aggregating all same-type, same-hop neighbor information into a single input, only one LLM call is required per $(s, i, \tau_k)$ triple, as opposed to an exponential number of node-to-node queries.

4. Formal Algorithmic and Mathematical Description

The process of generating relation tokens can be summarized by the following formalism (with all notation and steps as described in (Li et al., 22 Nov 2025)):

1. Node and Neighbor Token Computation:

$u_s = \begin{cases} \text{LLM}(x_s), & \text{if $x_s$ is text-rich} \ \operatorname{MeanPool} \{\text{LLM}(x_n)\ |\ n \in \mathcal{N}_s\}, & \text{otherwise} \end{cases}$

$$v_{s, \tau_k}^i = \operatorname{MeanPool}\{u_t\ |\ t \in \mathcal{N}_{s,\tau_k}^i\}$$

2. Relation Token Generation:

$$u_{s, \tau_k}^i = \text{LLM}(\text{RelationPrompt}^{(i)}(u_s, v_{s, \tau_k}^i))$$

3. Type-Mixing Transformer (per hop $i$, over $L$ layers):

$$\tilde{U}_s^{(i,\ell-1)} = \operatorname{LayerNorm}(U_s^{(i,\ell-1)})$$

$$\hat{H}_s^{(i,\ell)} = \operatorname{MHA}(\tilde{U}_s^{(i,\ell-1)}) + U_s^{(i,\ell-1)}$$

$$U_s^{(i,\ell)} = \operatorname{FFN}(\operatorname{LayerNorm}(\hat{H}_s^{(i,\ell)})) + \hat{H}_s^{(i,\ell)}$$

$$h_s^i = \sum_{\tau_k} \alpha_{s, \tau_k}^i \cdot \hat{u}_{s,\tau_k}^i,\ \text{where}\ \alpha_{s,\tau_k}^i = \operatorname{softmax}(u_s \cdot \hat{u}_{s,\tau_k}^i)$$

4. Hop-Mixing Transformer ($L$ layers):

$H_s^{(\ell-1)} = [h_s^0, h_s^1, ..., h_s^K]$

$$\hat{H}_s^{(\ell)} = \operatorname{MHA}(\operatorname{LN}(H_s^{(\ell-1)})) + H_s^{(\ell-1)}$$

$$H_s^{(\ell)} = \operatorname{FFN}(\operatorname{LN}(\hat{H}_s^{(\ell)})) + \hat{H}_s^{(\ell)}$$

$$z_s = \hat{h}_s^0 + \sum_{i=1}^K \gamma^i \cdot \hat{h}_s^i,\ \text{with } \gamma^i = \operatorname{softmax}( [\hat{h}_s^0 \parallel \hat{h}_s^i]W )$$

5. End-to-End Algorithm and Integration

A compact pseudocode representation, as provided in (Li et al., 22 Nov 2025), is:

for each node s in graph: if text-rich: u_s ← LLM(x_s) else: u_s ← MeanPool(LLM(x_n) for n in Neighbors(s)) for i in 1..K: for each node-type τ_k: N = {t | t is i-hop neighbor of s of type τ_k} v = MeanPool(u_t for t in N) prompt = buildRelationPrompt(u_s, v, i, pathPatterns(s, N)) u_{s,τ_k}^i = LLM(prompt) U_s^i = {u_{s,τ_k}^i | τ_k ∈ nodeTypes} return {u_s, {U_s^i for i=1..K}}

The produced set of node and relation tokens $(u_s, \{U_s^i\})$ are passed to the Hop-level Relation Graph Transformer, which sequentially aggregates across types (within hops) and across hops, yielding the representations $z_s$ consumed by pre-training and fine-tuning heads.

6. Training Paradigm and Loss Functions

The ELLA framework incorporating the LLM-Aware Relation Tokenizer is trained in a two-stage scheme:

• Pre-training (Self-supervised Link Prediction):

For every relation type $r$, positive and negative node pairs are sampled. The loss is:

$$\mathcal{L}_{pre} = -\sum_{(s, t)\in \text{Pos}} \log \text{sim}(s, t) - \sum_{(s', t')\in \text{Neg}} \log (1 - \text{sim}(s', t'))$$

where $$\text{sim}(s, t) = \operatorname{Sigmoid}(z_s W_{\tau(s)} \cdot z_t W_{\tau(t)})$$. During tokenization, CoT prompts are used to ensure semantic alignment with this objective.

• Fine-tuning (Multi-type Node Classification):

With all encoders frozen, a shallow classifier head operates on $z_s$, optimized by:

$$\mathcal{L}_{ft} = -\frac{1}{N} \sum_{s\in V_{label}} \sum_{c=1}^C y_{s, c} \log \hat{y}_{s, c}$$

7. Empirical Performance and Ablation

Comparative ablation studies highlight the critical importance of several architectural and prompt-engineering aspects:

• Substituting each relation token with a representative neighbor’s node token yields a significant drop (IMDB director Mi-F1 from 79.3 to 77.9; IMDB movie 80.5 to 78.8).
• Providing all relation tokens jointly while ignoring hop separation catastrophically degrades performance (~47 Mi-F1), underscoring the necessity of hop-wise decomposition.
• Removing CoT (“Steps: …”) from the prompt induces a minor decrease (IMDB director 79.3 to 78.8).
• Omitting meta-path descriptions in the prompt results in performance at ~78.9 Mi-F1.
• These findings confirm the superiority of LLM-derived relation tokens over node-based aggregations and demonstrate the important roles of meta-path priors and CoT-guided prompt tuning (Li et al., 22 Nov 2025).

A summary table excerpting these empirical findings:

Ablation Variant	IMDB Director Mi-F1	IMDB Movie Mi-F1
Full Relation Tokenizer	79.3	80.5
Replace with Neighbor Token	77.9	78.8
No Hop Separation	~47	~47
No CoT in Prompt	78.8	--
No Meta-path in Prompt	~78.9	--

8. Significance and Integration in Heterogeneous Graph Learning

The LLM-Aware Relation Tokenizer operationalizes frozen LLM reasoning within complex, large-scale heterogeneous graphs by means of explicit path- and type-aware prompts, guided pooling, and hop-separated token construction. Its architectural design addresses semantic representation, computational efficiency, and alignment between unsupervised pre-training and task-specific fine-tuning. When coupled with the Hop-level Relation Graph Transformer, it achieves strong performance with up to a 4x speedup relative to other LLM-based graph encoders while scaling to models of up to 13B parameters (Li et al., 22 Nov 2025).

This suggests that prompt-optimized, relation-centric LLM inference modules—when judiciously pooled and integrated—offer a tractable and effective route to modeling the latent semantics of relational data domains where textual and topological variety co-occur.