Original: M ATRIX S HUFFLE -E XCHANGE N ETWORKS FOR H ARD 2D T ASKS
Revised: M ATRIX S HUFFLE -E XCHANGE N ETWORKS FOR H ARD 2D T ASKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Convolutional neural networks have become the main tools for processing twodimensional data.
Revised: Convolutional neural networks have become the main tools for processing twodimensional data.

Original: They work well for images, yet convolutions have a limited receptive ﬁeld that prevents its applications to more complex 2D tasks.
Revised: They work well for images, yet convolutions have a limited receptive ﬁeld that prevents its applications to more complex 2D tasks.

Original: We propose a new neural model, called Matrix Shufﬂe-Exchange network, that can efﬁciently exploit long-range dependencies in 2D data and has comparable speed to a convolutional neural network.
Revised: We propose a new neural model, called Matrix Shufﬂe-Exchange network, that can efﬁciently exploit long-range dependencies in 2D data and has comparable speed to a convolutional neural network.

Original: It is derived from Neural Shufﬂe-Exchange network and has O (log n ) layers and O ( n 2 log n ) total time and space complexity for processing a n × n data matrix.
Revised: It is derived from Neural Shufﬂe-Exchange network and has O (log n ) layers and O ( n 2 log n ) total time and space complexity for processing a n × n data matrix.

Original: We show that the Matrix Shufﬂe-Exchange network is well-suited for algorithmic and logical reasoning tasks on matrices and dense graphs, exceeding convolutional and graph neural network baselines.
Revised: We show that the Matrix Shufﬂe-Exchange network is well-suited for algorithmic and logical reasoning tasks on matrices and dense graphs, exceeding convolutional and graph neural network baselines.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Data often comes in a form of two-dimensional matrices.
Revised: Data often comes in a form of two-dimensional matrices.

Original: Neural networks are often used for processing such data usually involving convolution as the primary processing method.
Revised: Neural networks are often used for processing such data usually involving convolution as the primary processing method.

Original: But convolutions are local, capable of analyzing only neighbouring positions in the data matrix.
Revised: But convolutions are local, capable of analyzing only neighbouring positions in the data matrix.

Original: That is good for images since the neighbouring pixels are closely related, but not sufﬁcient for data having more distant relationships.
Revised: That is good for images since the neighbouring pixels are closely related, but not sufﬁcient for data having more distant relationships.

Original: In this paper, we consider the problem how to efﬁciently process 2D data in a way that allows both local and long-range relationship modelling and propose a new neural architecture, called Matrix
Revised: In this paper, we consider the problem how to efﬁciently process 2D data in a way that allows both local and long-range relationship modelling and propose a new neural architecture, called Matrix

Original: Shufﬂe-Exchange network, to this end.
Revised: Shufﬂe-Exchange network, to this end.

Original: The complexity of the proposed architecture is O ( n 2 log n ) for processing n × n data matrix, which is signiﬁcantly lower than O ( n 4 ) if one would use the attention(Bahdanau et al., 2014; Vaswani et al., 2017) in its pure form.
Revised: The complexity of the proposed architecture is O ( n 2 log n ) for processing n × n data matrix, which is signiﬁcantly lower than O ( n 4 ) if one would use the attention (Bahdanau et al., 2014; Vaswani et al., 2017) in its pure form.

Original: The architecture is derived from the Neural Shufﬂe-Exchange networks(Freivalds et al., 2019; Draguns et al., 2020) by lifting their architecture from 1D to 2D.
Revised: The architecture is derived from the Neural Shufﬂe-Exchange networks (Freivalds et al., 2019; Draguns et al., 2020) by lifting their architecture from 1D to 2D.

Original: We validate our model on tasks with differently structured 2D input/output data.
Revised: We validate our model on tasks with differently structured 2D input/output data.

Original: It can handle complex data inter-dependencies present in algorithmic tasks on matrices such as transposition, rotation, arithmetic operations and matrix multiplication.
Revised: It can handle complex data inter-dependencies present in algorithmic tasks on matrices such as transposition, rotation, arithmetic operations and matrix multiplication.

Original: Our model reaches the perfect accuracy on test instances of the same size it was trained on and generalizes on much larger instances.
Revised: Our model reaches the perfect accuracy on test instances of the same size it was trained on and generalizes on much larger instances.

Original: In contrast, a convolutional baseline can be trained only on small instances and does not generalize.
Revised: In contrast, a convolutional baseline can be trained only on small instances and does not generalize.

Original: The generalization capability is an important measure for algorithmic tasks to say that the model has learned an algorithm, not only ﬁtted the training data.
Revised: The generalization capability is an important measure for algorithmic tasks to say that the model has learned an algorithm, not only ﬁtted the training data.

Original: Our model can be used for processing graphs by representing a graph with its adjacency matrix.
Revised: Our model can be used for processing graphs by representing a graph with its adjacency matrix.

Original: It has a signiﬁcant 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.
Revised: It has a signiﬁcant 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.

Original: We demonstrate that the proposed model can infer important local and non-local graph concepts by evaluating it on component labelling, triangle ﬁnding and transitivity tasks.
Revised: We demonstrate that the proposed model can infer important local and non-local graph concepts by evaluating it on component labelling, triangle ﬁnding and transitivity tasks.

Original: 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 ﬁnd these concepts even on small graphs.
Revised: 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 ﬁnd these concepts even on small graphs.

Original: The model can perform complex logical reasoning required to solve Sudoku puzzles.
Revised: The model can perform complex logical reasoning required to solve Sudoku puzzles.

Original: 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 speciﬁcally tailored for logical reasoning tasks.
Revised: 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 speciﬁcally tailored for logical reasoning tasks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Convolutional Neural Networks (CNN) are the primary tools for processing data with a 2D gridlike topology.
Revised: Convolutional Neural Networks (CNN) are the primary tools for processing data with a 2D gridlike topology.

Original: For instance, VGG(Simonyan & Zisserman, 2014) and ResNet(He et al., 2016a) enable high-accuracy image classiﬁcation.
Revised: For instance, VGG (Simonyan & Zisserman, 2014) and ResNet (He et al., 2016a) enable high-accuracy image classiﬁcation.

Original: Convolution is an inherently local operation that limits
Revised: Convolution is an inherently local operation that limits

Original: CNN use for long-range dependency modelling.
Revised: CNN use for long-range dependency modelling.

Original: The problem can be mitigated by using dilated (atrous) convolutions that have an expanded receptive ﬁeld.
Revised: The problem can be mitigated by using dilated (atrous) convolutions that have an expanded receptive ﬁeld.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Unfortunately, its application to 2D data is limited due to it’s high O ( n 4 ) time complexity for n × n input matrix.
Revised: Unfortunately, its application to 2D data is limited due to it’s high O ( n 4 ) time complexity for n × n input matrix.

Original: 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).
Revised: 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).

Original: Reformer (Kitaev et al., 2020) uses locality-sensitive hashing to approximate attention in time
Revised: Reformer (Kitaev et al., 2020) uses locality-sensitive hashing to approximate attention in time

Original: O ( n log n ) .
Revised: O ( n log n ) .

Original: Linformer (Wang et al., 2020) uses a linear complexity approximation to the original attention by creating a low-rank factorization of the attention matrix.
Revised: Linformer (Wang et al., 2020) uses a linear complexity approximation to the original attention by creating a low-rank factorization of the attention matrix.

Original: Sparse attention achieves great results on language modelling tasks, yet their application to complex data, where attending to entire input is required, is limited.
Revised: Sparse attention achieves great results on language modelling tasks, yet their application to complex data, where attending to entire input is required, is limited.

Original: Graph Convolutional Neural Networks (Micheli, 2009; Atwood & Towsley, 2016) generalizes the convolution operation from the grid to graph data.
Revised: Graph Convolutional Neural Networks (Micheli, 2009; Atwood & Towsley, 2016) generalizes the convolution operation from the grid to graph data.

Original: They have emerged as powerful tools for processing graphs with complex relations (see Wu et al.
Revised: They have emerged as powerful tools for processing graphs with complex relations (see Wu et al.

Original: for great reference)
Revised: for great reference)

Original: 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.
Revised: 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.

Original: Nonetheless, Xu et al.
Revised: Nonetheless, Xu et al.

Original: (2018) has shown that Graph Convolutional Networks may not distinguish some simple graph structures.
Revised: (2018) has shown that Graph Convolutional Networks may not distinguish some simple graph structures.

Original: To alleviate this problem they introduce Graph Isomorphism Network, that is as powerful as WeisfeilerLehman graph isomorphism test and is the most expressive among Graph Neural Network models.
Revised: 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.

Original: Neural algorithm synthesis and induction is a widely explored topic for 1D sequence problems(Abolaﬁa 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.
Revised: Neural algorithm synthesis and induction is a widely explored topic for 1D sequence problems (Abolaﬁa 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.

Original: Shin et al.
Revised: Shin et al.

Original: (2018) has proposed
Revised: (2018) has proposed

Original: Karel program synthesis from the input-output image pairs and execution trace.
Revised: Karel program synthesis from the input-output image pairs and execution trace.

Original: Differentiable
Revised: Differentiable

Original: Neural Computer(Graves et al., 2016), which employs external memory, has been applied to the
Revised: Neural Computer (Graves et al., 2016), which employs external memory, has been applied to the

Original: SGRDLU puzzle game, shortest path ﬁnding and traversal tasks on small synthetic graphs.
Revised: SGRDLU puzzle game, shortest path ﬁnding and traversal tasks on small synthetic graphs.

Original: 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.
Revised: 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.

Original:  1D S HUFFLE -E XCHANGE N ETWORKS
Revised:  1D S HUFFLE -E XCHANGE N ETWORKS

Original: Here we review the Neural Shufﬂe-Exchange (NSE) network for sequence-to-sequence processing, recently introduced by Freivalds et al.
Revised: Here we review the Neural Shufﬂe-Exchange (NSE) network for sequence-to-sequence processing, recently introduced by Freivalds et al.

Original: and revised by Draguns et al.
Revised: and revised by Draguns et al.

Original: This architecture offers an efﬁcient alternative to the attention mechanism and allows modelling of long-range dependencies in sequences of length n in O ( n log n ) time.
Revised: This architecture offers an efﬁcient alternative to the attention mechanism and allows modelling of long-range dependencies in sequences of length n in O ( n log n ) time.

Original: NSE network is a neural adaption of well-known Shufﬂe-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).
Revised: NSE network is a neural adaption of well-known Shufﬂe-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).

Original: The NSE network works for sequences of length n = 2 k , where k ∈ Z + , and it consists of alternating
Revised: The NSE network works for sequences of length n = 2 k , where k ∈ Z + , and it consists of alternating

Original: Switch and Shufﬂe layers.
Revised: Switch and Shufﬂe layers.

Original: Although all the levels are of the same structure, a network formed of 2 k − 2 Switch and Shufﬂe layers can learn a broad class of functions, including arbitrary permutation of elements.
Revised: Although all the levels are of the same structure, a network formed of 2 k − 2 Switch and Shufﬂe layers can learn a broad class of functions, including arbitrary permutation of elements.

Original: Such a network is called a Beneš block.
Revised: Such a network is called a Beneš block.

Original: A deeper and more expressive network may be obtained by stacking several Beneš blocks; for most tasks, two blocks are enough.
Revised: A deeper and more expressive network may be obtained by stacking several Beneš blocks; for most tasks, two blocks are enough.

Original: The ﬁrst k − 1 Switch and Shufﬂe layers of the Beneš block form Shufﬂe-Exchange block, the rest of the k − 1 Switch and Inverse Shufﬂe layers form its mirror counterpart.
Revised: The ﬁrst k − 1 Switch and Shufﬂe layers of the Beneš block form Shufﬂe-Exchange block, the rest of the k − 1 Switch and Inverse Shufﬂe layers form its mirror counterpart.

Original: In the Beneš block, only
Revised: In the Beneš block, only

Original: Switch layers have learnable parameters and weight sharing is employed between layers of the same
Revised: Switch layers have learnable parameters and weight sharing is employed between layers of the same

Original: Shufﬂe-Exchange block.
Revised: Shufﬂe-Exchange block.

Original: A single Shufﬂe-Exchange block is sufﬁcient to guarantee that any input is connected to any output through the network, i.e. the network has ‘receptive ﬁeld’ spanning the whole sequence.
Revised: A single Shufﬂe-Exchange block is sufﬁcient to guarantee that any input is connected to any output through the network, i.e. the network has ‘receptive ﬁeld’ spanning the whole sequence.

Original: 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.
Revised: 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.

Original: The Residual Switch Unit (RSU) proposed by Draguns et al.
Revised: The Residual Switch Unit (RSU) proposed by Draguns et al.

Original: (2020) works the best.
Revised: (2020) works the best.

Original: RSU has two inputs [ i 1 , i 2 ] , two outputs [ o 1 , o 2 ] , and two linear transformations on the feature dimension.
Revised: RSU has two inputs [ i 1 , i 2 ] , two outputs [ o 1 , o 2 ] , and two linear transformations on the feature dimension.

Original: After the ﬁrst transformation, root mean square layer normalization(RMSNorm)(Zhang & Sennrich, 2019) and Gaussian Error Linear Unit (GELU)(Hendrycks & Gimpel, 2016) follow.
Revised: After the ﬁrst transformation, root mean square layer normalization (RMSNorm) (Zhang & Sennrich, 2019) and Gaussian Error Linear Unit (GELU) (Hendrycks & Gimpel, 2016) follow.

Original: By default, the hidden representation has two times more feature maps than the input.
Revised: By default, the hidden representation has two times more feature maps than the input.

Original: The second linear transformation is applied after GELU and its output is scaled by a scalar h .
Revised: The second linear transformation is applied after GELU and its output is scaled by a scalar h .

Original: Additionally, the output of the unit is connected to its input using a residual connection that is scaled by a learnable parameter s .
Revised: Additionally, the output of the unit is connected to its input using a residual connection that is scaled by a learnable parameter s .

Original: RSU is deﬁned as follows:
Revised: RSU is deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: In the above equations, Z , W are weight matrices of size 2 m × 4 m and 4 m × 2 m , respectively, where m is the number of feature maps; s is a vector of size 2 m and b is a bias vector – all of those are learnable parameters; (cid:12) denotes element-wise vector multiplication and σ is the sigmoid function.
Revised: In the above equations, Z , W are weight matrices of size 2 m × 4 m and 4 m × 2 m , respectively, where m is the number of feature maps; s is a vector of size 2 m and b is a bias vector – all of those are learnable parameters; (cid:12) denotes element-wise vector multiplication and σ is the sigmoid function.

Original: s is initialized as σ − 1 ( r ) and h is initialized as √ 1 − r 2 ∗ 0 .
Revised: s is initialized as σ − 1 ( r ) and h is initialized as √ 1 − r 2 ∗ 0 .

Original: 25 , where r = 0 .
Revised: 25 , where r = 0 .

Original: The Shufﬂe layer performs the perfect shufﬂe permutation of the sequence elements.
Revised: The Shufﬂe layer performs the perfect shufﬂe permutation of the sequence elements.

Original: In this permutation, the destination address of an element is obtained by cyclic bit rotation of its source address.
Revised: In this permutation, the destination address of an element is obtained by cyclic bit rotation of its source address.

Original: The rotation to theright is used for the Shufﬂe layer and to the left for Inverse Shufﬂe layer. Shufﬂe layers do not have learnable parameters.
Revised: The rotation to the

Original:  T HE M ATRIX S HUFFLE -E XCHANGE N ETWORK
Revised:  T HE M ATRIX S HUFFLE -E XCHANGE N ETWORK

Original: The naive way to employ NSE on 2D data would be by transforming an input matrix into a sequence using raster scan ﬂattening (e.g., tf.reshape), but such a solution has two main drawbacks.
Revised: The naive way to employ NSE on 2D data would be by transforming an input matrix into a sequence using raster scan ﬂattening (e.g., tf.reshape), but such a solution has two main drawbacks.

Original: Firstly, NSE requires more than 8 seconds for 1M element sequence (which matches ﬂattened 1024x1024 matrix) in inference mode(Draguns et al., 2020), making its use unpractical on large matrices.
Revised: Firstly, NSE requires more than 8 seconds for 1M element sequence (which matches ﬂattened 1024x1024 matrix) in inference mode (Draguns et al., 2020), making its use unpractical on large matrices.

Original: Secondly, raster scan ﬂattening 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.
Revised: Secondly, raster scan ﬂattening 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.

Original: We propose Matrix Shufﬂe-Exchange (Matrix-SE) network – an adoption of NSE for 2D data, that is signiﬁcantly faster, generalizes on large matrices and retains the property of receptive ﬁeld spanning the whole input matrix (see Appendix B for more details).
Revised: We propose Matrix Shufﬂe-Exchange (Matrix-SE) network – an adoption of NSE for 2D data, that is signiﬁcantly faster, generalizes on large matrices and retains the property of receptive ﬁeld spanning the whole input matrix (see Appendix B for more details).

Original: The model is suitable for processing n × n input array, where n = 2 k for k ∈ Z +
Revised: The model is suitable for processing n × n input array, where n = 2 k for k ∈ Z +

Original: and each element is vector of m feature maps.
Revised: and each element is vector of m feature maps.

Original: Inputs that don’t fulﬁl this requirement has to be padded to the closest 2 k × 2 k array.
Revised: Inputs that don’t fulﬁl this requirement has to be padded to the closest 2 k × 2 k array.

Original: The Matrix-SE network works by rearranging the input matrix into a 1D sequence, then applying several interleaved Quaternary Switch and Quaternary Shufﬂe layers and then converting data back to 2D, see Fig.
Revised: The Matrix-SE network works by rearranging the input matrix into a 1D sequence, then applying several interleaved Quaternary Switch and Quaternary Shufﬂe layers and then converting data back to 2D, see Fig.

Original: Rearranging of n × n input to n 2 sequence has to be done in such way that enables generalization on larger matrices.
Revised: Rearranging of n × n input to n 2 sequence has to be done in such way that enables generalization on larger matrices.

Original: 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.
Revised: 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.

Original: The resulting ordering is the same as one would get from a depth-ﬁrst traversal of a quadtree.
Revised: The resulting ordering is the same as one would get from a depth-ﬁrst traversal of a quadtree.

Original: We call this transformation a Z-Order ﬂatten.
Revised: We call this transformation a Z-Order ﬂatten.

Original: The left matrix in Fig.
Revised: The left matrix in Fig.

Original: 2 shows matrix elements indexed according to Z-Order curve.
Revised: 2 shows matrix elements indexed according to Z-Order curve.

Original: The Z-Order unﬂatten transforms n 2 output sequence of the last Quaternary Switch layer to n × n representation according to Z-Order curve indexing.
Revised: The Z-Order unﬂatten transforms n 2 output sequence of the last Quaternary Switch layer to n × n representation according to Z-Order curve indexing.

Original: The middle part involving Quaternary Switch (QSwitch) and Quaternary Shufﬂe (QShufﬂe) layers is structured as one or more Beneš blocks the same way as in the NSE network.
Revised: The middle part involving Quaternary Switch (QSwitch) and Quaternary Shufﬂe (QShufﬂe) layers is structured as one or more Beneš blocks the same way as in the NSE network.

Original: The QSwitch and
Revised: The QSwitch and

Original: QShufﬂe layers differ from Switch and Shufﬂe layers in NSE, by performing the computations in groups of four instead of two.
Revised: QShufﬂe layers differ from Switch and Shufﬂe layers in NSE, by performing the computations in groups of four instead of two.

Original: Such principle (using 4-to-4 switches instead of 2-to-2) has been proved for Shufﬂe-Exchange computer networks(Dally & Towles, 2004).
Revised: Such principle (using 4-to-4 switches instead of 2-to-2) has been proved for Shufﬂe-Exchange computer networks (Dally & Towles, 2004).

Original: That reduces network depth 2x times preserving the theoretical properties of NSE.
Revised: That reduces network depth 2x times preserving the theoretical properties of NSE.

Original: Each Beneš block consists of interleaved k −  QSwitch and QShufﬂe layers that form Shufﬂe-Exchange block, followed by k − 1 QSwitch and
Revised: Each Beneš block consists of interleaved k − 1 QSwitch and QShufﬂe layers that form Shufﬂe-Exchange block, followed by k − 1 QSwitch and

Original: Inverse QShufﬂe layers that form mirror counterpart of the ﬁrst Shufﬂe-Exchange block.
Revised: Inverse QShufﬂe layers that form mirror counterpart of the ﬁrst Shufﬂe-Exchange block.

Original: In contrast with NSE, we ﬁnalize Beneš block with one more QSwitch layer, making it more alike to Beneš computer networks.
Revised: In contrast with NSE, we ﬁnalize Beneš block with one more QSwitch layer, making it more alike to Beneš computer networks.

Original: Last QSwitch layer serves as an output for the Beneš block and improves model expressiveness.
Revised: Last QSwitch layer serves as an output for the Beneš block and improves model expressiveness.

Original: Weight sharing is employed between QSwitch layers of same Shufﬂe-Exchange block.
Revised: Weight sharing is employed between QSwitch layers of same Shufﬂe-Exchange block.

Original: The last QSwitch layer of the Beneš block does not participate in weight sharing.
Revised: The last QSwitch layer of the Beneš block does not participate in weight sharing.

Original: If the network has more than one Beneš block, each block receives a distinct set of weights.
Revised: If the network has more than one Beneš block, each block receives a distinct set of weights.

Original: In the QSwitch layer, we divide adjacent elements of the sequence (corresponds to ﬂattened 2D input) into non-overlapping 4 element tuples.
Revised: In the QSwitch layer, we divide adjacent elements of the sequence (corresponds to ﬂattened 2D input) into non-overlapping 4 element tuples.

Original: That is implemented by reshaping the input sequence into a 4 times shorter sequence where 4 elements are concatenated along the feature dimension as [ i 1 , i 2 , i 3 , i 4 ] .
Revised: That is implemented by reshaping the input sequence into a 4 times shorter sequence where 4 elements are concatenated along the feature dimension as [ i 1 , i 2 , i 3 , i 4 ] .

Original: Then Quaternary Switch Unit (QSU) is applied to each tuple.
Revised: Then Quaternary Switch Unit (QSU) is applied to each tuple.

Original: QSU is based on RSU but has 4 inputs [ i 1 , i 2 , i 3 , i 4 ] and 4 outputs [ o 1 , o 2 , o 3 , o 4 ] .
Revised: QSU is based on RSU but has 4 inputs [ i 1 , i 2 , i 3 , i 4 ] and 4 outputs [ o 1 , o 2 , o 3 , o 4 ] .

Original: The rest of the Unit structure is left unchanged from RSU.
Revised: The rest of the Unit structure is left unchanged from RSU.

Original: The QSU is deﬁned as follows:
Revised: The QSU is deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: In the above equations, Z , W are weight matrices of size 4 m × 8 m and 8 m × 4 m , respectively; s is a vector of size 4 m and b is a bias vector – all of those are learnable parameters; h is a scalar value; (cid:12) denotes element-wise vector multiplication and σ is the sigmoid function.
Revised: In the above equations, Z , W are weight matrices of size 4 m × 8 m and 8 m × 4 m , respectively; s is a vector of size 4 m and b is a bias vector – all of those are learnable parameters; h is a scalar value; (cid:12) denotes element-wise vector multiplication and σ is the sigmoid function.

Original: We initialize s and h same as in RSU.
Revised: We initialize s and h same as in RSU.

Revised: Weights sharing is employed between all QSUs of the same QSwitch layer.

Original: The QShufﬂe layer rearranges the elements according to the cyclic digit rotation permutation.
Revised: The QShufﬂe layer rearranges the elements according to the cyclic digit rotation permutation.

Original: This permutation for a sequence S is deﬁned 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.
Revised: This permutation for a sequence S is deﬁned 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.

Original: For the ﬁrst k − 1 QShufﬂe layers of the Beneš block, we apply rotation to the right – to the left for the remaining k − 1 layers.
Revised: For the ﬁrst k − 1 QShufﬂe layers of the Beneš block, we apply rotation to the right – to the left for the remaining k − 1 layers.

Original:  E VALUATION
Revised:  E VALUATION

Original: We have implemented our proposed architecture in TensorFlow.
Revised: We have implemented our proposed architecture in TensorFlow.

Original: The code will be published on
Revised: The code will be published on

Original: GitHub.
Revised: GitHub.

Original: All models are trained and evaluated on single Nvidia RTX 2080 Ti (11 GB)
Revised: All models are trained and evaluated on single Nvidia RTX 2080 Ti (11 GB)

Original: GPU using
Revised: GPU using

Original: Rectiﬁed Adam (RAdam) optimizer (Liu et al., 2019) and softmax cross-entropy loss.
Revised: Rectiﬁed Adam (RAdam) optimizer (Liu et al., 2019) and softmax cross-entropy loss.

Original: We evaluate the Matrix-SE model on several 2D tasks.
Revised: We evaluate the Matrix-SE model on several 2D tasks.

Original: Our main emphasis is on hard, previously unsolved tasks: algorithmic tasks on matrices and graphs and logical reasoning task – solving Sudoku puzzles.
Revised: Our main emphasis is on hard, previously unsolved tasks: algorithmic tasks on matrices and graphs and logical reasoning task – solving Sudoku puzzles.

Original: The proposed architecture is compared to ResNet and GNN baselines which are reasonable alternatives that could be used for these tasks.
Revised: The proposed architecture is compared to ResNet and GNN baselines which are reasonable alternatives that could be used for these tasks.

Original: 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.
Revised: 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.

Original: We additionally evaluate the speed of the reported models and show that the Matrix-SE network achieves comparable or better results.
Revised: We additionally evaluate the speed of the reported models and show that the Matrix-SE network achieves comparable or better results.

Original: Model is also evaluated on CIFAR-10 dataset and the results are given in Appendix E.
Revised: Model is also evaluated on CIFAR-10 dataset and the results are given in Appendix E.

Original: 5. I NFERRING MATRIX ALGORITHMS
Revised: 5. I NFERRING MATRIX ALGORITHMS

Original: We selected four matrix tasks of varying difﬁculty – transpose, rotation by 90 degrees, element-wise
Revised: We selected four matrix tasks of varying difﬁculty – transpose, rotation by 90 degrees, element-wise

Original: XOR of two matrices, and matrix squaring.
Revised: XOR of two matrices, and matrix squaring.

Original: For the transpose and rotation tasks, we generate a random input matrix from alphabet 1-11 and expect the network to learn the required element permutation.
Revised: For the transpose and rotation tasks, we generate a random input matrix from alphabet 1-11 and expect the network to learn the required element permutation.

Original: 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.
Revised: 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.

Original: All tasks, except matrix squaring, has simple O ( n 2 ) algorithms.
Revised: All tasks, except matrix squaring, have simple O ( n 2 ) algorithms.

Original: Matrix squaring is the hardest of the selected tasks.
Revised: Matrix squaring is the hardest of the selected tasks.

Original: We generate a random binary square matrix and ask the model to output its square (matrix multiplication with itself) modulo 2.
Revised: We generate a random binary square matrix and ask the model to output its square (matrix multiplication with itself) modulo 2.

Original: That is a challenging task for the Matrix-SE model having the computation power of O ( n 2 log n ) , while the currently lowest asymptotic complexity for the matrix multiplication algorithm is O ( n 2 . 373 )
Revised: That is a challenging task for the Matrix-SE model having the computation power of O ( n 2 log n ) , while the currently lowest asymptotic complexity for the matrix multiplication algorithm is O ( n 2 . 373 )

Original: (Le Gall, 2014b).
Revised: (Le Gall, 2014b).

Original: Nonetheless, fast approximation algorithm for matrix multiplication of ˜ O ( n 2 ) complexity exists(Manne & Pal, 2014).
Revised: Nonetheless, fast approximation algorithm for matrix multiplication of ˜ O ( n 2 ) complexity exists (Manne & Pal, 2014).

Original: The lower bound for matrix multiplication for real and complex numbers in the model of arithmetic circuits is Ω( n 2 log n ) proven by Raz (2002).
Revised: The lower bound for matrix multiplication for real and complex numbers in the model of arithmetic circuits is Ω( n 2 log n ) proven by Raz (2002).

Original: Matrix squaring has the same complexity as matrix multiplication since these problems can be easily reduced one to the other.
Revised: Matrix squaring has the same complexity as matrix multiplication since these problems can be easily reduced one to the other.

Original: We selected matrix squaring for evaluation because it has only one input.
Revised: We selected matrix squaring for evaluation because it has only one input.

Original: 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).
Revised: 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).

Original: Such ResNet-29 has a receptive ﬁeld of size 59 × 59 that is enough to span the whole input on the largest matrix seen in training, making it on par with the Matrix-SE network that has a full receptive ﬁeld for any input size.
Revised: Such ResNet-29 has a receptive ﬁeld of size 59 × 59 that is enough to span the whole input on the largest matrix seen in training, making it on par with the Matrix-SE network that has a full receptive ﬁeld for any input size.

Original: Both models are trained on inputs up to size 32 × 32 for 500k steps and evaluated on size up to 1024 × 1024 .
Revised: Both models are trained on inputs up to size 32 × 32 for 500k steps and evaluated on size up to 1024 × 1024 .

Original: Matrix-SE reaches perfect accuracy on the test examples of size up to 32 × 32 for all tasks, see
Revised: Matrix-SE reaches perfect accuracy on the test examples of size up to 32 × 32 for all tasks, see

Original: Table 1.
Revised: Table 1.

Original: It generalizes on larger inputs for all tasks except matrix squaring.
Revised: It generalizes on larger inputs for all tasks except matrix squaring.

Original: ResNet-29 struggles to learn these tasks even for size 32 × 32 and doesn’t generalize.
Revised: ResNet-29 struggles to learn these tasks even for size 32 × 32 and doesn’t generalize.

Original: Matrix-SE also converges signiﬁcantly faster than ResNet-29 (see Appendix D for more details).
Revised: Matrix-SE also converges signiﬁcantly faster than ResNet-29 (see Appendix D for more details).

Original: 5. I NFERRING GRAPH ALGORITHMS
Revised: 5. I NFERRING GRAPH ALGORITHMS

Original: A graph can be easily represented by its adjacency matrix making Matrix-SE model a good ﬁt for tasks on graphs, especially for tasks where edges have additional information.
Revised: A graph can be easily represented by its adjacency matrix making Matrix-SE model a good ﬁt for tasks on graphs, especially for tasks where edges have additional information.

Original: We consider three graph algorithmic tasks: connected component labelling, transitivity, and triangle ﬁnding.
Revised: We consider three graph algorithmic tasks: connected component labelling, transitivity, and triangle ﬁnding.

Original: For the connected component labelling task, we initialize a labelled graph with random edge labels in range 2-100.
Revised: For the connected component labelling task, we initialize a labelled graph with random edge labels in range 2-100.

Original: The task is to label all edges of each connected component with the lowest label among  all the component’s edges.
Revised: The task is to label all edges of each connected component with the lowest label among [Table] all the component’s edges.

Original: For the transitivity task, the goal is to add edges to a directed graph for
Revised: For the transitivity task

Original: [Table]
Revised: [Table]

Original:  every two vertices that have a transitive path of length 2 between them.
Revised: the goal is to add edges to a directed graph for every two vertices that have a transitive path of length 2 between them.

Original: The hardest of the considered graph algorithmic tasks is the triangle ﬁnding.
Revised: The hardest of the considered graph algorithmic tasks is the triangle ﬁnding.

Original: The fastest currently known algorithm relies on matrix multiplication and has asymptotic complexity of O ( n 2 . 373 ) , where n is the vertex count (Alon et al., 1997).
Revised: The fastest currently known algorithm relies on matrix multiplication and has asymptotic complexity of O ( n 2 . 373 ) , where n is the vertex count (Alon et al., 1997).

Original: Triangle ﬁnding is a fruitful ﬁeld for quantum computing where an algorithm of quantum query complexity ˜ O
Revised: Triangle ﬁnding is a fruitful ﬁeld for quantum computing where an algorithm of quantum query complexity ˜ O

Original: ( n 5 / 4 ) (Le Gall, 2014a) has been achieved.
Revised: ( n 5 / 4 ) (Le Gall, 2014a) has been achieved.

Original: Therefore it is interesting if Matrix-SE can infer an O ( n 2 log n ) time algorithm, at least approximately.
Revised: Therefore it is interesting if Matrix-SE can infer an O ( n 2 log n ) time algorithm, at least approximately.

Original: For this task, we generate random complete bipartite graphs and add a few random edges.
Revised: For this task, we generate random complete bipartite graphs and add a few random edges.

Original: Although dense, such graphs have only a few triangles.
Revised: Although dense, such graphs have only a few triangles.

Original: The goal is to return all the edges belonging to any triangle.
Revised: The goal is to return all the edges belonging to any triangle.

Original: [Table]
Revised: [Table]

Original: 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).
Revised: 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).

Original: Since Graph Neural
Revised: Since Graph Neural

Original: Networks are a natural choice for graph task, we also compare Matrix-SE with Graph Isomorphism
Revised: Networks are a natural choice for graph task, we also compare Matrix-SE with Graph Isomorphism

Original: 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.
Revised: 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.

Original: For this purpose, we use GIN implementation from DGL library 4 with 5 Graph Isomorphism layers where a 3-layer perceptron is used at each layer.
Revised: For this purpose, we use GIN implementation from DGL library 4 with 5 Graph Isomorphism layers where a 3-layer perceptron is used at each layer.

Original: The ﬁnal graph representation is mapped to the output using another 3-layer perceptron.
Revised: The ﬁnal graph representation is mapped to the output using another 3-layer perceptron.

Original: The GIN network used in our experiments has 4.7M learnable parameters.
Revised: The GIN network used in our experiments has 4.7M learnable parameters.

Original: Graph Neural Networks are more suitable for vertex-centric tasks, where all the features are represented in the vertices.
Revised: Graph Neural Networks are more suitable for vertex-centric tasks, where all the features are represented in the vertices.

Original: To represent edge labels into GIN we follow a common practise(Yang et al., 2020; Veliˇckovi´c et al., 2017) and concatenate vertex features with element-wise multiplication between the edge and vertex features.
Revised: To represent edge labels into GIN we follow a common practice (Yang et al., 2020; Veliˇckovi´c et al., 2017) and concatenate vertex features with element-wise multiplication between the edge and vertex features.

Original: Both models are trained on graphs up to 32 vertices up to the moment where the loss hasn’t decreased for 20k steps.
Revised: Both models are trained on graphs up to 32 vertices up to the moment where the loss hasn’t decreased for 20k steps.

Original: Then models are tested on graphs with up to 1024 vertices.
Revised: Then models are tested on graphs with up to 1024 vertices.

Original: For the proposed tasks,
Revised: For the proposed tasks,

Original: Matrix-SE reaches perfect accuracy on test examples with up to 32 vertices and generalizes to moderately larger graphs.
Revised: Matrix-SE reaches perfect accuracy on test examples with up to 32 vertices and generalizes to moderately larger graphs.

Original: We ﬁnd that Graph Isomorphism Network struggles to solve these tasks even on graph sizes seen in training time.
Revised: We ﬁnd that Graph Isomorphism Network struggles to solve these tasks even on graph sizes seen in training time.

Original: Also, GIN consumes more GPU memory than Matrix-SE – the graph for the transitivity task with 1024 vertices don’t ﬁt in GPU memory.
Revised: Also, GIN consumes more GPU memory than Matrix-SE – the graph for the transitivity task with 1024 vertices don’t ﬁt in GPU memory.

Original: 5. S OLVING S UDOKU
Revised: 5. S OLVING S UDOKU

Original: Solving a Sudoku puzzle (see Appendix A for more details) requires sophisticated reasoning and can be a challenging task for a human player.
Revised: Solving a Sudoku puzzle (see Appendix A for more details) requires sophisticated reasoning and can be a challenging task for a human player.

Original: 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.
Revised: 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.

Original: The complexity of a Sudoku puzzle is characterized by how many numbers are already given in the puzzle.
Revised: The complexity of a Sudoku puzzle is characterized by how many numbers are already given in the puzzle.

Original: Park (2018) has proposed solving easy Sudoku puzzles (24-36 givens) using a convolutional neural network.
Revised: Park (2018) has proposed solving easy Sudoku puzzles (24-36 givens) using a convolutional neural network.

Original: Wang et al.
Revised: Wang et al.

Original: (2019) achieves close-to-perfect accuracy on slightly easier Sudoku puzzles (27-37 givens) by encoding each puzzle as MAXSat instance and then applying SATNet to it.
Revised: (2019) achieves close-to-perfect accuracy on slightly easier Sudoku puzzles (27-37 givens) by encoding each puzzle as MAXSat instance and then applying SATNet to it.

Original: Recurrent Relational Networks(Palm et al., 2018) can solve signiﬁcantly harder Sudokus (17givens) - they encode Sudoku puzzle as a graph and recursively employ Graph Neural Network for 32 steps.
Revised: Recurrent Relational Networks (Palm et al., 2018) can solve signiﬁcantly harder Sudokus (17givens) - they encode Sudoku puzzle as a graph and recursively employ Graph Neural Network for 32 steps.

Original: At each step, hidden representation is mapped to the solution by a 3-layer perceptron.
Revised: At each step, hidden representation is mapped to the solution by a 3-layer perceptron.

Original: Loss is obtained as the sum of a softmax cross-entropy loss of output at each step.
Revised: Loss is obtained as the sum of a softmax cross-entropy loss of output at each step.

Original: [Table]
Revised: [Table]

Original: We validate the Matrix-SE model on hard Sudoku dataset (17-34 givens) created by Palm et al.
Revised: We validate the Matrix-SE model on hard Sudoku dataset (17-34 givens) created by Palm et al.

Original: (see example from test set in Appendix A) and easy Sudoku dataset (24-36 givens) by Park (2018).
Revised: (see example from test set in Appendix A) and easy Sudoku dataset (24-36 givens) by Park (2018).

Original: Matrix-SE model is trained the same way as Recurrent Relational Network – we initiate a model with 2 Beneš blocks and 96 feature maps (3.5M total learnable parameters) and then apply it for 10 steps.
Revised: Matrix-SE model is trained the same way as Recurrent Relational Network – we initiate a model with 2 Beneš blocks and 96 feature maps (3.5M total learnable parameters) and then apply it for 10 steps.

Original: At each step, hidden representation is mapped to output by one GELU layer and one linear layer then the softmax cross-entropy loss is applied.
Revised: At each step, hidden representation is mapped to output by one GELU layer and one linear layer then the softmax cross-entropy loss is applied.

Original: We represent a Sudoku puzzle in its natural form – as a matrix – and pad it to size 16 × 16 .
Revised: We represent a Sudoku puzzle in its natural form – as a matrix – and pad it to size 16 × 16 .

Original: The model achieves 96.6% total accuracy on hard Sudoku dataset and 100% total accuracy on easy Sudoku dataset using 30 relational steps at evaluation, see Table 3.
Revised: The model achieves 96.6% total accuracy on hard Sudoku dataset and 100% total accuracy on easy Sudoku dataset using 30 relational steps at evaluation, see Table 3.

Original: 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 to a graph, signiﬁcantly simplifying the training pipeline.
Revised: 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 to a graph, signiﬁcantly simplifying the training pipeline.

Original: 5. P ERFORMANCE OF THE MODEL
Revised: 5. P ERFORMANCE OF THE MODEL

Original: Here we compare the training and evaluation speed of Matrix-SE network with other neural models.
Revised: Here we compare the training and evaluation speed of Matrix-SE network with other neural models.

Original: As a benchmark, we chose two algorithmic tasks – matrix squaring and graph transitivity task.
Revised: As a benchmark, we chose two algorithmic tasks – matrix squaring and graph transitivity task.

Original: The comparisons are done on single Nvidia RTX 2080Ti (11 Gb)
Revised: The comparisons are done on single Nvidia RTX 2080Ti (11 Gb)

Original: GPU using Tensorﬂow.
Revised: GPU using Tensorﬂow.

Original: We report an average of 300 steps on a single element batch.
Revised: We report an average of 300 steps on a single element batch.

Original: On the matrix squaring task, Matrix-SE (3.5M parameters) is compared to ResNet-29 (4.3M parameters) and Residual Shufﬂe-Exchange (RSE) network (3.5M parameters), see Fig. 3.
Revised: On the matrix squaring task, Matrix-SE (3.5M parameters) is compared to ResNet-29 (4.3M parameters) and Residual Shufﬂe-Exchange (RSE) network (3.5M parameters), see Fig. 3.

Original: Matrix-SE model works approximately 9x faster than RSE on large matrices and is on par with a convolutional neural network.
Revised: Matrix-SE model works approximately 9x faster than RSE on large matrices and is on par with a convolutional neural network.

Original: It’s worth mentioning that Matrix-SE is fast despite it has varying depth depending on the input size and its receptive ﬁeld grows to cover the whole sequence while ResNet-29 has a ﬁxed depth and receptive ﬁeld.
Revised: It’s worth mentioning that Matrix-SE is fast despite it has varying depth depending on the input size and its receptive ﬁeld grows to cover the whole sequence while ResNet-29 has a ﬁxed depth and receptive ﬁeld.

Original: To compare with graph networks, we chose the transitivity task for evaluation.
Revised: To compare with graph networks, we chose the transitivity task for evaluation.

Original: We use Matrix-SE and Graph Isomorphism Network (GIN) of conﬁgurations given in Section 5.2.
Revised: We use Matrix-SE and Graph Isomorphism Network (GIN) of conﬁgurations given in Section 5.2.

Original: As depicted in Figure 4, the Matrix-SE model is slightly faster, both during training and evaluation, and it can process 2x larger dense graphs than GIN.
Revised: As depicted in Figure 4, the Matrix-SE model is slightly faster, both during training and evaluation, and it can process 2x larger dense graphs than GIN.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original:  A BLATION STUDY
Revised:  A BLATION STUDY

Original: The most prominent feature that distinguishes the Matrix-SE from its predecessor – RSE is that it uses
Revised: The most prominent feature that distinguishes the Matrix-SE from its predecessor – RSE is that it uses

Original: Z-Order ﬂatten to transform data from a 2D matrix to 1D sequence.
Revised: Z-Order ﬂatten to transform data from a 2D matrix to 1D sequence.

Original: We study whether it is necessary by substituting Z-Order ﬂatten with raster scan ﬂatten.
Revised: We study whether it is necessary by substituting Z-Order ﬂatten with raster scan ﬂatten.

Original: We train both versions on matrix squaring, transpose and triangle ﬁnding tasks for 500k steps (for triangle ﬁnding, ﬁrst 150k are shown) on matrices of size up to 32 × 32 .
Revised: We train both versions on matrix squaring, transpose and triangle ﬁnding tasks for 500k steps (for triangle ﬁnding, ﬁrst 150k are shown) on matrices of size up to 32 × 32 .

Original: We can see in Fig.
Revised: We can see in Fig.

Original: 6 that the version with Z-Order ﬂatten generalizes to larger instances, and the other does not.
Revised: 6 that the version with Z-Order ﬂatten generalizes to larger instances, and the other does not.

Original: 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.
Revised: 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.

Original: For the matrix squaring task, we use Matrix-SE with 2 Beneš blocks and 96 feature maps as a baseline, but, for triangle ﬁnding, Matrix-SE with 2 Beneš blocks and 192 feature maps.
Revised: For the matrix squaring task, we use Matrix-SE with 2 Beneš blocks and 96 feature maps as a baseline, but, for triangle ﬁnding, Matrix-SE with 2 Beneš blocks and 192 feature maps.

Original: The results are depicted on Fig.
Revised: The results are depicted on Fig.

Original: 5, Fig.
Revised: 5, Fig.

Original: The proposed baseline models work best for the chosen tasks.
Revised: The proposed baseline models work best for the chosen tasks.

Original: A larger and deeper model in most cases yield better results, but when reasonable accuracy is obtained, diminishing improvements are observed.
Revised: A larger and deeper model in most cases yield better results, but when reasonable accuracy is obtained, diminishing improvements are observed.

Original: On the matrix squaring and triangle ﬁnding tasks Z-Order ﬂatten increase the convergence speed and accuracy compared to raster scan ﬂatten.
Revised: On the matrix squaring and triangle ﬁnding tasks Z-Order ﬂatten increase the convergence speed and accuracy compared to raster scan ﬂatten.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We have introduced the Matrix Shufﬂe-Exchange model, which is of complexity O ( n 2 log n ) and can model long-range dependencies in 2D data.
Revised: We have introduced the Matrix Shufﬂe-Exchange model, which is of complexity O ( n 2 log n ) and can model long-range dependencies in 2D data.

Original: 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.
Revised: 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.

Original: It has shown considerable logical reasoning capabilities validated by ability to solve Sudoku puzzles, where it is on par with state-of-the-art neural networks.
Revised: It has shown considerable logical reasoning capabilities validated by ability to solve Sudoku puzzles, where it is on par with state-of-the-art neural networks.

Original: This is the initial introduction of the Matrix Shufﬂe-Exchange model, and we further expect its significant improvements by ﬁne-tuning its constituents, combination with other models and application to other tasks.
Revised: This is the initial introduction of the Matrix Shufﬂe-Exchange model, and we further expect its significant improvements by ﬁne-tuning its constituents, combination with other models and application to other tasks.

Original: A S UDOKU PUZZLE
Revised: A S UDOKU PUZZLE

Original: Sudoku is a popular logic-based puzzle is given as 9 × 9 board with at least 17 givens.
Revised: Sudoku is a popular logic-based puzzle is given as 9 × 9 board with at least 17 givens.

Original: The player is required to ﬁll each row, column, and nine 3 × 3 subregions of the 9 × 9 board with a digit from 1 to 9.
Revised: The player is required to ﬁll each row, column, and nine 3 × 3 subregions of the 9 × 9 board with a digit from 1 to 9.

Original: The same digit may not appear twice in the same row, column or subregion.
Revised: The same digit may not appear twice in the same row, column or subregion.

Original: McGuire et al.
Revised: McGuire et al.

Original: (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.
Revised: (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.

Original: Fig.
Revised: Fig.

Original: 8 depicts an example of an unﬁlled Sudoku puzzle with 17 givens and its solution obtained by the Matrix Shufﬂe-Exchange network.
Revised: 8 depicts an example of an unﬁlled Sudoku puzzle with 17 givens and its solution obtained by the Matrix Shufﬂe-Exchange network.

Original: For training, we used Sudoku dataset with hard puzzles having 17-36 givens(Palm et al., 2018).
Revised: For training, we used Sudoku dataset with hard puzzles having 17-36 givens (Palm et al., 2018).

Original: Training data was augmented by applying validity preserving transformations – transposition, stack 5 and band 6 permutations.
Revised: Training data was augmented by applying validity preserving transformations – transposition, stack 5 and band 6 permutations.

Original: [Figure]
Revised: [Figure]

Original: B R ECEPTIVE FIELD
Revised: B R ECEPTIVE FIELD

Original: The receptive ﬁeld of a neural network is deﬁned as the size of the region in the input that inﬂuences a single output feature.
Revised: The receptive ﬁeld of a neural network is deﬁned as the size of the region in the input that inﬂuences a single output feature.

Original: It has been proved that Beneš interconnection computer network with 4-toswitches and 2 log n − 1 switch and 2 log n − 2 shufﬂe layers for an input sequence of length n 2 can implement any input-output permutation(Dally & Towles, 2004).
Revised: It has been proved that Beneš interconnection computer network with 4-toswitches and 2 log n − 1 switch and 2 log n − 2 shufﬂe layers for an input sequence of length n 2 can implement any input-output permutation (Dally & Towles, 2004).

Original: Since the Matrix Shufﬂe-Exchange network has the same interconnection structure as the Beneš computer network, it has the same theoretical properties.
Revised: Since the Matrix Shufﬂe-Exchange network has the same interconnection structure as the Beneš computer network, it has the same theoretical properties.

Original: One can check that Matrix Shufﬂe-Exchange has the receptive ﬁeld spanning the whole input by building a quaternary tree backwards from a single output element, as depicted in
Revised: One can check that Matrix Shufﬂe-Exchange has the receptive ﬁeld spanning the whole input by building a quaternary tree backwards from a single output element, as depicted in

Original: Fig.
Revised: Fig.

Original: Even after the ﬁrst Shufﬂe-Exchange block, each input element participates in the construction of the feature.
Revised: Even after the ﬁrst Shufﬂe-Exchange block, each input element participates in the construction of the feature.

Original: C R ES N ET -29 MODEL
Revised: C R ES N ET -29 MODEL

Original: It is known (Silver 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).
Revised: It is known (Silver 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).

Original: Because of that, ResNet seems a suitable choice for the algorithmic task on matrices.
Revised: Because of that, ResNet seems a suitable choice for the algorithmic task on matrices.

Original: For all algorithmic tasks on matrices, we use pre-activation ResNet-29(He et al., 2016b) with 3x3 convolutional kernels and 128 feature maps.
Revised: For all algorithmic tasks on matrices, we use pre-activation ResNet-29 (He et al., 2016b) with 3x3 convolutional kernels and 128 feature maps.

Original: Instead of Batch Normalization(Ioffe & Szegedy, 2015)
Revised: Instead of Batch Normalization (Ioffe & Szegedy, 2015)

Original: we use Layer Normalization(Ba et al., 2016), as Batch Normalization gives bad results even on small matrices.
Revised: we use Layer Normalization (Ba et al., 2016), as Batch Normalization gives bad results even on small matrices.

Original: The network consists of 14 residual blocks, that are prepended with embedding layer and 3x3 convolution.
Revised: The network consists of 14 residual blocks, that are prepended with embedding layer and 3x3 convolution.

Original: Fig.
Revised: Fig.

Original: 10 depicts a residual block of the ResNet-29 network.
Revised: 10 depicts a residual block of the ResNet-29 network.

Original: The last residual block is followed by Layer Normalization,
Revised: The last residual block is followed by Layer Normalization,

Original: ReLU and a linear transformation.
Revised: ReLU and a linear transformation.

Original: The network doesn’t contain any bottlenecks or pooling layers.
Revised: The network doesn’t contain any bottlenecks or pooling layers.

Original: Fig.
Revised: Fig.

Original: 11 depicts the full ResNet-29 structure.
Revised: 11 depicts the full ResNet-29 structure.

Original: D C ONVERGENCE SPEED
Revised: D C ONVERGENCE SPEED

Original: Here we report Matrix Shufﬂe-Exchange (Matrix-SE) with 3.5M learnble parameters and ResNet- (ResNet) with 4.5M learnable parameters test error depending on the training step for algorithmic tasks.
Revised: Here we report Matrix Shufﬂe-Exchange (Matrix-SE) with 3.5M learnble parameters and ResNet- (ResNet) with 4.5M learnable parameters test error depending on the training step for algorithmic tasks.

Original: Both models were trained using Rectiﬁed Adam optimizer with learning rate 1e-4 and batch size 32.
Revised: Both models were trained using Rectiﬁed Adam optimizer with learning rate 1e-4 and batch size 32.

Original: Fig.
Revised: Fig.

Original: 12 and 13 report test error on matrix squaring (squaring), transpose (transpose), rotate by 90 (rotate90) degrees and bitwise XOR (bitwise XOR) tasks.
Revised: 12 and 13 report test error on matrix squaring (squaring), transpose (transpose), rotate by 90 (rotate90) degrees and bitwise XOR (bitwise XOR) tasks.

Original: Matrix Shufﬂe-Exchange converges signiﬁcantly faster than ResNet-29 and achieves 0 error test rate on matrices of the same size as in training.
Revised: Matrix Shufﬂe-Exchange converges signiﬁcantly faster than ResNet-29 and achieves 0 error test rate on matrices of the same size as in training.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: E CIFAR-10 IMAGE CLASSIFICATION
Revised: E CIFAR-10 IMAGE CLASSIFICATION

Original: We have evaluated Matrix Shufﬂe-Exchange performance on CIFAR-10 image classiﬁcation task.
Revised: We have evaluated Matrix Shufﬂe-Exchange performance on CIFAR-10 image classiﬁcation task.

Original: CIFAR-10(Krizhevsky et al., 2009) is a conventional image classiﬁcation task that consists oftrain and 10000 test 32 × 32 images in 10 mutually exclusive classes.
Revised: CIFAR-10 (Krizhevsky et al., 2009) is a conventional image classiﬁcation task that consists oftrain and 10000 test 32 × 32 images in 10 mutually exclusive classes.

Original: We use Matrix Shufﬂe-Exchange with 92 feature maps and 2 Beneš blocks.
Revised: We use Matrix Shufﬂe-Exchange with 92 feature maps and 2 Beneš blocks.

Original: Between both blocks, dropout with rate 0.5 is applied.
Revised: Between both blocks, dropout with rate 0.5 is applied.

Original: The model has a total of 3.5M learnable parameters.
Revised: The model has a total of 3.5M learnable parameters.

Original: We augment the training data by randomly ﬂipping the input image and train Matrix Shufﬂe-Exchange for 200k iterations with batch size 32.
Revised: We augment the training data by randomly ﬂipping the input image and train Matrix Shufﬂe-Exchange for 200k iterations with batch size 32.

Original: The model without additional pretraining achieves 78% test accuracy, which is on par with a feed-forward network that is trained on a highly augmented dataset(Lin et al., 2015).
Revised: The model without additional pretraining achieves 78% test accuracy, which is on par with a feed-forward network that is trained on a highly augmented dataset (Lin et al., 2015).

Original: 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).
Revised: 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).

Original: The accuracy of our model could plausibly be improved using more augmentation or a larger dataset; otherwise, the model overﬁts.
Revised: The accuracy of our model could plausibly be improved using more augmentation or a larger dataset; otherwise, the model overﬁts.