YOSO: You Only Spectralize Once

Taking a Spectral Detour to Accelerate Graph Neural Network

This repository contains the PyTorch implementation for the paper "You Only Spectralize Once (YOSO): Taking a Spectral Detour to Accelerate Graph Neural Network".

YOSO is a Graph Neural Network (GNN) training scheme designed to significantly reduce computational overhead. It achieves this by primarily operating in a compressed spectral subspace. The key steps involve:

1. Projecting input node features onto a learnable orthonormal Graph Fourier basis ($U_l$).
2. Using a fixed sensing matrix ($\Phi$), informed by compressed sensing principles, to retain only $M$ significant spectral coefficients (where $M \ll N$, $N$ being the total number of nodes).
3. Performing all subsequent GNN computations (message passing and feature transformation) entirely in this reduced $M$-dimensional spectral domain.
4. Computing Sparse Spectral Coefficients: At the output layer, for the GNN's compressed output $T^{(L)}$, the sparse spectral coefficients are computed by iteratively solving a compressed sensing recovery problem. This is computed in each training iteration.
5. Reconstructing the full $N$-dimensional graph embeddings for predictions and task loss.
6. A joint loss function optimizes both the task-specific GNN loss and a reconstruction loss.

YOSO aims to achieve substantial training time reduction, particularly by minimizing non-computational overheads like data movement and sampling, while maintaining competitive model accuracy.

Requirements

• Python 3.9+
• PyTorch (e.g., 2.0.0+ recommended, tested with CUDA if available)
• OGB (ogb)
• PyTorch Geometric (torch-geometric) - for Reddit/Cora datasets
• scikit-learn (sklearn)
• NumPy (numpy)

Datasets

The code supports automatic download and preprocessing for the following datasets:

• Node Classification:
  • ogbn-arxiv (via OGB)
  • ogbn-products (via OGB)
  • reddit (via PyTorch Geometric)
• Link Prediction:
  • ogbl-ppa (via OGB)
  • ogbl-citation2 (via OGB)

Datasets will be downloaded to the directory specified by the --data_dir argument (default: data/).

Running the Code

The main script for running experiments is main.py. Configure experiments using command-line arguments defined in global_.py.

Example Commands:

• Node Classification on ogbn-arxiv:

  python main.py \
      --task node_classification \
      --dataset ogbn-arxiv \
      --data_dir data/ \
      --hidden_dim 128 \
      --lr 0.01 \
      --eta_U 0.001 \
      --alpha 0.5 \
      --beta 0.5 \
      --lambda_ 0.1 \
      --epochs 200 \
      --M 256 \
      --optimizer adam \
      --solver_lr 0.01 \
      --solver_iterations 15

• Link Prediction on ogbl-ppa:

  python main.py \
      --task link_prediction \
      --dataset ogbl-ppa \
      --data_dir data/ \
      --hidden_dim 128 \
      --output_dim_lp 64 \
      --lr 0.005 \
      --eta_U 0.0005 \
      --alpha 0.5 \
      --beta 0.5 \
      --lambda_ 0.05 \
      --epochs 200 \
      --M 128 \
      --optimizer adam \
      --solver_lr 0.01 \
      --solver_iterations 15

File Structure

• main.py: Main script to execute training and evaluation.
• model.py: Contains the PyTorch implementation of the YOSO model.
• dataloader.py: Handles downloading, loading, and basic structural preparation of datasets.
• preprocessing.py: Implements specialized preprocessing steps for YOSO.
• global_.py: Defines and parses command-line arguments.