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

TL;DR

This paper tackles robustness of graph neural networks under train-test distribution shift by reconstructing node features at test time without retraining. It proposes a control-theoretic framework that treats node features as inputs and predictions as states, using a neural controller $f_{\theta}$ guided by a neural Lyapunov function $V_{\phi}$ to ensure stable convergence to ground-truth labels. The method enforces Lyapunov stability through a joint training objective and counterexample generation via an SMT solver, enabling a single, stable test-time adjustment that yields class-representative embeddings for labeled nodes. Empirical evaluation on five datasets demonstrates improved accuracy and clearer separation between classes, with visualizations and Lyapunov analyses corroborating stability. This approach provides a theoretically grounded, parameter-free post-hoc boost to GNN robustness under distribution shifts, highlighting a practical control-theoretic direction for model-agnostic robustness.

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.

Figures (5)

• Figure 1: The overall process of our method. We input the original prediction $\hat{Y}t$ into the neural controller to obtain adjusted node features, and then feed these adjusted node features back into the well-trained GNN to obtain the prediction $\hat{Y}{t+1}$ at time $t+1$. Simultaneously, we input the original prediction $\hat{Y}_t$ into the neural Lyapunov function to obtain the Lyapunov value. We then use these values to compute the Lyapunov loss and update the parameters of both the neural controller and the neural Lyapunov function. Afterward, we use a solver to check whether there are states in the state space that do not satisfy the Lyapunov conditions. If such states exist, we add them to the training set and continue training until none remain.
• Figure 2: The architecture of the neural controller and the neural Lyapunov function and Lyapunov loss. We simultaneously train the neural Lyapunov function and the neural controller to obtain a control system that ensures stability and effectively guides the GNN predictions to gradually approach the ground truth.
• Figure 3: Visualization of the Cora dataset after dimensionality reduction with t-SNE. (a) Node embeddings derived from the original node features; (b) Node embeddings derived from the reconstructed node features by FRGNN; (c) Node embeddings derived from the reconstructed node features by our method. Orange points represent the training node embeddings, while blue points depict the test node embeddings. The red dashed line highlights that our method effectively aligns the embeddings of test nodes more closely with those of the training nodes compared to the original GNN and FRGNN.
• Figure 4: Visualization of the Cora dataset after dimensionality reduction with t-SNE. (a) Node embeddings derived from the original node features; (b) Node embeddings derived from the reconstructed node features by FRGNN; (c) Node embeddings derived from the reconstructed node features by our method. Different colore indicate different categories. Our method shows better clustering of node embeddings for each class than the original GNN and FRGNN.
• Figure 5: Visualization of Lyapunov function value. Dimension 1 represents dimension 1 of the t-SNE reduced state, Dimension 2 represents dimension 2 of the t-SNE reduced state, and the color represents the value of the Lyapunov function. The Lyapunov function values are zero near the equilibrium point and increase as they move away, indicating that the learned neural Lyapunov function meets the Lyapunov stability criterion.

Theorems & Definitions (6)

• Definition 1
• Theorem 1
• proof
• Remark
• Remark
• proof