Matrix Shuffle-Exchange Networks for Hard 2D Tasks
Authors: Emīls Ozoliņš
Karlis Freivalds
Agris Šostaks
Venue:  (2021)

Abstract: 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.

Under review as a conference paper at ICLR 2021
MATRIX SHUFFLE-EXCHANGE NETWORKS FOR HARD 2D TASKS
Anonymous authors Paper under double-blind review
ABSTRACT
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 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 con- volutional neural network. It is derived from Neural and has (log n)) layers and O(n? log n) total time and space complexity for pro- cessing a is well-suited for algorithmic and logical reasoning tasks on matrices and dense graphs, exceeding convolutional and graph neural network baselines. Its distinct advantage the capability of retaining full long-range dependency modelling when generalizing larger instances much larger than could be processed with models equipped with a dense attention mechanism.
network
data matrix. We show that the Matrix Shuffle- Shuffle-Xcchange network
INTRODUCTION
Data often comes form of two-dimensional matrices. Neural networks are often used for process- ing such data usually involving convolution as the primary processing method. But convolutions are local, capable of analyzing only neighbouring positions in the data matrix. That good for images since the neighbouring pixels are closely related, but not sufficient for data having more distant relationships.
In this paper, we consider the problem how to efficiently process 2D data in a way that allows both local and long-range relationship modelling and propose new neural architecture, called Matrix Chuffle-Exchange network, to this end. The complexity of the proposed architecture is O(n log n) for processing n n data matrix, which is significantly lower than O(n*) if one would use the attention (Bahdanau et al. 2014| Vaswani et al. (017)) in its pure form. The architecture is derived from the Neural Shuffle Exciaange networks (Freivalds et al. 2019: Draguns et al. 2020) by lifting their architecture from ID to 2D.
We validate our model on tasks with differently structured 2D input/output data. It can handle complex data dnterddependencies present in algorithmic tasks on matrices such transposition, rotation, arithmetic operations and matrix multiplication. Our model reaches the perfect accuracy on test instances of the same size it was trained on and generalizes on much larger instances. In contrast, a convolutional baseline can be trained only on small instances and does not generalize. The generalization capability is an important measure for algorithmic tasks to say that the model has leameed an algorithm, not only fitted the training data.
Our model can be used for processing graphs by representing a graph with its adjacency matrix. has a significant advantage over graph neural networks (GNN) in case of dense graphs having additional data associated with graph edges (for example, edge length) since GNNs typically attach data to vertices, not edges. We demonstrate that the proposed model can infer important local and non-local graph concepts by evaluating it on component labelling, triangle finding and transitivity tasks. reaches the perfect accuracy test instances of the same size it was trained on and generalizes on larger instances while the GNN baseline struggles to find these concepts even on small graphs. The model can perform complex logical reasoning required to solve Sudoku puzzles. It achieves 100% correct solutions on easy puzzles and 96.6% on hard puzzles which is on par with the state-of-thh-an deep learning model which was specifically tailored for logical reasoning tasks.

---

Under review as a conference paper at ICLR 2021
2
RELATED WORK
Convolutional Neural Networks (CNN) are the primary tools for processing data with a 2D grid- like topology. For instance, VGG (Simonyan & Zissserman, 2014) and ResNet (He et al. (016a) enable high-accuracy image classification. Convolution is an inherently local operation that limits CNN use for long-ran range dependency modelling. The problem can be mitigated by using dilated (atrous) convolutions that have an expanded receptive field. Such an approach works well on image segmentation Yu & Koltun| 2015) but is not suitable for algorithmic tasks where generalization on larger inputs crucial.
The attention mechanism (Bahdanau et al. 2014] Vaswani et al., 2017) is widespread way of solving the long-range dependency problem sequence tasks. Unfortunately, its application to 2D data is limited due to it's s high Onn*) time complexity for n x n input matrix. Various sparse attention mechanisms have been proposed to deal with the quadratic complexity of dense attention by attending only to small predetermined subset of locations (Child et al. | [Blltayy et al. 2020 Zaheer et al., 2020). Reformer (Kitaev et al. 2020) uses locality-sensitive hashing to approximate attention in time Oin log n). Linformer (Wang et al. 2020) uses linear complexity approximation to the original attention by creating a low-rank factorization of the attention matrix. Sparse attention achieves great results on language modelling tasks, yet their application complex data, where attending to entire input is required, is limited.
Graph Convolutional Neural Networks (Micheli 2009
Atwood & Towsley 2016) generalizes
the convolution operation from the grid to graph data. They have emerged as powerful tools for processing graphs with complex relations (see (W et al. (2019) for great reference). Such networks have successfully been applied to image segmentation (Gong et al., 2019), program reasoning (Allamanis et al. 2018). and combinatorial optimization (Li et al. 2018) tasks. Nonetheless, Xu (2018) has shown that Graph Convolutional Networks may not distinguish some simple graph structures. To alleviate this problem they introduce Graph Isomorphism Network, that is as powerful as Seisfeiler-Lehmane graph isomorphism test and is the most expressive among Graph Neural Network models.
Neural algorithm synthesis and induction is a widely explored topic for ID sequence problems (Abolafia et al. 2020 Freivalds et al. 2019 Freivalds & Liepins 2018 Kaiser & Sutskever 2015
Draguns et al.
2020) but for 2D data, only few works exist. Shin et al. (2018) has proposed Karel program synthesis from the input-output image pairs and execution trace. Differentiable
Neural Computer (Graves al.] 20166), which employs external memory, has been applied to the SGRDLU puzzle game, shortest path finding and traversal tasks on small synthetic graphs Several neural network architectures have been developed for learning to play board games (Silver al.) 2018), including chess (David et al. (2116) and go (Silver et al. 2016: 2017), often using complex architectures or reinforcement learning.
ID SHUFFLE-EXCHANCE NETWORKS
Here we review the Neural Shuffle-Exchange Exchange (NSE) network for sequence-to-scquence processing, recently introduced by Freivalds et al. (2019) and revised by Draguns et al. (2020). This architec- ture offers an efficient alternative to the attention mechanism and allows modelling of long-range dependencies in sequences of length n in On log n) time. NSE network
neural adaption of well-known Shuffle-Exchange Exchange and Beneš interconnections networks, that allows linking any two devices using a logarithmic number of switching layers (see Daily & Towles (2004) for an excellent introduction).
The NSE network works for sequences of length n = 2*, where E Z+, and it consists of alternating Switch and Shuffle layers. Although all the levels are of the same structure, a network formed of 24- 2 Switch and Shuffle layers can learn a broad class functions, including arbitrary permutation of elements. Such network is called Beneš block. A deeper and more expressive network may be obtained by stacking several Beneš blocks; for most tasks, two blocks are enough.
The first k- 1 Switch and Shuffle layers of the Beneš block form Shuffle-Exchange block, the rest of the l- 1 Switch and Inverse Shuffle layers form its mirror counterpart. In the Beneš block, only Switch layers have learnable parameters and weight sharing is employed between layers of the same
2

---

Under review as a conference paper at ICLR 2021
Shuffle-Eae Shuffle-Eccaange Exchange block. A single Chuffle-Exchanee block is sufficient to guarantee that any input is connected to any output through the network, i.e. the network has 'receptive field' spanning the whole sequence.
In the Switch layer, elements of the input sequence are divided into adjacent non-overlapping pairs, and the Switch Unit is applied each pair. The Residual Switch Unit (RSU) proposed by (Daauuss et al. (2020) works the best. RSU has two inputs [ 22). two outputs [o,,0], and two linear transformations on the feature dimension. After the first transformation, root mean square layer normalization (RMSNorm) (Zhang & Sennrich, 2019) and Gaussian Error Linear Unit (GELU) (Hendrycks & Gimpel 2016) follow. By default, the hidden representation has two times more feature maps than the imput. The second linear transformation is applied after GELU and its output is scaled by scalar h. Additionally, the output of the unit connected to its input using a residual connection that is scaled by learnable parameter 8. RSU is defined as follows:
[1 12]
g GELU (RMISNOM Zi))
Wg +
[01, 02] as) @ithoec
In the above equations, Z, W are weight matrices of size 2m x 4m and 4m x 2mn, respectively, where is the number of feature maps; is a vector of size 2m and b is a bias vector - all of those are learnable parameters; @ denotes element-wist vector multiplication and the sigmoid function. is initialized as o (r) and h initialized as vI- -2 0.25, where r 0.9.
The Shuffle layer performs the perfect shuffle permutation of the sequence elements. In this permu- tation, the destination address an element is obtained by eyelic bit rotation of its source address. The rotation to the right is used for the Shuffle layer and to the left for Inverse Shuffle layer. Shuffle layers do not have learnable parameters.
4 THE MATRIX
NETWORK
The naive way to employ NSE on 2D data would be by transforming an input matrix into a sequence using raster scan flattening (e.g., tf. reshhap)), but such solution has two main drawbacks. Firstly, NSE requires more than 8 seconds for IM element sequence (which matches flattened 1024x1024 matrix) in inference mode (Draguns et al. 2020), making its use unpractical on large matrices. Secondly, raster scan flattening doesn't preserve element locality the inputs can be of different sizes, limiting the networks ability to generalize on large inputs an essential requirement for algorithmic tasks.
We propose Matrix Shuffle-Exchange (Matrix-S)) network - an adoption of NSE for 2D data, that is significantly faster, generalizes large matrices and retains the property of receptive field spanning the whole input matrix (see Appendix [8] for more details). The model is suitable for processing n x n input array, where 2* for l Z and each element is vector of m feature maps. Inputs that don't fulfil this requirement has to be padded to the closest 2 x 2² array.
The Matrix-SE network works by rearranging the input matrix into 1D sequence, then applying several interleaved Quaternary Switch and Quaternary Shuffle layers and then converting data back 2D, see Fig. Rearranging of n X n input to n² sequence has to be done in such way that enables generalization on larger matrices. For this purpose, we read out values according to the Z-Order curve, as has recursive structure and it preserves element locality regardless of the input size | The resulting ordering is the same as one would get from depth-fire first traversal of quadtree. We call this transformation a Z-Order flatten. The left matrix in Fig. 2 shows matrix elements indexed according to Z-Order curve. The Z-Order unflatten transforms n² output sequence of the last Quaternary Switch layer to n X n representation according to Z-Order curve indexing.
The middle part involving Quaternary Switch (QSwitch) and Quaternary Shuffle (QShuffle) layers is structured as one or more Beneš blocks the same way as in the NSE network. The QSwitch and QShuffle layers differ from Switch and Shuffle layers in NSE, by performing the computations in groups of four instead of two. Such principle (using 4-to-4 switches instead of 2-to-2) has been proved for Shuffle-Exchanpe computer networks (Dally & Towles 2004). That reduces network
https://enn wik wwtp:e/wwwwiiiiiii.o..o..//
3

---

Under review as a conference paper at ICLR 2021
depth 2x times preserving the theoretical properties of NSE. Each Beneš block consists of interleaved - QSwitch and QShuffle layers that form Shuhane-Exchange block, followed by k - QSwitch and block. In contrast
Inverse QShuffle layers that form mirror counterpart of the first Shuffle-Exchanee with NSE, we finalize Beneš block with one more QSwitch layer, making it more alike to Beneš computer networks. Last QSwitch layer serves expressiveness. Weight sharing is employed between QS layers of same Ehuffle-Exchange block. The last QSwitch layer of the Beneš block does participate in weight sharing. If the network has more than one Beneš block, each block receives a distinct set of weights.
output for the Benes block and improves model
block
Banad block
11 Shawlorlstt
hpst
2
Figure 1: Matrix Shuffle-Exchange model consists of one or more Beneš block that is enclosed by Z-Order flatten, and Z-Order unflatten transformations. For a 2 x 2 input array, Beneš block has total of 2k Quaternary Switch and 21 2 Quaternary Shuffle layers, giving rise to O(n log n) time complexity.
In the QSwitch layer, we divide adjacent elements of the sequence (corresponds to flattened 2D input) into non-overlapping 4 element tuples. That is implemented by reshaping the input sequence into a times shorter sequence where elements are concatenated along the feature dimension as [E 22, { Then Quaternary Switch Unit (QSU) is applied each tuple. QSU is based RSU but has 4 inputs {i] 1² 13, id] and 4 outputs (o o², o². o The rest of the Unit structure is left unchanged from RSU. The QSU is defined as follows:
i1, i2, 23, 21}
g = (RLU(MM Zi))
Wg +
(o, 02, 0g 0a] = a(s)
In the above equations, Z. W are weight matrices of size 4m X 8m and 8m x 4m, respectively; s is vector of size 4m and is bias vector all of those learnable parameters; h is a scalar value; C denotes element-wise vector multiplication and or is the sigmoid function. We initialize s and same in RSU. Weights sharing is employed between all QSUs of the same QSwitch layer.
00
01
10
11
00
10
01
11
02
03
12
20
30
21
20
21
30
31
02
12
03
13
22
23
32
33
22
32
23
33
Figure 2: Depiction how the Quaternary Shuffle layer permutes the matrix (technically it permutes ID sequence representation of matrix). The left image shows the matrix elements ordered according to Z-Order curve, where the ordering indices are represented by base-4 numbers. The right image shows element order after the Quaternary Shuffle. Quaternary Shuffle permutation be interpreted as splitting matrix rows into two halves (white and green) and interleaving the halves, then applying the same transformation columns (white and red).
The QShuffle layer rearranges the elements according to the cyclic digit rotation permutation. This permutation for sequence is defined S(x] =Slgrotate(1; k)], where grotate(2, k) applies
4

---

Under review as a conference paper at ICLR 2021
eyclic shift by one digit to which is treated as quaternary (base 4) number and le is its length in quaternary digits ? For the first h QShuffle layers of the Beneš block, we apply rotation to the right- to the left for the remaining k layers.
EVALUATION
We have implemented our proposed architecture in TensorFlow. The code will be published on GitHub. All models are trained and evaluated on single Nvidia RTX 2080 Ti (11 GB) GPU using Rectified Adam (RAdam) optimizer (Liu et al.. 2019) and softmax cross-entropy loss.
We evaluate the Matrix-SE model on several 2D tasks. Our main emphasis is on hard, previously unsolved tasks: algorithmic tasks on matrices and graphs and logical reasoning task solving Sudoku puzzles. The proposed architecture is compared to ResNet and GNN baselines which are reasonable alternatives that could be used for these tasks. For algorithmic tasks, to achieve generalization, we use curriculum learning introduced by Freivalds & Liepins] (2018) where models of different sizes are instantiated and trained simultaneously. We additionally evaluate the speed of the reported models and show that the Matrix-SE SE network achieves comparable or better results. Model is also evaluated on CIFAR-10 dataset and the results are given in Appendix
5.1 INFERRING MATRIX ALGORITMM
We selected four matrix tasks of varying difficulty- transpose, rotation by 90 degrees, element-wise XOR of two matrices, and matrix squaring. For the transpose and rotation tasks, we generate random input matrix from alphabet 1-11 and expect the network learn the required element permutation. For the XOR task, we generate two binary matrices, place them side by side, add separator symbols between them, then apply padding to obtain a square matrix. All tasks, except matrix squaring, have simple (n²) algorithms.
Matrix squaring is the hardest of the selected tasks. We generate random binary square matrix and ask the model to output its square (matrix multiplication with itself) modulo 2. That is a challenging task for the Matrix-SE model having the computation power (n log n),, while the currently lowest asymptotic complexity for the matrix multiplication algorithm is (( (Le Gall 2014b).
Nonetheless, fast approximation algorithm for matrix multiplication of Ó(n)) complexity exists (Manne & Pal 2014). The lower bound for matrix multiplication for real and complex numbers in the model of arithmetic circuits is s(n2 log n) proven by Raz (2002). Matrix squaring has the same complexity as matrix multiplication since these problems can be easily reduced one to the other. We selected matrix squaring for evaluation because it has only one input.
We evaluate two models - Matrix-SE network with 2 Beneš blocks, 96 feature maps having 3.5M learnable parameters and ResNet-29 model with 128 feature maps, kernel size of 3x3 having 4.5M learnable parameters (the model structure is given in Appendix (C) Such ResNet-29 has a receptive field of size x 59 that is enough to span the whole input on the largest matrix seen training, making it on par with the Matrix-SE network that has a full receptive field for any input size Both models are trained on inputs up to size 32 x 32 for 500k steps and evaluated on size up to 1024 x 1024. Matrix-SE reaches perfect accuracy
the test examples of size up to 32 x 32 for all tasks, see Table T It generalizes on larger inputs for all tasks except matrix squaring. ResNet-29 struggles to learn these tasks even for size 32 x 32 and doesn't 't generalize. Matrix also converges significantly faster than ResNet-29 (see Appendix D for more details).
5.2 INFERRING GRAPH ALGORITIIMS
graph can be easily represented by its adjacency matrix making Matrix-SB model a good fit for tasks on graphs, especially for tasks where edges have additional information. We consider three graph algorithmic tasks: connected component labelling, transitivity, and triangle finding.
For the connected component labelling task, we initialize a labelled graph with random edge labels in range 2-0 The task is to label all edges each connected component with the lowest label among
nnte that it matches the definition of le above derived from 7
We used library proposed by Araujo al. 2219)) for calculating receptive field
5

---

Under review as a conference paper at ICLR 2021
Table I: Element-wise test accuracy (average of runs) on algorithmic tasks matrices depending on input length. Both models are trained using curriculum learning on matrices up to size 32 32.
Size of input matrix (n
Model
Task
Train size
Generalization size
4
8
16
32
64
128
256
512
1024
Transpose
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
Rotate by 90
1.0
1.0
1.0
1.0
1.0
1.0
0.8
0.41
0.19
Matrix-SE
Bitwise XOR
1.0
1.0
1.0
1.0
0.96
0.89
0.78
0.67
0.56
Matrix squaring
1.0
1.0
1.0
1.0
0.49
0.49
0.48
0.45
0.41
Transpose
1.0
1.0
1.0
0.99
0.36
0.15
0.08
0.06
0.05
Rotate by 90
1.0
1.0
1.0
0.99
0.17
0.08
0.06
0.06
0.05
ResNect-29
Bitwise XOR
1.0
1.0
1.0
1.0
0.4
0.31
0.28
0.26
0.26
Matrix squaring
1.0
1.0
1.0
0.88
0.59
0.55
0.52
0.51
0.50
all the component's 's edges. For the transitivity task, the goal is add edges directed graph for every two vertices that have a transitive path of length 2 between them. The hardest of the considered graph algorithmic tasks is the triangle finding. The fastest currently known algorithm relies on matrix multiplication and has asymptotic complexity of (1222373), where is the vertex count (Alon et al. 1997). Triangle finding is a fruitful field for quantum computing where an algorithm of quantum query complexity 5/2) (Le Gall 2014a) has been achieved. Therefore it is interesting Matrix-SE can infer On log n.) time algorithm, at least approximately. For this task, we generate random complete bipartite graphs and add a few random edges. Although dense, such graphs have only a few triangles. The goal is to return all the edges belonging to any triangle.
Table 2: Element-wise test accuracy (average of runs) on algorithmic tasks on graphs depending on vertex count. Both models are trained using curriculum learning graphs with up to 32 vertices.
Vertex count
Model
Task
Train size
Generalization size
8
16
32
64
128
256
512
1024
Component Labeling Triangle Finding
1.0
1.0
1.0
1.0
0.99
0.97
0.96
0.91
Matrix-SE
1.0
1.0
1.0
1.0
1.0
0.98
0.93
0.87
Transitivity
1.0
1.0
1.0
0.96
0.88
0.82
0.77
0.71
Component Labeling
0.87
0.92
0.94
0.91
0.82
0.74
0.67
0.63
GIN
Triangle Finding
0.91
0.95
0.95
0.87
0.51
0.63
0.77
0.89
Transitivity
0.82
0.93
0.92
0.80
0.56
0.56
0.79
We use Matrix-SE network with 2 Beneš blocks, 192 feature maps and 14M learnable parameters (similar results can be achieved with 96 feature maps and 3.5M parameters). Since Graph Neural Networks natural choice for graph task, we also compare Matrix-SE with Graph Isomorphism Network (GIN) (Xu et al. 2018), that is proved (Xu et al. 2018) to be the most expressive among the message passing based Graph Neural Networks. For this purpose, use GIN implementation from DGL libraryF with 5 Graph Isomorphism layers where 3-layer perceptron is used at each layer. The final graph representation mapped to the output using another 3-layer perceptron. The GIN network used in our experiments has 4.7M learnable parameters. Graph Neural Networks are more suitable for vertex-centric tasks, where all the features are represented the vertices. To represent edge labels into GIN we follow a common practice (Yang et al.) 2020 Veličković et al., (017) and concatenate vertex features with element-wise multiplication between the edge and vertex features.
Both models are trained on graphs up to 32 vertices up to the moment where the loss hasn't decreased for 20k steps. Then models are tested on graphs with up to 1024 vertices. For the proposed tasks, Matrix-SE reaches perfect accuracy on test examples with up to 32 vertices and generalizes to moderately larger graphs. We find that Graph Isomorphism Network struggles to solve these tasks

---

Under review as a conference paper at ICLR 2021
graph sizes seen in training time. Also, GIN consumes GPU memory than Matrix- the graph for the transitivity task with 1024 vertices don't 't fit in GPU memory.
5.3 SOLVING SUDOKU
Solving a Sudoku puzzle (see Appendix A for more details) requires sophisticated reasoning and can be a challenging task for a human player. Since contemporary solvers, typically, use a backtracking search to solve Sudoku (Norvig] 2009 Crook 2009). it's interesting to see if neural models are capable of such task without backtracking.
The complexity of Sudoku puzzle is characterized by how many numbers are already given in the puzzle. Park (018)) has proposed solving easy Sudoku puzzles (24-36 givens) using convolutional neural network. Wang al. (2019) achieves close-to-perfece accuracy on slightly easier Sudoku puzzles (27-37 givens) by encoding each puzzle as MAXSat instance and then applying SATNet to it. Recurrent Relational Networks (Palm al. 2018) can solve significantly harder Sudokus (17-34 givens) they encode Sudoku puzzle as graph and recursively employ Graph Neural Network for 32 steps. At each step, hidden representation is mapped to the solution by a 3-layer perceptron. Loss obtained the sum of a softmax cross- cross-entropy loss of output each step.
Table 3: Comparison of different neural Sudoku solvers. Sudoku puzzles with fewer givens are harder. The accuracy represents the fraction of completely correctly solved test puzzles.
Model
Givens
Accuracy %
Matrix Shuffle-Excaage (this work)
17-34
96.6
Recurrent Relational Networks (Palm et al.
2018
17-34
96.6
Matrix -xcihlnge (this work)
24-36
100.0
Convolutional Neural Network (Park 2018
24-36
70
SATNet (Wang 2019
27-37
98.3
We validate the Matrix-SE model on hard Sudoku dataset (17-34 givens) created by Palm al. (2018) (see example from test set in Appendix(A) and easy Sudoku dataset (24-36 givens) by Park (2018 8). Matrix-SE model is trained the same way as Recurrent Relational Network we imitiate model with 2 Beneš blocks and 96 feature maps (3.5M total learnable parameters) and then apply it for 10 steps. At each step, hidden representation mapped to output by one GELU layer and one linear layer then the softmax cross-entrond cross-entropy loss is applied. We represent Sudoku puzzle in its natural form matrix and pad it to size 16 x and 100% total accuracy It's worth mentioning that Matrix-SE model requires fewer steps than Recurrent Relational Network (64 steps) and Sudoku doesn't need to be converted graph, significantly simplifying the training pipeline.
16. The model achieves 96.6% total accuracy on hard Sudoku dataset easy Sudoku dataset using 30 relational steps at evaluation, see Table 3
5.4 PERFORMANCE OF THE MODEL
Here we compare the training and evaluation speed of Matrix-SB network with other neural models. As benchmark, we chose two algorithmic tasks matrix squaring and graph transitivity task. The comparisons are done on single Nvidia RTX 2080Ti (11 Gb) GPU using Tensorflow. We report an average of 300 steps on single element batch.
On the matrix squaring task, Matrix-SE (3.5M parameters) is compared ResNet-29 (4.3M parame- ters) and Residual Shuhfle-Fxchange (RSE) network (3.5M parameters), see Fig. 3 Matrix-SE model works approximately 9x faster than RSE on large matrices and on par with convolutional neural network. It's worth mentioning that Matrix-SE is fast despite it has varying depth depending on the input size and its receptive field grows to cover the whole sequence while ResNet-29 has a fixed depth and receptive field.
To compare with graph networks, we chose the transitivity task for evaluation. We use Matrix-ST
and Graph Isomorphism Network (GIN) of configurations given in Section[522] As depicted in Figure the Matrix-SE model is slightly faster, both during training and evaluation, and it can process 2x larger dense graphs than GIN.
7

---

Under review as a conference paper at ICLR 2021
Declerenn RSC nen
-0.
Vatroda
6.20
8:135
8,002
3.07
0.00
18
172
Figure 3: Training and evaluation speed on ma- trix squaring task for single input instance.
Figure 4: Training and evaluation speed on graph transitivity task for a single input instance.
ABLATION STUDY
The most prominent feature that distinguishes the Matrix-SE from its predecessor - RSE that uses Z-Order flatten transform data from 2D matrix to ID sequence. We study whether it is necessary by substituting Z-Order flatten with raster scan flatten. We train both versions on matrix squaring, transpose and triangle finding tasks for 500k steps (for triangle finding, first 150k shown) matrices of size up to 32 x 32. We can see in Fig. 6 that the version with Z-Order flatten generalizes to larger instances, and the other does not.
We also assess the effect of the model size in terms of feature map count and the model depth (Beneš blocks) on the convergence speed and accuracy. For the matrix squaring task, we use Matrix-SE with 2 Beneš blocks and 96 feature maps as a baseline, but, for triangle finding, Matrix-SE with 2 Beneš blocks and 192 feature maps. The results work best for the chosen tasks. A larger and deeper model in most cases yield better results, but when reasonable accuracy is obtained, diminishing improvements are observed. On the matrix squaring and triangle finding tasks Z-Order flatten increase the convergence speed and accuracy compared to raster scan flatten.
depicted on Fig. 5 Fig. 7 The proposed baseline models
fatur
Figure 5: The effect of model Figure 6: Generalization impact size on the accuracy on the ma- of Z-Order flatten compared to trix squaring task.
Figure 7: The effect of model size on the accuracy for the tri- raster scan flatten matrix tran- angle finding task.
sitivity task.
7 DISCUSSION
We have introduced the Matrix Shuffle-Fxchangg model, which is of complexity Onn² log n) and can model long-range dependencies 2D data. It is capable of inferring matrix and graph algorithms purely from input-output examples exceeding convolutional and graph neural network baselines in terms of accuracy, speed and generalization power. It has shown considerable logical reasoning capabilities validated by ability solve Sudoku puzzles, where it is par with state-of-the-arr neural networks.
This is the initial introduction of the Matrix Shuffle-Exchange model, and we further expect its signif- icant improvements by fine- tuning its constituents, combination with other models and application to other tasks.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Daniel A Abolatia, Rishabh Singh, Manzil Zaheer, and Charles Sutton. Towards modular algorithm induction. arXiv preprint arXiv:2203. 2020.
Miltiadis Allamanis, Mare Brockschmidt, and Mahmoud Khademi. Learning to represent programs with graphs. In International Conference on Learning Representations, 2018.
Noga Alon, Raphael Yuster, and Uri Zwick. Finding and counting given length cycles. Algorithmica, 17(3):209-23,, 1997.
Andre Araujo, Wade Norris, and Jack Sim. Computing receptive fields of convolutional neural networks. Distill, 2019. doi: 10.23915ddistill00001.. https://distill.pub//2/ http://iiilllpp22200000oooomnnee fields.
James Atwood and Don Towsley. Diffusion-convolutional neural networks. In Advances in neural information processing systems, pp. 1993-2001, 2016.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXivv1607.064500, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. arXiv preprint arXiv:/409.0473, 2014.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- document transformer. arXiv preprint 255 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 19G1::9904.00000 2019.
JF Crook. A pencil-and-paper algorithm for solving sudoku puzzles. Notices of the AMS, 56(4): 460-468, 2009.
William James Dally and Brian Patrick Towles. Principles and practices of interconnection networks. Elsevier, 2004.
Omid E David, Nathan S Netanyahu, and Lior Wolf. Deepchess: End-to-end deep neural network for automatic learning in chess. In International Conference on Artificial Neural Networks, pp. 88-96. Springer, 2016.
Andis Draguns, Emils Ozolinš, Agris Šostaks, Matiss Apinis, and Karlis Freivalds. Residual Shuffle- Exchange Networks for Fast Processing of Long Sequences, 2020.
Karlis Freivalds and Renars Liepins. Improving the Neural GPU architecture for algorithm learning. The ICML workshop Neural Abstract Machines & Program Induction v2 (NAMPI 2018), 2018.
Karlis Freivalds, Emils Ozolinš, and Agris Šostaks. Neural Shuffle-Exchange Networks Sequence Processing in O log n) Time. In Advances in Neural Information Processing Systems, pp. 6626-6637, 2019.
Ke Gong, Yiming Gao. Xiaodan Liang, Xiaohui Shen, Meng Wang, and Liang Lin. Graphonomy: Universal human parsing via graph transfer learning. In CVPR, 2019.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñssa, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016a.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference computer vision, pp. 630-645. Springer, 201 20.

---

Under review as a conference paper at ICLR 2021
Dan Hendrycks and Kevin Gimpel. Gaussian enror arXiv:1606.08415. 2016.
linear units (gelus). arXiv preprint
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: 2015.
Eukasz Kaiser and Ilya Sutskever. Neural GPUs learn algorithms. arXiv preprint arXiv: 151.0255, 2015.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint 2XXv:20001004551, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
François Le Gall. Improved quantum algorithm for triangle finding via combinatorial arguments. In 2014 IEEE 55th Annual Symposium on Foundations of Computer Science, pp. 216-225. IEEE, 2014a.
François Gall. Powers of tensors and fast matrix multiplication. In Proceedings of the 39th international symposium on symbolic and algebraic computation, pp. 296-303, 2014b.
Zhuwen Li, Qifeng Chen, and Vladlen Koltun. Combinatorial optimization with graph convolutional networks and guided tree search. In Advances in Neural Information Processing Systems, pp. 539-548, 2018.
Zhouhan Lin, Roland Memisevic, and Kishore Konda. How far can we go without convolution: Improving fully-connected networks. arXiv preprint arXiv:1511.02580 2015.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance the adaptive learning rate and beyond. arXiv preprint arXiv: 1.Ki:v9908.0265. 2019.
Shiva Manne and Manjish Pal. Fast approximate matrix multiplication by solving linear systems. arXiv preprint arXiv:1408.4283 2014.
Gary McGuire, Bastian Tugemann, and Gilles Civario. There is 16-clue sudoku: Solving the sudoku minimum number of clues problem via hitting set enumeration. Experimental Mathematics, 2302)1190-217, 2014.
Alessio Micheli. Neural network for graphs: A contextual constructive approach. IEEE Transactions on Neural Networks, 20(3):498-511 511, 2009.
Peter Norvig. Solving every sudoku puzzle. Preprint, 2009.
Rasmus Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. Advances in Neural Information Processing Systems, pp. 3368 -378, 2018.
Kyubyong Park. Can convolutional neural networks crack sudoku puzzles? https::/
com/ com/Kyubyongee Kyubyong/uuuokk.... 2018
Ran Raz. On the complexity of matrix product. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing, pp. 144-151, 2002.
Eui Chul Shin, Illia Polosukhin, and Dawn Song. Improving neural program synthesis with inferred execution traces. In Advances in Neural Information Processing Systems, pp. 8917-8926, 2018.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwiesser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):444, 2016.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv:1712.01815, 2017.
10

---

Under review as a conference paper at ICLR 2021
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou. Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-11 44, 2018.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.155: 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 2017.
Po-Wei Wang, Priya L Donti, Bryan Wilder, and Zico Kolter. Satnet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver. arXiv preprint 19Xiv:990..12144 2019.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. 2020.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S Yu. A comprehensive survey on graph neural networks. arXiv preprint arXiv:1901.0059. 2019.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint 1rXiv:1810.00880 2018.
Lei Yang, Dapeng Chen, Xiaohang Zhan, Rui Zhao, Chen Change Loy, and Dahua Lin. Learning to cluster faces via confidence and connectivity estimation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13369-13378, 2020.
Fisher Yu and Vladlen Koltun. Multi-scale context aggregation by dilated convolutions. arXiv preprint 1rXiv:1511.07122,, 2015.
Manzil Zaheer, Guru Guruganesh. Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007. /4062, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, pp. 22360-12371 2019.

---

Under review as a conference paper at ICLR 2021
A SUDOKU PUZZLE
Sudoku is popular logic-based puzzle is given as 9 x 9 board with at least 17 givens. 'he player is required to fill each row, column, and nine 3 subregions of the 9 board with digit from to 9. The same digit may not appear twice in the same row, column or subregion. McGuire et al. (2014) proved that there doesn't exist Sudoku puzzles with less than 17 givens, making them the hardest among all the classical Sudoku puzzles easier puzzles can be obtained by providing more givens. Fig. depicts an example of an unfilled Sudoku puzzle with 17 givens and its solution obtained by the Matrix Shuffle-xxchnnge network. For training, we used Sudoku dataset with hard puzzles having 17-36 givens (Palm et al. 2018). Training data was augmented by applying validity preserving transformations transposition, stack and permuuatttooiionss
2
4
9
7
3
8
5
6
2
4
8
6
8
2
6
9
4
5
3
7
1
4
5
6
2
7
3
9
8
5
7
3
5
6
4
7
2
3
8
9
4
7
8
5
3
9
2
4
6
8
2
3
9
8
4
6
1
7
5
5
8
3
1
4
2
9
5
7
8
6
3
6
2
6
9
7
3
2
8
4
5
7
3
8
5
4
6
9
2
Figure 8: Sudoku puzzle with 17 given clues from the test set (Palm et al. obtained by the Matrix Shuffle- Extla--e network.
2018 and its solution
$3 rows and 3 subregions
63 columns and 3 subregions
12

---

Under review as a conference paper at ICLR 2021
B RECEPTIVE FIELD
The receptive field of a neural network defined as the size of the region in the input that influences single output feature. It has been proved that Beneš interconnection computer network with 4-to-4 switches and log 1 switch and 2 log - shuffle layers for an input sequence of length n² can implement any input-output permutation (Dally Towles 2004). Since the Matrix Shuffle- Exchaange network has the same interconnection structure as the Benes computer network, it has the same theoretical properties. One can check that Matrix Euuffne-Exchageeeeee has the receptive field spanning the whole input by building a quaternary tree backwards from single output element, as depicted in Fig- [ Even after the first. Cruufle-Exxchanee block, each input element participates in the construction of the feature.
Figure The receptive field for the single output feature on x input and one Benek block. Orange connections and green Switch Units depicts components that participate in the considered output feature; grey connections and Switch Units ones that don't.
13

---

Under review as a conference paper at ICLR 2021
C RESNET-29 MODEL
is known (iiverr et al. 2016 2017 2018) that convolutional neural networks, more precise ResNet can estimate moves in many grid-based board games (e.g., chess, go, shogi). Because of that, ResNet seems suitable choice for the algorithmic task on matrices. For all algorithmic tasks on matrices, we use pre-activation ResNet-29 (He al. 2016b) with 3x3 convolutional kemells and 128 feature maps. Instead of Batch Normalization (loffe & Szegedy (013)) we use Layer Normalization (Ba et al. 2016), as Batch Normalization gives bad results even on small matrices. The network consists of 14 residual blocks, that are prepended with embedding layer and 3x3 convolution. Fig. deppiccssss residual block of the ResNet-29 network. The last residual block is followed by Layer Normalization, ReLU and linear transformation. The network doesn't 't contain any bottlenecks pooling layers. Fig. depicts the full ResNet-29 structure.
Input
Embedding
Layer Norm
Residual block
3x3 Conv
Norm
Recidual
ReLU
Layer Norm
3x3 Conv
ReLu
Unnea
Ourpuu
Figure 10: Residual block used in ResNet-29.
Figure 11: ResNet-29 full architecture for algo- rithmic tasks on matrices.
14

---

Under review as a conference paper at ICLR 2021
D CONVERGENCE SPEED
Here we report Matrix Shuffle-Ecchang (Matrix-SE) with 3.5M learnble parameters and ResNet-29 (ResNet) with 4.5M learnable parameters test errоr depending on the training step for algorithmic tasks. Both models were trained using Rectified Adam optimizer with learning rate le-4 and batch size 32. Fig- 12 and report test error on matrix squaring (squaring), transpose (transpose), rotate by 90 (rotate90) degrees and bitwise XOR (bitwise XOR) tasks. Matrix Shuffle- Exeefle-Exchaggee converges significantly faster than ResNet-29 and achieves 0 error test rate on matrices of the same size as in training.
002
Bosatedo Reshet
Squaring:
Biniise XOR Matrix SE
Transpose Notrix SE
error
... Binwise XOR Reshet
... Transpose ResNet
error
10
Figure 12: Error rate depending on training step test set for matrix rotation by 90 degrees and bitwise XOR operation.
Figure 13: Error rate depending on training step
test set for matrix transpose and squaring.
15

---

Under review as a conference paper at ICLR 2021
E CIFAR-I0 10 IMAGE CLASSIFICATION
We have evaluated Matrix Shuffle-Exchange performance on CIFAR-1O image classification task. CIFAR-IO (Krizhevsky et 2009) is conventional image classification task that consists of 50000 train and 10000 test 32 x 32 images in 10 mutually exclusive classes. We use Matrix Chuffle-Exchange with 92 feature maps and 2 Beneš blocks. Between both blocks, dropout with rate 0.5 applied. The model has a total of 3.5M learnable parameters.
We augment the training data by randomly flipping the input image and train Matrix Shuffle-xxchange for 200k iterations with batch size 32. The model without additional pretraining achieves 78% test aceuracy, which is on par with a feed-forward network that is trained on a highly augmented dataset (Lin et al. 2015). But our model has somewhat worse accuracy than a typical fully-convolutional neural network, consisting of 9 convolutional layers that can achieve approximately 90% accuracy without additional pretraining (Hendrycks & Gimpel 2016). The accuracy of our model could plausibly improved using more augmentation or a larger dataset; otherwise, the model overfits.
16


Reviews:
[{"review_id": "lPoA46QwraO", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper adapts the recently introduced Neural Shuffle-Exchange (NSE) network for 2D inputs. This is done by introducing two changes; flattening 2D input using a z-order curve, and using 4-to-4 switches rather than 2-to-2. Experiments on synthetic data show that the proposed model can outperform baselines.\n\n### Strengths\n\n1. The extension of NSE to 2D domains is an interesting research direction.\n2. The use of the Z-order curve to map from 2D to 1D such that locality is approximately preserved is interesting.\n3. The model performs well for the tasks explored here.\n\n### Weaknesses\n\n1. The overall contributions are limited. While flattening 2D data using a z-order curve is interesting, it is a well known technique for spatial indexing. The use of z-order curves is also not entirely unprecedented in the field. For example, [1, 2] also use z-order curves in the context of 2D data, and [3] in the context of 3D data.\n2. The paper states that naive flattening of 2D data is expensive since it results in long sequences that are very slow to process with a NSE. The paper then proposes a technique which flattens 2D into an equally long sequence. This is very confusing. I'm assuming that reported speed improvements are due to the use of 4-to-4 rather than 2-to-2 switch, is this correct? Wouldn't this technique also be applicable to the original NSE?\n3. The experiments are limited, consisting exclusively of synthetic data and a single baseline for each task. While an experiment involving images (CIFAR-10) is presented in the Appendix, the description is brief, and the model performs no better than a feed-forward network. This is concerning as CV seems like a natural domain for exploration, especially given overlapping motivation  between this work and recent literature exploring transformers in the vision domain [4, 5]. As such, while the results are promising, they are not enough to convince me of that Matrix-SE is likely to be useful for practical  problems of interest.\n4. The presentation of technical and experimental details could be improved.\n   1. I'm confused about how weights are shared. In particular, while the figure make it clear that weights are shared among Shuffle layers in the same block (similarly for Inverse Shuffle layers), I'm unsure if they are shared between units, although I suspect they are not.\n   2. The original NSE paper pads \"the input sequence to that length by placing the sequence at a random position and adding zeros on both ends.\" Is a similar technique used in this work? If not, and weights are not shared between units (Point 4.1), I'm confused about how the model can generalize to larger matrices.\n   3. The paper does not contain information about how hyper-parameters were tuned, or information about stability of results. How robust is the model to different hyper-parameter settings? Are reported metrics averaged over multiple runs?\n   4. It should be clarified in the paper that Z-Order curves are only approximately locality preserving and contain large discontinuous jumps.\n   5. In section 5.2, the paper states that \"Graph Isomorphism Network struggles to solve these tasks.\" However, on the hardest task, GIN outperforms Matrix-SE on triangle-finding on the largest graphs.\n\n### Recommendation\n\nI recommend rejection. While I believe the extension of NSE to 2D domains is a worthwhile pursuit, I do not believe this paper represents significant progress is this regard.\n\n### Minor Issues\n\n1. Section 5.1: \"All tasks, except matrix squaring, **has** simple O(n2) algorithms.\" => \"All tasks, except matrix squaring, **have** simple O(n2) algorithms.\"\n2. Section 5.2: \"a common **practise**\" =>  \"a common **practice**\"\n3. Inconsistent spacing is used around parenthesis.\n\n### References\n\n1. Zhang, Jianjin, et al. \"Z-order recurrent neural networks for video prediction.\" 2019 IEEE International Conference on Multimedia and Expo (ICME). IEEE, 2019.\n2. Kumar Jayaraman, Pradeep, et al. \"Quadtree convolutional neural networks.\" *Proceedings of the European Conference on Computer Vision (ECCV)*. 2018.\n3. Corcoran, Thomas, et al. \"A spatial mapping algorithm with applications in deep learning-based structure classification.\" arXiv preprint arXiv:1802.02532 (2018).\n4. Parmar, Niki, et al. \"Image transformer.\" *arXiv preprint arXiv:1802.05751* (2018).\n5. Child, Rewon, et al. \"Generating long sequences with sparse transformers.\" *arXiv preprint arXiv:1904.10509* (2019).", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZBXCK-41L_6", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work proposes the Neural Shuffle-Exchange Network to capture both local and global dependencies for 2D data. The idea extends the 1D Neural Shuffle-Exchange Network to its 2D application. The proposed method first converts 2D data to 1D following the Z-order, then apply several Quaternary Switch and Quaternary Shuffle layers, and finally convert the data back to 2D space. The experimental results show that the proposed method can obtain better performance with reasonable computational cost. \n\nStrengths:\n+ The proposed method is very interesting. It studies how to capture long-term dependencies for 2D data in an efficient way. \n+ It uses the Z-order curve to flatten 2D data into 1D. It can preserve locality information for features. \n+ The experimental results show the proposed method can obtain better performance.\n\nWeaknesses:\n-  This work extends the Neural Shuffle-Exchange Network 1D to 2D. Compared with the 1D NSE, its technical contribution is incremental, and the novelty is limited.  \n-  Even though the Z-order curve can locality information, it is not clear whether the spatial information in 2D data can be preserved. It should be discussed. \n-  The attention mechanism can learn non-local dependencies effectively. It should be compared in the experimental studies. In addition, it is true that the vanilla version of attention has O(n^4) complexity but there are several improved variants with competitive performance but lower computational cost.\n- The proposed method is also applied to graph data and is compared with GIN. However, for graph data, I would suggest conducting experiments on benchmark datasets.\n\n=====Update after rebuttal=====\n\nI have read the authors' rebuttal. I still believe the novelty is limited, and hence I keep my score unchanged. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "8SOKJGKQ1lE", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The paper proposes a network architecture called Matrix Shuffle-Exchange (Matrix-SE) that can learn many logical reasoning tasks on 2D data and graph. It has complexity O(n^2 log n) for 2D input of size n x n, which is much smaller than the complexity of naive attention applied to 2D data (O(n^4)). The proposed architecture is an adaptation of the Neural Shuffle-Exchange network architecture (Freivalds et al., 2019), moving from 1D to 2D data. This adaptation is done by using a Z-order iteration of the 2D input, then performing radix-4 shuffle and radix-4 exchange, instead of radix-2. This model is shown to be able to solve several hard tasks on 2D data, such as inferring algorithms on binary matrices (transpose, rotation, bitwise XOR, matrix squaring), graph operations (component labeling, triangle finding, transitivity), and solving Sudoku puzzles. The experiments show impressive results and the model's ability to generalize to test inputs of larger sizes than those in the training set.\n\n======\n\nStrengths:\n1. Wide variety of algorithmic and logical reasoning tasks. I enjoyed reading the experiment section, and the tasks are fun and creative.\n2. Impressive generalization on larger sizes. On those tasks mentioned above, the Matrix-SE model generalizes well to 2D arrays that have larger sizes than those in the training set, out-performing baselines (ResNet, Graph Isomorphism Network).\n3. Simple model design. The generalization from Neural Shuffle-Exchange network to Matrix-SE is natural and straightforward.\n4. The paper is well-written and easy to read.\n\nWeaknesses:\n1. Lack of theoretical characterization of the model. More concretely, what kind of operations can be represented by Matrix-SE? The theoretical motivation seems to be from the classic result that Benes networks can represent any permutation. However, it's not clear how expressive the proposed model is.\n2. More realistic tasks. Do the logical reasoning tasks in more realistic scenario, such as modeling social networks represented as graphs?\n\n======\n\nOverall, I vote for accepting. The proposed model is simple, can perform logical reasoning tasks on 2D data, and generalizes well beyond sizes that are in the training set.\n\n======\n\nAdditional feedback and questions:\n- Do all the matrices in section 5.1 have binary values?\n- How is accuracy defined in Table 1 and 2. Does the output matrix have to match the label exactly? Or is it accuracy per element?\n- In Section 5.4, why is Residual-SE so much slower (9x) than Matrix-SE? I think it should only be 2x slower, because the depth of Residual-SE is twice that of the Matrix-SE.\n\n==== After rebuttal: Thank you for clarifying details. I maintain my rating. \nShowing that the inductive bias from Z-order curve and shuffle-exchange allows generalization to larger sizes is very interesting. Overall I vote for accepting.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers are split.  Two reviewers consider the technical contribution of the paper to be insufficient, and raise concerns about comparisons with Transformers or using more standard benchmarks for GNN experiments.   The other considers the experiments convincing and the method worth publishing.   My own view is that this work is not ready for inclusion in the conference.  In particular, I think this paper would be much stronger with either: 

1: a more practical task to illustrate where this method might be applied in earnest,
2: more analysis and baselines on the synthetic data.  Synthetic data can be enough for a new method if it illuminates the functioning and the benefits and drawbacks.  In this paper, we have synthetic data with little analysis, and imo (concurring with R5) insufficient baselines.  For example, while a vanilla Transformer probably could not do the matrix problems (with the matrices encoded naively), one might expect Transformers with sparse attention to do quite well on e.g. transpose and 90 degree rotation, especially given the training curriculum and proper positional embeddings; a convolutional network seems like a strawman.  I also agree with R5 that standard benchmarks for GNN exist, and these might be appropriate (or at least there should be some discussion of why they are not).
 3: some theoretical discussion of what the proposed model can do that other methods fundamentally cannot.

I do think this is interesting work, and encourage the authors to revise and resubmit.
Decision: Reject