GraphKAN: Graph Kolmogorov-Arnold Networks

repository·main·Indexed 43 days ago

https://github.com/willhua127/graphkan-graph-kolmogorov-arnold-networks

An experimental implementation of Kolmogorov-Arnold Networks (KAN) adapted for Graph Neural Networks (GNNs). The library requires torch and pytorch_geometric and recommends implementing KAN layers in the latent feature space via an initial nn.Linear projection to ensure training signals are received.

Tokens
622
Snippets
3
Records
4
Agent score
6%

What's inside GraphKAN

  1. How to implement KAN for graphs effectively

    main

    Based on experimental observations, GraphKAN should be implemented in the latent feature space rather than applying KAN layers directly to input features.

    Critical Requirement: You must use a nn.Linear layer to project input features into a low-dimensional latent space before passing them to the KAN layer. If you skip this linear projection and apply KANLayer directly to input features, the model will fail to receive any training signal.

    Recommended Pattern:

    1. Use nn.Linear(in_feat, hidden_feat) for initial projection.
    2. Use KANLayer on the resulting latent features.
    # RECOMMENDED: Project to latent space first
    self.lin_in = nn.Linear(in_feat, hidden_feat, bias=use_bias)
    # ... then use KAN on the output of self.lin_in
    
    # NOT RECOMMENDED: Direct KAN on input features (will lack training signal)
    # self.lin_in = KANLayer(in_feat, hidden_feat, grid_feat, addbias=use_bias)
  2. Choosing an optimizer for GraphKAN

    main

    When training GraphKAN, there is a trade-off between stability and convergence speed between SGD and Adam:

    • SGD (ASGD): Provides much higher stability during training.
    • Adam (AdamW): Converges significantly faster (e.g., ~200 epochs on Cora) but is less stable than SGD.

    If using SGD, expect much longer training times (e.g., ~10,000 epochs on Cora).

  3. Configure device for CPU training

    main

    If GPU (CUDA) training fails, you can force the model to use the CPU by explicitly setting the device in your arguments.

    Change the device assignment from the default auto-detection to a hardcoded CPU device:

    # Default (auto-detect)
    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
    
    # Change to this for CPU training
    args.device = torch.device('cpu')