How to Make LMs Strong Node Classifiers?

Technical Summary: AUGLM – Augmented Graph Language Model for Node Classification

Problem Statement
While Language Models (LMs) have shown promise in graph learning tasks, existing approaches for node classification on Text-Attributed Graphs (TAGs) face a fundamental trade-off between performance and architectural flexibility. State-of-the-art (SOTA) methods like InstructGLM and LLaGA achieve strong results by modifying the LM's architecture—specifically by expanding the vocabulary with unique tokens for each node or injecting topology-aware embeddings into the token space. While effective, these modifications compromise the generality of the backbone LM, preventing seamless multi-task learning across diverse datasets and restricting the use of proprietary or closed-source models where architecture access is unavailable. The central question addressed by this work is: How can off-the-shelf, text-to-text instruction-tuned LMs achieve competitive node classification performance without any architectural modifications?

Methodology: AUGLM
The authors propose AUGLM (Augmented Graph Language Model), a framework that enhances standard LMs through data augmentation strategies rather than architectural changes. AUGLM reformulates node classification as a text-to-text task, preserving the LM's original structure while enriching the input context. The framework relies on three core components:

1. Retrieval-Based Aggregation:
   Instead of relying on simple multi-hop neighbor aggregation (which can introduce noise or cause "neighbor explosion"), AUGLM employs two complementary retrieval strategies to provide rich contextual information:

   • Topological Retrieval: Utilizes Personalized PageRank (PPR) to identify the most relevant neighbors for a target node based on graph structure, rather than fixed hop distances. The textual features of these top-K PPR neighbors are concatenated to form a structural context.
   • Prototypical Semantic Retrieval: Inspired by Retrieval-Augmented Generation (RAG) and Graph Transformers, this module treats the textual features of all non-target nodes as an external corpus. It retrieves "prototypes"—representative examples for each class—by first training a lightweight GNN to generate node predictions. The top-confident examples for each class are selected as prototypes. The system then retrieves the most semantically similar prototype (and its associated PPR neighbors) to the target node's text.

2. Classifier Guidance (Candidate Pruning):
   To address the LM's limited inherent understanding of graph topology, AUGLM incorporates guidance from a specialist model. A pre-trained lightweight GNN predicts a set of top-$I$ candidate labels for each node. These candidates are converted into text and included in the prompt. This strategy effectively prunes the search space, allowing the LM to focus its reasoning on distinguishing between highly relevant, topology-informed candidates rather than the entire label space.

3. Training and Optimization:
   The framework involves three parameterized modules: a lightweight GNN ($\psi$) for generating prototypes and candidates, a semantic retriever ($\phi$), and the backbone LM ($\theta$).

   • The GNN is pre-trained independently.
   • The semantic retriever is trained using a distribution-matching loss (KL divergence) to align retrieval probabilities with the LM's likelihood of generating the correct ground-truth label.
   • The backbone LM is fine-tuned using standard token-wise negative log-likelihood (NLL) loss on the augmented input prompts.
   • Crucially, the computational graphs for the retriever and the LM are decoupled via stop-gradient operators during specific training phases, ensuring efficient backpropagation.

Key Contributions

• Architectural Preservation: Unlike prior SOTA methods, AUGLM requires no modifications to the LM's architecture or vocabulary, enabling the use of any off-the-shelf text-to-text LM.
• Dual-Strategy Augmentation: The introduction of a hybrid retrieval mechanism combining topological (PPR) and semantic (prototype-based) information to overcome the limitations of simple neighbor aggregation.
• Lightweight Guidance: A novel method for transferring topological inductive bias from a GNN to an LM via candidate label pruning, enhancing performance without altering the LM's internal parameters.
• Multi-Task Flexibility: The framework supports joint training on multiple diverse datasets without the need for dataset-specific token expansions, a capability often lost in architecture-modifying approaches.

Experimental Results
The authors evaluated AUGLM on four real-world TAG datasets: Cora, Pubmed, ogbn-arxiv, and ogbn-products.

• Performance: Equipped with Flan-T5 backbones, AUGLM consistently outperformed SOTA text-output node classifiers (including InstructGLM and LLaGA).
• Comparison with Vector-Output Models: AUGLM achieved performance comparable to, and in some cases surpassing (e.g., on ogbn-products), the best vector-output models (such as TAPE and GLEM+RevGAT), despite operating purely in the text domain.
• Ablation Studies: Experiments confirmed that all three components (topological retrieval, semantic retrieval, and label pruning) contribute positively to performance, though their relative importance varies by dataset.
• Joint Training: A model jointly trained on all four datasets achieved performance comparable to models trained separately on each dataset, demonstrating the framework's ability to generalize across diverse graph structures without performance degradation.

Significance and Claims
The paper positions AUGLM as a crucial step toward bridging the gap between specialized, task-specific node classifiers and general-purpose foundation models. By demonstrating that high-performance graph learning can be achieved through input augmentation rather than architectural modification, the work suggests a path toward more versatile and widely applicable graph learning models. The authors claim that this approach retains the benefits of instruction tuning—specifically the ability to jointly train on diverse datasets—while overcoming the inflexibility of previous graph-LM integration methods. The work concludes that while current LMs have inherent limitations in understanding graph structures, these can be effectively mitigated through strategic data augmentation, paving the way for unified models that excel in multiple graph learning tasks.