Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond Parameter Finetuning: Test-Time Representation Refinement for Node Classification

Jiaxin Zhang, Yiqi Wang, Siwei Wang, Xihong Yang, Yu Shi, Xinwang Liu, En Zhu

TL;DR

TTReFT targets graph out-of-distribution generalization by shifting test-time adaptation from updating model parameters to refining node representations. It introduces three core components—uncertainty-guided node selection, a low-rank representation intervention (LoReFT), and an intervention-aware masked autoencoder (IAMAE)—to drive test-time refinement without forgetting pre-trained knowledge, with theoretical guarantees under distribution shifts. Empirically, TTReFT consistently improves OOD performance across five benchmarks, maintains 100% ID accuracy, and offers substantial efficiency advantages over parameter-finetuning baselines. Overall, the work establishes representation finetuning as a practical, effective paradigm for graph TTT with strong empirical and theoretical support.

Abstract

Graph Neural Networks frequently exhibit significant performance degradation in the out-of-distribution test scenario. While test-time training (TTT) offers a promising solution, existing Parameter Finetuning (PaFT) paradigm suffer from catastrophic forgetting, hindering their real-world applicability. We propose TTReFT, a novel Test-Time Representation FineTuning framework that transitions the adaptation target from model parameters to latent representations. Specifically, TTReFT achieves this through three key innovations: (1) uncertainty-guided node selection for specific interventions, (2) low-rank representation interventions that preserve pre-trained knowledge, and (3) an intervention-aware masked autoencoder that dynamically adjust masking strategy to accommodate the node selection scheme. Theoretically, we establish guarantees for TTReFT in OOD settings. Empirically, extensive experiments across five benchmark datasets demonstrate that TTReFT achieves consistent and superior performance. Our work establishes representation finetuning as a new paradigm for graph TTT, offering both theoretical grounding and immediate practical utility for real-world deployment.

Beyond Parameter Finetuning: Test-Time Representation Refinement for Node Classification

TL;DR

TTReFT targets graph out-of-distribution generalization by shifting test-time adaptation from updating model parameters to refining node representations. It introduces three core components—uncertainty-guided node selection, a low-rank representation intervention (LoReFT), and an intervention-aware masked autoencoder (IAMAE)—to drive test-time refinement without forgetting pre-trained knowledge, with theoretical guarantees under distribution shifts. Empirically, TTReFT consistently improves OOD performance across five benchmarks, maintains 100% ID accuracy, and offers substantial efficiency advantages over parameter-finetuning baselines. Overall, the work establishes representation finetuning as a practical, effective paradigm for graph TTT with strong empirical and theoretical support.

Abstract

Graph Neural Networks frequently exhibit significant performance degradation in the out-of-distribution test scenario. While test-time training (TTT) offers a promising solution, existing Parameter Finetuning (PaFT) paradigm suffer from catastrophic forgetting, hindering their real-world applicability. We propose TTReFT, a novel Test-Time Representation FineTuning framework that transitions the adaptation target from model parameters to latent representations. Specifically, TTReFT achieves this through three key innovations: (1) uncertainty-guided node selection for specific interventions, (2) low-rank representation interventions that preserve pre-trained knowledge, and (3) an intervention-aware masked autoencoder that dynamically adjust masking strategy to accommodate the node selection scheme. Theoretically, we establish guarantees for TTReFT in OOD settings. Empirically, extensive experiments across five benchmark datasets demonstrate that TTReFT achieves consistent and superior performance. Our work establishes representation finetuning as a new paradigm for graph TTT, offering both theoretical grounding and immediate practical utility for real-world deployment.
Paper Structure (38 sections, 1 theorem, 36 equations, 5 figures, 7 tables)

This paper contains 38 sections, 1 theorem, 36 equations, 5 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Problem Formulation
  4. PaFT for TTT on Graphs
  5. ReFT Paradigm on Graphs
  6. Methodology
  7. Framework Overview
  8. Uncertainty-Guided Node Selection
  9. Low-Rank Representation Intervention
  10. Intervention-Aware Masked Autoencoder
  11. Intervention-aware Masking
  12. Optimization Objective
  13. Complexity Analysis
  14. Theoretical Analysis
  15. Experiment
  16. ...and 23 more sections

Key Result

Theorem 1

(Effectiveness of Test-Time Intervention under Orthogonal Shift). Let $(X, Y) \sim P$ be training samples with labels generated by $Y = \text{softmax}(A X W)$. Under test distribution $P_t$ induced by $\tilde{X} = Q X$ for orthogonal $Q$, and given an intervention $\Phi(\tilde{X})$ with $U, V$ chose

Figures (5)

  • Figure 1: A comparison of the working mechanisms of PaFT and ReFT on graphs. (1) Parameter Finetuning (PaFT): The model parameters are updated during test-time training, which can lead to catastrophic forgetting, as illustrated by correctly classified training samples (blue nodes) being misclassified (turning red) after adaptation. (2) Representation Finetuning (ReFT): The pre-trained model parameters remain frozen. Adaptation is achieved by applying targeted, learnable interventions to a sparse subset of node representations (highlighted in yellow), successfully correcting misclassifications without forgetting pre-trained knowledge.
  • Figure 2: The overall framework of TTReFT. Our framework operates in three stages: (1) Pre-training: A GNN model is trained on source-domain data. All parameters are frozen after this stage. (2) Test-Time Representation Finetuning: For unlabeled test data, nodes with high predictive uncertainty (highlighted in yellow) are selected. A learnable Low-Rank Representation Intervention (LoReFT) is applied to them. The intervention parameters are optimized using a novel Intervention-Aware Masked Autoencoder (IAMAE) loss, which dynamically masks features based on local intervention density. (3) Inference: The frozen pre-trained model and the learned interventions are combined for final prediction, with interventions only applied to the selected uncertain nodes.
  • Figure 3: Ablation studies on key components. (a) Comparison of different self-supervised objectives for guiding intervention optimization. (b) Impact of node selection strategies on pubmed dataset. The bars represent accuracy when applying interventions to: no nodes (baseline), all nodes, test nodes, and our uncertainty-guided subset (top 10% entropy).
  • Figure 4: Sensitivity analysis of hyperparameters on the citeseer dataset. The $\star$ marks the optimal value for each parameter.
  • Figure 5: Visualization of representation refinement using t-SNE.

Theorems & Definitions (1)

  • Theorem 1