Control the GNN: Utilizing Neural Controller with Lyapunov Stability for Test-Time Feature Reconstruction

TL;DR

This paper introduces a control-theoretic approach using Lyapunov stability to reconstruct node features in GNNs during testing, improving prediction accuracy under distribution shifts.

Contribution

It proposes a novel neural controller based on Lyapunov stability to reconstruct node features, ensuring stable and accurate predictions during test time.

Findings

Significant performance improvements across multiple datasets.

Theoretical guarantee of prediction proximity to ground truth.

Effective handling of distribution shifts in GNN testing.

Abstract

The performance of graph neural networks (GNNs) is susceptible to discrepancies between training and testing sample distributions. Prior studies have attempted to mitigating the impact of distribution shift by reconstructing node features during the testing phase without modifying the model parameters. However, these approaches lack theoretical analysis of the proximity between predictions and ground truth at test time. In this paper, we propose a novel node feature reconstruction method grounded in Lyapunov stability theory. Specifically, we model the GNN as a control system during the testing phase, considering node features as control variables. A neural controller that adheres to the Lyapunov stability criterion is then employed to reconstruct these node features, ensuring that the predictions progressively approach the ground truth at test time. We validate the effectiveness of our approach through extensive experiments across multiple datasets, demonstrating significant performance improvements.