A Hierarchical Quantized Tokenization Framework for Task-Adaptive Graph Representation Learning

Foundation models in language and vision benefit from a unified discrete token interface that converts raw inputs into sequences for scalable pre-training and inference. For graphs, an effective tokenizer should yield reusable discrete codes that capture both node semantics and relational structure across scales, yet prior quantization-based graph tokenizers typically combine residual vector quantization (RVQ) levels with fixed rules and often focus on a single structural view, limiting cross-task transfer. We present a hierarchical quantized tokenization framework with task-conditioned routing and dual-view token streams. It produces multi-scale codes and two synchronized sequences: a local stream that preserves node-level information and a diffusion-style multi-hop stream that summarizes connectivity. A lightweight router learns task-dependent mixtures over RVQ depths to select an appropriate granularity, while a gated cross-attention module aligns and fuses the two streams into a single token sequence without altering the downstream backbone encoder. Experiments on node classification and link prediction show consistent gains over strong quantized baselines at matched compute, with ablations verifying contributions from hierarchical quantization, adaptive routing, and fusion.

💡 Research Summary

This paper addresses a critical gap in graph foundation models (GFMs): the lack of a unified, reusable discrete token interface that can capture both node semantics and multi‑scale relational structure. Existing quantization‑based graph tokenizers rely on residual vector quantization (RVQ) with fixed rules for combining hierarchical codebooks, and they typically focus on a single structural view (either local or global). Consequently, the token representations are not adaptable to the varying granularity requirements of different downstream tasks such as node classification (which benefits from fine‑grained local information) and link prediction (which benefits from coarse‑grained multi‑hop connectivity).

To overcome these limitations, the authors propose TAU (Task‑Adaptive Unified Graph Tokenizer), a hierarchical quantized tokenization framework that introduces three key innovations:

1. Hierarchical RVQ for Dual‑View Features – The graph is first encoded with a GCN to obtain node embeddings H. In parallel, Personalized PageRank (PPR) diffusion is applied to H, producing multi‑hop enhanced features H_PPR. Both H and H_PPR are independently quantized using an M‑level RVQ pipeline, yielding two parallel token streams: a local stream (from H) and a global stream (from H_PPR). This dual‑view design ensures that both fine‑grained semantic cues and broader connectivity patterns are available for downstream processing.

2. Task‑Adaptive Quantization Routing (TAQR) – Rather than concatenating or equally weighting all RVQ levels, TAQR learns a scalar routing weight w(m) for each depth m. For each level, the quantized tokens are mean‑pooled to a summary vector s_m, which is fed through a two‑layer MLP to produce logits z(m). A temperature‑controlled softmax converts these logits into routing probabilities. The final token representation is a weighted sum over depths: C = Σ_m w(m)·C(m). This soft routing allows the tokenizer to emphasize coarse codes for tasks that need global context and fine residual codes for tasks that need detailed local information, without altering the backbone architecture.

3. Dual Cross‑Attention Fusion – A single-direction cross‑attention would force one token set to always be the query and the other the context, which is suboptimal across tasks. TAU therefore employs a bidirectional cross‑attention mechanism: the local token set attends to the global set and vice‑versa. Learned projection matrices W_Q, W_K, W_V generate queries, keys, and values for both streams, and the two attention outputs are combined through a gating module. This fusion aligns the two views while preserving their complementary information.

The resulting weighted, fused token sequence is fed unchanged into a downstream GFM backbone, which can be a pre‑trained Transformer or a GNN. Because the tokenizer’s parameters (codebooks, routing MLP, attention projections) are lightweight, they can be fine‑tuned jointly with the backbone or kept frozen for zero‑shot transfer.

Experimental Evaluation – The authors benchmark TAU on several node classification datasets (Cora, PubMed, ogbn‑arxiv) and link prediction datasets (ogbl‑collab, ogbl‑citation). They compare against strong quantization baselines such as GQT, VQ‑GNN, and other RVQ‑based tokenizers, matching compute budgets (FLOPs, memory). TAU consistently outperforms these baselines, achieving 1–3 percentage‑point gains in accuracy or ROC‑AUC. Notably, when a tokenizer trained on node classification is frozen and reused for link prediction, performance drops dramatically for the baselines, whereas TAU’s adaptive routing mitigates this degradation.

Ablation Studies – Three ablations are presented: (i) removing TAQR and using uniform averaging of levels, (ii) using only a single cross‑attention direction, and (iii) discarding the PPR‑enhanced stream. Each ablation leads to a measurable performance drop (≈1.5–2.2 pp), confirming that hierarchical quantization, task‑adaptive routing, and dual‑view attention each contribute meaningfully. Visualizations of the learned routing weights show that node‑classification tasks assign higher weight to deeper RVQ levels (fine residuals), while link‑prediction tasks allocate more weight to shallow levels (coarse structure), providing an interpretable signal of task‑specific granularity.

Significance and Limitations – TAU demonstrates that a graph tokenizer can be both multi‑scale and task‑aware while remaining compatible with existing foundation‑model backbones. The routing weights are directly interpretable, offering insight into which structural scales are most informative for a given downstream objective. However, the current implementation relies on explicit PPR computation, which may become a bottleneck for extremely large graphs. Additionally, the routing module is trained per‑task, so true zero‑shot transfer across completely unseen tasks may still require fine‑tuning.

Future Directions – The authors suggest extending TAQR to meta‑learning or reinforcement‑learning regimes to learn task‑agnostic routing policies, applying the framework to hypergraphs or dynamic graphs, and exploring alternative multi‑scale signal generators (e.g., Laplacian spectral filters) to reduce the computational overhead of PPR.

In summary, the paper presents a well‑motivated, technically sound, and empirically validated framework that advances graph tokenization toward a universal, scalable interface for graph foundation models.