Multiresolution Matrix Factorization and Wavelet Networks on Graphs

Truong Son Hy, Risi Kondor

Keywords: Multiresolution matrix factorization, wavelet theory, graph neural networks, Stiefel manifold optimization, reinforcement learning, molecular representation learning

TL;DR: A new learning algorithm based on Stiefel manifold optimization and Reinforcement learning to solve the problem of Multiresolution matrix factorization and produce a wavelet basis on graphs.

Abstract: Multiresolution Matrix Factorization (MMF) is unusual amongst fast matrix factorization algorithms in that it does not make a low rank assumption. This makes MMF especially well suited to modeling certain types of graphs with complex multiscale or hierarchical strucutre. While MMF promises to yields a useful wavelet basis, finding the factorization itself is hard, and existing greedy methods tend to be brittle. In this paper, we propose a ``learnable'' version of MMF that carfully optimizes the factorization with a combination of reinforcement learning and Stiefel manifold optimization through backpropagating errors. We show that the resulting wavelet basis far outperforms prior MMF algorithms and provides the first version of this type of factorization that can be robustly deployed on standard learning tasks. Furthermore, we construct the wavelet neural networks (WNNs) learning graphs on the spectral domain with the wavelet basis produced by our MMF learning algorithm. Our wavelet networks are competitive against other state-of-the-art methods in molecular graphs classification and node classification on citation graphs.

Track: Original Research Track