Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 4: Ok but not good enough - rejection
review_content: 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. 

Strengths:
+ The proposed method is very interesting. It studies how to capture long-term dependencies for 2D data in an efficient way. 
+ It uses the Z-order curve to flatten 2D data into 1D. It can preserve locality information for features. 
+ The experimental results show the proposed method can obtain better performance.

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

=====Update after rebuttal=====

I have read the authors' rebuttal. I still believe the novelty is limited, and hence I keep my score unchanged. 
review_title: Official Blind Review #5

final_score: rating: 4: Ok but not good enough - rejection
review_content: 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.

### Strengths

1. The extension of NSE to 2D domains is an interesting research direction.
2. The use of the Z-order curve to map from 2D to 1D such that locality is approximately preserved is interesting.
3. The model performs well for the tasks explored here.

### Weaknesses

1. 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.
2. 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?
3. 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.
4. The presentation of technical and experimental details could be improved.
   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.
   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.
   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?
   4. It should be clarified in the paper that Z-Order curves are only approximately locality preserving and contain large discontinuous jumps.
   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.

### Recommendation

I 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.

### Minor Issues

1. Section 5.1: "All tasks, except matrix squaring, **has** simple O(n2) algorithms." => "All tasks, except matrix squaring, **have** simple O(n2) algorithms."
2. Section 5.2: "a common **practise**" =>  "a common **practice**"
3. Inconsistent spacing is used around parenthesis.

### References

1. Zhang, Jianjin, et al. "Z-order recurrent neural networks for video prediction." 2019 IEEE International Conference on Multimedia and Expo (ICME). IEEE, 2019.
2. Kumar Jayaraman, Pradeep, et al. "Quadtree convolutional neural networks." *Proceedings of the European Conference on Computer Vision (ECCV)*. 2018.
3. Corcoran, Thomas, et al. "A spatial mapping algorithm with applications in deep learning-based structure classification." arXiv preprint arXiv:1802.02532 (2018).
4. Parmar, Niki, et al. "Image transformer." *arXiv preprint arXiv:1802.05751* (2018).
5. Child, Rewon, et al. "Generating long sequences with sparse transformers." *arXiv preprint arXiv:1904.10509* (2019).
review_title: Interesting direction, but more work needed

final_score: rating: 8: Top 50% of accepted papers, clear accept
review_content: 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.

======

Strengths:
1. Wide variety of algorithmic and logical reasoning tasks. I enjoyed reading the experiment section, and the tasks are fun and creative.
2. 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).
3. Simple model design. The generalization from Neural Shuffle-Exchange network to Matrix-SE is natural and straightforward.
4. The paper is well-written and easy to read.

Weaknesses:
1. 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.
2. More realistic tasks. Do the logical reasoning tasks in more realistic scenario, such as modeling social networks represented as graphs?

======

Overall, 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.

======

Additional feedback and questions:
- Do all the matrices in section 5.1 have binary values?
- 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?
- 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.

==== After rebuttal: Thank you for clarifying details. I maintain my rating. 
Showing that the inductive bias from Z-order curve and shuffle-exchange allows generalization to larger sizes is very interesting. Overall I vote for accepting.

review_title: Simple model that performs well on logical reasoning tasks on 2D data, and generalizes to larger sizes than trained.

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