Riemannian Liquid Spatio-Temporal Graph Network

Liangsi Lu^1†, Jingchao Wang^2†, Zhaorong Dai³, Hanqian Liu⁴, Yang Shi^1✉

¹ Guangdong University of Technology ² Peking University ³ South China Agricultural University ⁴ Sun Yat-Sen University

RLSTG is a Riemannian geometry-based liquid spatio-temporal graph neural network for spatio-temporal graph data prediction tasks, supporting traffic flow prediction and link prediction tasks.

Environment Requirements

Quick Installation

# Create conda environment
conda env create -f environment.yaml

# Activate environment
conda activate rlstg

Dependencies

• Python 3.12
• PyTorch 2.5.1 (with CUDA 12.4 support)
• PyTorch Geometric
• torch-spatiotemporal
• Ray (for hyperparameter search)
• Other dependencies see environment.yaml

Supported Datasets

Traffic Datasets (Node Regression Tasks)

• METR-LA: Los Angeles area traffic sensor data
• PEMS03: California traffic data (PEMS system, region 3)
• PEMS04: California traffic data (PEMS system, region 4)

Social Network Datasets (Link Prediction Tasks)

• Enron: Enron email network dataset

Quick Start

Basic Usage

# Run RLSTG model on METR-LA dataset (default configuration)
python main.py --dataset metrla --model RLSTG

# Run on PEMS03 dataset
python main.py --dataset pems03 --model RLSTG --epochs 300

# Link prediction on Enron dataset
python main.py --dataset enron --model RLSTG --epochs 200

Running Modes

1. Hyperparameter Search Mode (Default)

python main.py --dataset metrla --mode search

This mode automatically searches for optimal hyperparameter configurations.

2. Evaluation Mode

python main.py --dataset metrla --mode eval

Uses the best configuration found to perform evaluation.

3. Full Mode

python main.py --dataset metrla --mode all

First performs hyperparameter search, then evaluates using the best configuration.

Configuration Parameters

Dataset Parameters

• --dataset: Choose dataset, options: ['metrla', 'pems03', 'pems04', 'enron']
• --vl_size: Validation set ratio, default 0.15 (15%)
• --ts_size: Test set ratio, default 0.15 (15%)
• --batch_size: Batch size, default 200
• --min_distance: Minimum interval for selecting time steps in traffic datasets, default 3
• --remove_ratio: Ratio of edges to remove from graph, default 0.0

Training Parameters

• --epochs: Number of training epochs, default 500
• --patience: Early stopping patience, default 10
• --metric: Evaluation metric to optimize, default 'MAE'
• --seed: Random seed, default 42
• --load_weight_only: Load only model weights (not training history)

Hyperparameter Search Parameters

• --ntrials: Number of trials for best configuration, default 3
• --max_ray: Maximum number of parallel processes for Ray search, default 8

Other Parameters

• --savedir: Results save directory, default './results'
• --debug: Debug mode
• --num_nodes: Number of nodes for visualization, default 3

Model Configuration

RLSTG Model Configuration

Model configurations are defined in the config.py file, main parameters include:

Geometric Space Configuration

• hyperbolic_dim: Hyperbolic space dimension
• spherical_dim: Spherical space dimension
• euclidean_dim: Euclidean space dimension
• hyperboloid_curvature: Hyperboloid curvature, default 1e-6
• sphere_curvature: Sphere curvature, default 1.0

Network Structure Configuration

• backbone_layers: Number of backbone network layers, default 1
• backbone_activation: Activation function, default 'tanh'
• dropout: Dropout ratio, default 0
• ode_unfold: ODE unfolding steps, default 4
• ode_solver: ODE solver, default 'gd'
• epsilon: Numerical stability parameter, default 1

Optimizer Configuration

• lr: Learning rate, search range [1e-6, 1e-2]
• weight_decay: Weight decay, default 1e-6
• optimizer: Optimizer type, default 'Adam'

Custom Configuration

To modify model configuration, edit the get_conf_rlstg function in config.py:

def get_conf_rlstg(input_dim, output_dim, time_dim, edge_dim, data_name):
    for lr in [1e-3, 1e-2, 1e-4]:  # Modify learning rate search space
        for wd in [1e-6, 1e-5]:     # Modify weight decay search space
            for h, s, e in [(0,0,16), (2,0,14)]:  # Modify geometric space dimension combinations
                # ... other configurations

Usage Examples

Example 1: Traffic Flow Prediction on METR-LA Dataset

python main.py \
    --dataset metrla \
    --model RLSTG \
    --epochs 500 \
    --batch_size 200 \
    --patience 10 \
    --mode all \
    --seed 42

Example 2: Quick Experiment on PEMS03 Dataset

python main.py \
    --dataset pems03 \
    --epochs 100 \
    --patience 5 \
    --debug \
    --mode search

Example 3: Link Prediction on Enron Dataset

python main.py \
    --dataset enron \
    --model RLSTG \
    --epochs 200 \
    --metric MAE \
    --mode all

Output Description

Results Saving

• Experimental results are saved in the ./results/ directory.
• Each dataset has its own subdirectory.
• Includes model weights, training history, hyperparameter configurations, etc.

Log Information

During execution, displays:

• Dataset loading information (number of nodes, edges, feature dimensions, etc.)
• Training progress and performance metrics.
• Best hyperparameter configuration.
• Final evaluation results.

Notes

1. GPU Support: Code automatically detects CUDA availability, GPU acceleration recommended for training
2. Memory Requirements: Large datasets may require substantial memory, adjust batch_size as needed
3. Parallel Search: Hyperparameter search uses Ray for parallelization, adjust parallelism with --max_ray
4. Data Preprocessing: Traffic data is sampled according to min_distance parameter for time steps.
5. Randomness: Use --seed parameter to ensure reproducible experiments

Troubleshooting

Common Issues

1. CUDA out of memory: Reduce batch_size or use CPU mode.
2. Dependency version conflicts: Ensure using the provided environment.yaml file.
3. Data loading failure: Check if datasets are correctly downloaded and extracted.

Getting Help

python main.py --help

View all available parameters and descriptions.

Citation

If you find our work helpful, please star 🌟 this repo and cite 📑 our paper. Thanks for your support!

@misc{lu2026rlstg,
  title={Riemannian Liquid Spatio-Temporal Graph Network},
  author={Liangsi Lu and Jingchao Wang and Zhaorong Dai and Hanqian Liu and Yang Shi},
  year={2026},
  eprint={2601.14115},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2601.14115},
}