Matrix Shuffle-Exchange Networks for Hard 2D Tasks
This addresses the need for efficient long-range dependency modeling in 2D data processing, particularly for algorithmic and logical reasoning tasks, representing a novel method rather than an incremental improvement.
The paper tackles the problem of limited receptive fields in convolutional neural networks for complex 2D tasks by proposing Matrix Shuffle-Exchange networks, which efficiently model long-range dependencies with O(log n) layers and O(n^2 log n) complexity, achieving superior performance on algorithmic and logical reasoning tasks compared to convolutional and graph neural network baselines.
Convolutional neural networks have become the main tools for processing two-dimensional data. They work well for images, yet convolutions have a limited receptive field that prevents its applications to more complex 2D tasks. We propose a new neural model, called Matrix Shuffle-Exchange network, that can efficiently exploit long-range dependencies in 2D data and has comparable speed to a convolutional neural network. It is derived from Neural Shuffle-Exchange network and has $\mathcal{O}( \log{n})$ layers and $\mathcal{O}( n^2 \log{n})$ total time and space complexity for processing a $n \times n$ data matrix. We show that the Matrix Shuffle-Exchange network is well-suited for algorithmic and logical reasoning tasks on matrices and dense graphs, exceeding convolutional and graph neural network baselines. Its distinct advantage is the capability of retaining full long-range dependency modelling when generalizing to larger instances - much larger than could be processed with models equipped with a dense attention mechanism.