ChebMixer: Efficient Graph Representation Learning with MLP Mixer

TL;DR

ChebMixer presents an efficient graph representation learning approach by integrating fast Chebyshev polynomial spectral filtering to generate multiscale node tokens, refining them with an MLP-Mixer, and aggregating via Chebyshev interpolation. This avoids the quadratic self-attention bottleneck of graph Transformers while delivering rich multi-hop information for downstream tasks. Empirical results show state-of-the-art performance on graph node classification and medical image segmentation, along with favorable runtimes and robust ablation findings. The method suggests a unified framework for handling graph- and image-inspired data representations with scalable computation.

Abstract

Graph neural networks have achieved remarkable success in learning graph representations, especially graph Transformer, which has recently shown superior performance on various graph mining tasks. However, graph Transformer generally treats nodes as tokens, which results in quadratic complexity regarding the number of nodes during self-attention computation. The graph MLP Mixer addresses this challenge by using the efficient MLP Mixer technique from computer vision. However, the time-consuming process of extracting graph tokens limits its performance. In this paper, we present a novel architecture named ChebMixer, a newly graph MLP Mixer that uses fast Chebyshev polynomials-based spectral filtering to extract a sequence of tokens. Firstly, we produce multiscale representations of graph nodes via fast Chebyshev polynomial-based spectral filtering. Next, we consider each node's multiscale representations as a sequence of tokens and refine the node representation with an effective MLP Mixer. Finally, we aggregate the multiscale representations of nodes through Chebyshev interpolation. Owing to the powerful representation capabilities and fast computational properties of MLP Mixer, we can quickly extract more informative node representations to improve the performance of downstream tasks. The experimental results prove our significant improvements in a variety of scenarios ranging from graph node classification to medical image segmentation.

Figures (7)

• Figure 1: Applications of ChebMixer. Graph node classification and medical image segmentation.
• Figure 2: The proposed ChebMixer. ChebMixer first uses a novel neighborhood extraction module, $K$-hop extractor, to generate multiscale or multi-hop representations of graph nodes and treat them as a sequence of tokens. ChebMixer then refines the multi-hop representations of graph node with an effective MLP Mixer and develops a novel aggregator to aggregate the multiscale representations of the nodes.
• Figure 3: Architecture for medical image segmentation, which is composed of encoder, decoder, and skip connections. The encoder contains two modules: patch merging layer and GraphProp. The former performs downsampling, and the latter is constructed based on ChebMixer for feature extraction. The encoded feature representations in each encoder layer are fed into the corresponding CNN-decoder(i.e., the ResUpBlock in the figure, which has the function of up-sampling) via skip connections. The final output of the decoder goes through a head to obtain the segmentation result.
• Figure 4: Visual result comparison of our model and SOTA.
• Figure 5: Ablation experiment on $K$-hop aggregator. The "ChebNetII" represents the aggregation module in ChebNetII he2022chebnetii.