Matrix Shuffle-Exchange Networks for Hard 2D Tasks

1 INTRODUCTION . Data often comes in a form of two-dimensional matrices . Neural networks are often used for processing 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 is 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 a new neural architecture , called Matrix Shuffle-Exchange network , to this end . The complexity of the proposed architecture is O ( n2 log n ) for processing n × n data matrix , which is significantly lower than O ( n4 ) if one would use the attention ( Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ) in its pure form . The architecture is derived from the Neural Shuffle-Exchange networks ( Freivalds et al. , 2019 ; Draguns et al. , 2020 ) by lifting their architecture from 1D to 2D . We validate our model on tasks with differently structured 2D input/output data . It can handle complex data inter-dependencies present in algorithmic tasks on matrices such as 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 learned an algorithm , not only fitted the training data . Our model can be used for processing graphs by representing a graph with its adjacency matrix . It 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 . It reaches the perfect accuracy on 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-the-art deep learning model which was specifically tailored for logical reasoning tasks . 2 RELATED WORK . Convolutional Neural Networks ( CNN ) are the primary tools for processing data with a 2D gridlike topology . For instance , VGG ( Simonyan & Zisserman , 2014 ) and ResNet ( He et al. , 2016a ) enable high-accuracy image classification . Convolution is an inherently local operation that limits CNN use for long-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 is crucial . The attention mechanism ( Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ) is a widespread way of solving the long-range dependency problem in sequence tasks . Unfortunately , its application to 2D data is limited due to it ’ s high O ( n4 ) time complexity for n× n input matrix . Various sparse attention mechanisms have been proposed to deal with the quadratic complexity of dense attention by attending only to a small predetermined subset of locations ( Child et al. , 2019 ; Beltagy et al. , 2020 ; Zaheer et al. , 2020 ) . Reformer ( Kitaev et al. , 2020 ) uses locality-sensitive hashing to approximate attention in time O ( n log n ) . Linformer ( Wang et al. , 2020 ) uses a 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 to 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 Wu 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 et al . ( 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 Weisfeiler-Lehman graph isomorphism test and is the most expressive among Graph Neural Network models . Neural algorithm synthesis and induction is a widely explored topic for 1D 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 a 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 et al. , 2016 ) , 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 et al. , 2018 ) , including chess ( David et al. , 2016 ) and go ( Silver et al. , 2016 ; 2017 ) , often using complex architectures or reinforcement learning . 3 1D SHUFFLE-EXCHANGE NETWORKS . Here we review the Neural Shuffle-Exchange ( NSE ) network for sequence-to-sequence processing , recently introduced by Freivalds et al . ( 2019 ) and revised by Draguns et al . ( 2020 ) . This architecture offers an efficient alternative to the attention mechanism and allows modelling of long-range dependencies in sequences of length n in O ( n log n ) time . NSE network is a neural adaption of well-known Shuffle-Exchange and Beneš interconnections networks , that allows linking any two devices using a logarithmic number of switching layers ( see Dally & Towles ( 2004 ) for an excellent introduction ) . The NSE network works for sequences of length n = 2k , where k ∈ Z+ , and it consists of alternating Switch and Shuffle layers . Although all the levels are of the same structure , a network formed of 2k− 2 Switch and Shuffle layers can learn a broad class of functions , including arbitrary permutation of elements . Such a network is called a 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 k − 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 Shuffle-Exchange block . A single Shuffle-Exchange 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 to each pair . The Residual Switch Unit ( RSU ) proposed by Draguns et al . ( 2020 ) works the best . RSU has two inputs [ i1 , i2 ] , two outputs [ o1 , o2 ] , 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 input . The second linear transformation is applied after GELU and its output is scaled by a scalar h. Additionally , the output of the unit is connected to its input using a residual connection that is scaled by a learnable parameter s. RSU is defined as follows : i = [ i1 , i2 ] g = GELU ( RMSNorm ( Zi ) ) c = Wg + b [ o1 , o2 ] = σ ( s ) i + h c In the above equations , Z , W are weight matrices of size 2m × 4m and 4m × 2m , respectively , where m is the number of feature maps ; s is a vector of size 2m and b is a bias vector – all of those are learnable parameters ; denotes element-wise vector multiplication and σ is the sigmoid function . s is initialized as σ−1 ( r ) and h is initialized as √ 1− r2 ∗ 0.25 , where r = 0.9 . The Shuffle layer performs the perfect shuffle permutation of the sequence elements . In this permutation , the destination address of an element is obtained by cyclic 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 SHUFFLE-EXCHANGE 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.reshape ) , but such a solution has two main drawbacks . Firstly , NSE requires more than 8 seconds for 1M 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 if 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-SE ) network – an adoption of NSE for 2D data , that is significantly faster , generalizes on large matrices and retains the property of receptive field spanning the whole input matrix ( see Appendix B for more details ) . The model is suitable for processing n×n input array , where n = 2k for k ∈ Z+ and each element is vector of m feature maps . Inputs that don ’ t fulfil this requirement has to be padded to the closest 2k × 2k array . The Matrix-SE network works by rearranging the input matrix into a 1D sequence , then applying several interleaved Quaternary Switch and Quaternary Shuffle layers and then converting data back to 2D , see Fig . 1 . Rearranging of n× n input to n2 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 it has recursive structure and it preserves element locality regardless of the input size.1 The resulting ordering is the same as one would get from a depth-first traversal of a 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 n2 output sequence of the last Quaternary Switch layer to n× 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-Exchange computer networks ( Dally & Towles , 2004 ) . That reduces network 1https : //en.wikipedia.org/wiki/Z-order_curve depth 2x times preserving the theoretical properties of NSE . Each Beneš block consists of interleaved k−1 QSwitch and QShuffle layers that form Shuffle-Exchange block , followed by k−1 QSwitch and Inverse QShuffle layers that form mirror counterpart of the first Shuffle-Exchange block . In contrast with NSE , we finalize Beneš block with one more QSwitch layer , making it more alike to Beneš computer networks . Last QSwitch layer serves as an output for the Beneš block and improves model expressiveness . Weight sharing is employed between QSwitch layers of same Shuffle-Exchange block . The last QSwitch layer of the Beneš block does not participate in weight sharing . If the network has more than one Beneš block , each block receives a distinct set of weights . 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 4 times shorter sequence where 4 elements are concatenated along the feature dimension as [ i1 , i2 , i3 , i4 ] . Then Quaternary Switch Unit ( QSU ) is applied to each tuple . QSU is based on RSU but has 4 inputs [ i1 , i2 , i3 , i4 ] and 4 outputs [ o1 , o2 , o3 , o4 ] . The rest of the Unit structure is left unchanged from RSU . The QSU is defined as follows : i = [ i1 , i2 , i3 , i4 ] g = GELU ( RMSNorm ( Zi ) ) c = Wg + b [ o1 , o2 , o3 , o4 ] = σ ( s ) i + h c In the above equations , Z , W are weight matrices of size 4m× 8m and 8m× 4m , respectively ; s is a vector of size 4m and b is a bias vector – all of those are learnable parameters ; h is a scalar value ; denotes element-wise vector multiplication and σ is the sigmoid function . We initialize s and h same as in RSU . Weights sharing is employed between all QSUs of the same QSwitch layer . The QShuffle layer rearranges the elements according to the cyclic digit rotation permutation . This permutation for a sequence S is defined as S [ x ] = S [ qrotate ( x , k ) ] , where qrotate ( x , k ) applies cyclic shift by one digit to x which is treated as a quaternary ( base 4 ) number and k is its length in quaternary digits.2 For the first k − 1 QShuffle layers of the Beneš block , we apply rotation to the right – to the left for the remaining k − 1 layers .

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.