REOrdering Patches Improves Vision Models

We sincerely thank the reviewer for their detailed and insightful review. We are thrilled that you found our paper to have a novel premise, an elegant RL formulation, and to be well-written. The weaknesses you've identified are helpful and highlight areas where we can provide deeper analysis. We agree with your assessments and address each of your points below:

ARM scan ordering and spatial continuity: ARM does not use a single scan but instead uses an ensemble of four fixed, competing scan directions and sum their outputs to form the final representation: (1) forward row-major (top-left to bottom-right), (2) backward row-major (bottom-right to top-left), (3) forward column-major (top-left to bottom-right), and (4) backward column-major (bottom-right to top-left). Figure 1 in our paper visualizes row-major and column-major scan orders.

A standard row-major patch ordering is optimal for the row-forward scan, but it is simultaneously suboptimal for the row-backward scan and the two column-based scans. This creates an inherent architectural tension. The final output of the Mamba block depends on the sum of these four scans, so the ideal input sequence is not one that is perfect for a single scan, but one that provides the best compromise for all four simultaneously.

REOrder, meanwhile, does not re-learn a basic row-major or column-major sequence. Instead, it learns a permutation of that base ordering. Therefore, we learn a data- and model-specific patch sequence to maximize the combined effectiveness of the entire multi-scan ensemble. This provides improved performance against even a row-major or column-major patch order, as seen in Figure 5.

The information-theoretic approach is a weak prior: Our exploration of compressibility was intended to see if a simple, information-theoretic heuristic could serve as a useful prior. The resulting weak and model-dependent correlation is, in itself, a key finding of our paper: it demonstrates that no simple, universal prior across a search space spanning ~10^{365} options exists, thus motivating our more powerful (but complex) reinforcement learning approach. Your point about Hilbert vs. column ordering is excellent and likely highlights the limitations of applying a 1D compression algorithm (LZMA) to sequences derived from 2D data. While Hilbert curves preserve 2D locality better, this may not translate to higher compressibility for a 1D algorithm.

JPEG compression is a good suggestion for an alternative compression strategy to explore. Unfortunately, JPEG is designed for continuous-tone 2D images. Since the sequence models we test operate on 1D sequences, JPEG is not an applicable compression scheme that we can test. In fact, JPEG on shuffled 1D sequences, such as those following a Hilbert curve, would likely result in sequences that are less compressible than row-major due to sharp patch boundaries.

Interpretability of learned orderings: We thank the reviewer for this point, which prompted us to conduct a deeper analysis of the patterns learned by REOrder. The “keyboard” example in Figure 4 suggested a link to semantic importance, but we agree a broader analysis is needed. Our new investigation reveals that REOrder learns to align the spatial structure of a dataset with the specific inductive biases of different model architectures.

Our analysis is grounded in two key observations: (1) long-sequence models are sensitive to input order, and (2) ImageNet-1K and Functional Map of the World (FMoW) exhibit a strong center bias, where the most semantically critical content is usually in the middle of the image. To investigate this, we quantified how REOrder repositions patches from the image’s central region relative to their initial ordering. This can reveal whether the learned policy prioritizes central patches by moving them to the beginning (“front-loading”) or the end (“back-loading”) of the sequence. We discovered two distinct, model-dependent strategies:

1. Mamba prefers to front-load central tokens. REOrder systematically moved central patches towards the start of the sequence (on average, a shift of +6.8 positions). This “front-loading” strategy intuitively aligns with Mamba's recurrent state-space mechanism. By processing the most salient patches first, the model establishes a strong initial state ($h_t$​) that informs the processing of the entire sequence. This allows the model to build a robust understanding of the image’s core subject from the outset.

2. Longformer and Transformer-XL prefer to back-load central tokens. Both models learned to shift central patches towards the end of the sequence (average shifts of -20.4 and -2.2 positions, respectively). Transformer-XL’s recurrent memory can summarize information from earlier, less critical patches. By positioning central patches later in the sequence, these tokens can utilize information from the entire image rather than only from half of it. Comparatively, Longformer’s sliding-window attention, which accumulates context over time, would provide greater information for patches near the end. This could maximize the contextual information available for the local and global attention mechanisms when classifying objects.

This contrast between “front-loading” for Mamba and “back-loading” for Transformer-XL/Longformer on the same ImageNet dataset demonstrates that REOrder is not finding a single, trivial ordering. Instead, it is learning sophisticated, model-specific policies that exploit how each architecture processes information. We will add this detailed analysis and discussion to the final version.

We thank the reviewer for their detailed and constructive review of our work. We appreciate that you recognize the value in how our RL-based approach can lead to positive outcomes across a diverse set of architectures and datasets. The weaknesses you’ve identified are excellent points that give us a clear pathway to improving our paper. We address these in depth below:

Global ordering limits flexibility: We agree with the reviewer that learning a single global ordering for an entire dataset is a limitation and that dynamic, per-image ordering is an exciting direction for future work. Our primary goal in this paper was to first prove conclusively that (1) patch ordering effects are real and significant in modern long-sequence models, and (2) that reinforcement learning is a viable strategy for discovering superior orderings. By learning a single, global permutation, we deliberately isolated the effect of the ordering itself, providing a clean and stable experimental setup without the confounding variance of a more complex, dynamic policy network. One way to break global ordering-associated limitations is to modify REOrder to use a small policy network that takes image patch embeddings as input and generates the ordering logits $z$ dynamically. We will update our Limitations and Future Work section accordingly.

Lack of interpretability: We thank the reviewer for this point, which prompted us to conduct a deeper analysis of the patterns learned by REOrder. The “keyboard” example in Figure 4 suggested a link to semantic importance, but we agree a broader analysis is needed. Our new investigation reveals that REOrder learns to align the spatial structure of a dataset with the specific inductive biases of different model architectures.

Our analysis is grounded in two key observations: (1) long-sequence models are sensitive to input order, and (2) ImageNet-1K and Functional Map of the World (FMoW) exhibit a strong center bias, where the most semantically critical content is usually in the middle of the image. To investigate this, we quantified how REOrder repositions patches from the image’s central region relative to their initial ordering. This can reveal whether the learned policy prioritizes central patches by moving them to the beginning (“front-loading”) or the end (“back-loading”) of the sequence. We discovered two distinct, model-dependent strategies:

1. Mamba prefers to front-load central tokens. REOrder systematically moved central patches towards the start of the sequence (on average, a shift of +6.8 positions). This “front-loading” strategy intuitively aligns with Mamba's recurrent state-space mechanism. By processing the most salient patches first, the model establishes a strong initial state ($h_t$​) that informs the processing of the entire sequence. This allows the model to build a robust understanding of the image’s core subject from the outset.

2. Longformer and Transformer-XL prefer to back-load central tokens. Both models learned to shift central patches towards the end of the sequence (average shifts of -20.4 and -2.2 positions, respectively). Transformer-XL’s recurrent memory can summarize information from earlier, less critical patches. By positioning central patches later in the sequence, these tokens can utilize information from the entire image rather than only from half of it. Comparatively, Longformer’s sliding-window attention, which accumulates context over time, would provide greater information for patches near the end. This could maximize the contextual information available for the local and global attention mechanisms when classifying objects.

This contrast between “front-loading” for Mamba and “back-loading” for Transformer-XL/Longformer on the same ImageNet dataset demonstrates that REOrder is not finding a single, trivial ordering. Instead, it is learning sophisticated, model-specific policies that exploit how each architecture processes information. We will add this detailed analysis and discussion to the final version.

Limited choice of priors: This is a great point to highlight. Since the search space of possible permutations is prohibitively large (~10^{365} for 196 patches), we provide an information-theoretic strategy such that we can start with a strong prior for our learning process. To directly test the importance of this initialization and the robustness of our RL algorithm, we have performed the ablation study you suggested. We ran additional experiments on ImageNet-1K with Transformer-XL to compare REOrder's performance when initialized with structured priors versus a completely random permutation. The Top-1 accuracy results are as follows:

• Baseline (row-major): 60.81
• REOrder with a row-major prior: 61.40
• REOrder with a column-major prior: 62.78
• REOrder with a completely random prior: 53.10

These results demonstrate that the choice of prior is critical. While our RL algorithm can explore the permutation space, its ability to find a high-performing solution is significantly hampered by a random, unstructured start. This finding validates our core approach: using an information-theoretic strategy to select a sensible prior is crucial for making the learning process tractable and effective.

We thank the reviewer for their positive assessment of our paper’s quality, clarity, and the interesting nature of the problem. We appreciate the opportunity to clarify some key aspects of our work and address the valid questions raised. The central thesis of our paper is that many modern, efficient vision models, regardless of whether they are based on attention approximations, recurrence, or state-space models, sacrifice the permutation equivariance of the original Vision Transformer. Our goal was to investigate this shared sensitivity to patch ordering across different architectural families. We deliberately chose Longformer, Transformer-XL, and Mamba not as “ViT-variants” vs. “Mamba”, but as representative examples of three distinct strategies for achieving efficiency, all of which, as we demonstrate in Section 3.3, become sensitive to input order as a result. We address your questions in depth below:

A perceived lack of improvement in performance for ViT-variants: Our results show that REOrder provides significant and consistent performance gains for Transformer-XL as well. Specifically, on ImageNet-1K, REOrder improved Transformer-XL's accuracy by 1.50% over the fixed Hilbert curve ordering and 1.09% over the spiral ordering. On the Functional Map of the World dataset, REOrder improved Transformer-XL's best-performing baseline (column-major) by an additional 1.10%. While the gains in performance for T-XL and Longformer are not as drastic as they are for Mamba, they are statistically significant across the board.

Cost of REOrder: REOrder adds as many parameters as there are patches. For a standard 224x224 image with 16x16 patches, this is 196 parameters. This increases parameter count for a ViT-Large-equivalent model by ~$2.0\times10^{-4}$%. This makes REOrder an extremely efficient method that adds negligible runtime and memory overhead to a standard supervised learning setup.

Effectiveness for Longformer: This is an important nuance that speaks to the model's architecture. As we note in our conclusion, Longformer's design, which combines a sliding-window attention with global attention tokens, is a closer approximation of full self-attention compared to Transformer-XL or Mamba. Because it was the least susceptible to patch ordering changes among the three models to begin with (as seen in Figure 2), it is expected (and we observe that) that REOrder has less room for improvement. This finding supports our overall thesis that the degree of ordering sensitivity is tied to the degree to which a model deviates from full, permutation-equivariant self-attention.

Exploring a content-aware mechanism: We appreciate this suggestion. REOrder is a content-aware method insofar that the learned patch ordering is specific to a model/dataset pair. A permutation is then sampled based on these logits. This means the final patch order is determined by the learned relevance of different parts of the image content for the specific task, making it inherently content-aware. Interesting future work can include conditioning the RL-policy on every image rather than a dataset at large. We will ensure this aspect of our method is described more clearly in the main text.

Thank you for providing a detailed and insightful review for our paper. We appreciate your recognition of the novelty of exploring the effects of for-granted architectural decisions such as patch ordering, especially as the community moves increasingly to long-sequence models. Below, we answer the questions you posed in your review, present some extra experimental details, and attempt to clarify any misconceptions that may have arisen in our writeup:

Top-1 Accuracy on ImageNet-1K: Our goal with this work was to isolate the effects of patch ordering. To do so, we simplified the training recipe to its basics and kept it static across all runs: 100 epochs, $1\times10^{-4}$ learning rate, no augmentations, and a cosine learning rate decay. For this rebuttal, we increased the training length to 300 epochs. For T-XL, this improved performance to ~62% top-1 accuracy on ImageNet-1K. For T-XL with REOrder (row-major prior), this improved performance to ~63.5%. The ordering in our results from epoch 100 to epoch 300 did not change, nor did the relative magnitude of the results.

The original ViT paper achieved 77.91% top-1 accuracy on ImageNet-1K, but their training and evaluation setup is drastically different, which explains their higher performance. Dosovitskiy et. al. first pre-train on ImageNet at 224px resolution for 300 epochs, and then fine-tune on ImageNet at 384px resolution. In comparison, our work is trained fully-supervised from random initialization for 100 epochs (300 epochs in the above experiment). The gap in performance to the ViT paper is then largely explained by pre-training and high-resolution inference. Training for longer maintains the ordering of our results.

Selection of T-XL, Longformer, and ARM (Mamba): We selected Transformer-XL, Longformer, and Mamba (specifically the ARM implementation) as they represent three distinct and influential architectural approaches to handling long-sequence data. Transformer-XL introduces a segment-level recurrence mechanism and a novel relative positional encoding scheme. This design allows the model to capture dependencies beyond a fixed length by reusing hidden states from previous segments, effectively creating a memory. Longformer addresses the quadratic complexity of the Transformer's self-attention mechanism by employing a combination of a sliding window (local) attention and a global attention mechanism. Mamba introduces a structured state-space model that operates with linear-time complexity and constant memory scaling. Each of these models were SOTA on benchmarks such as WikiText-103, One Billion Word, etc. at the time of release and are still near-SOTA today.

Our work shows that regardless of the long-sequence modeling philosophy, patch order effects have a significant impact on downstream task performance.

Improvement to Mamba and bi-direction scan: This is a good point, and we apologize that our paper was not more clear about this point.The Mamba implementation we utilize throughout our work is ARM (Autoregressive Pretraining with Mamba in Vision) [0], a model that incorporates multi-directional scanning. As we detail in our Preliminaries (Section 3.3) and Methods (Section 4), ARM is specifically designed for vision tasks and extends the base Mamba architecture significantly. Specifically, for each layer, ARM processes the input by running four parallel, causal scans in different directions (left, right, up, down) and then combines their outputs. This architecture inherently includes the bidirectional (e.g., forward and backward) and multi-axis (e.g., row-wise and column-wise) processing that you mentioned. Our work shows that even with these multi-directional scanning mechanisms, the model's performance remains highly sensitive to the initial 1-D patch ordering. We show that the challenge of patch ordering is not a relic of simpler, unidirectional models but a persistent issue in state-of-the-art architectures like ARM. REOrder is able to drastically improve performance even for these SOTA SSM architectures.

Pre-training and out-of-distribution performance: REOrder is fully-supervised with respect to the images, while the patch ordering is learned without supervision with an RL setup. The core premise of our work is that the optimal patch ordering is not universal, but is instead dependent on the specific pairing of a model’s architecture and a dataset’s intrinsic properties. REOrder is designed to discover and exploit this pairing. However, your suggestion of pre-training brings up an interesting idea for future work: in lieu of an information-theoretic prior, a model and the REOrder policy would be pre-trained together on a large-scale dataset. The resulting learned policy would then serve as a prior for natural images in future tasks. During fine-tuning on a new downstream task, both the model weights and the policy logits could be updated, allowing the learned prior patch ordering to further adapt to the specific requirements of the new task.

[0] https://arxiv.org/abs/2406.07537v1

Your kind response has helped me understand several points. Thank you very much. However, I still have some questions regarding the following points.

Top-1 Accuracy on ImageNet-1K: I believe that the additional experiments you conducted have provided evidence that supports the usefulness of your proposed method. As you explained, we should take into account the fact that fine-tuning to high resolution was not performed. However, as I understand it, ViT typically achieves accuracy in the 70% range on ImageNet-1K even with standard settings. Therefore, I remain somewhat concerned about the performance gap that still exists. I suspect this difference may be due to the lack of data augmentation (CutMix, for instance), but I do not believe that there are any factors that restrict the use of these techniques with the proposed method.

Improvement to Mamba and bidirectional scan: I apologize for my lack of understanding regarding this point. My question was to ask how specifically superior the proposed method is compared to existing Mamba-type models optimized for image tasks, such as [1] [2]. I understand that the training conditions in this study were not necessarily optimized for performance, but It is unfortunate that we were unable to quantitatively understand how much the proposed method could potentially improve SOTA performance on ImageNet-1K. On the other hand, given the additional experiments mentioned above, I understand that sufficient qualitative improvement can be expected.

In any case, the review-discussion process has deepened my understanding of this research, and I view it favorably overall. I will carefully consider the final score, taking into account discussions with other reviewers.

[1] Liu et al., VMamba: Visual State Space Model, NeurIPS 2024

[2] Zhu et al., Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model, ICML 2024

Thank you for your time and thoroughness in your review! Your feedback is incredibly valuable. We appreciate that you find our work interesting and find the writing to be well-done. Based on your review, we have conducted some additional experiments and re-organized the paper for clarity. We discuss each point in depth below:

Including more supervised learning experiments for the baselines: We previously ran these experiments but under-estimated the importance of including them in the main body of the paper. We will put these in the main body of the paper for the camera-ready version. You suggested running two experiments: (1) picking a static, random permutation once at the start of training, and (2) taking a previously-learned ordering and statically training with it. In addition to what you suggested, we added an additional experiment: (3) For every iteration during training, pick a random permutation ordering.

1. tests if any arbitrary permutation produces better or worse results than row-major or REOrder.
2. tests whether the learned permutation is what matters most, or the process of guided exploration added by REOrder
3. tests whether if unguided exploration is better than guided exploration

For the results below, we use Transformer-XL Base with all of the same parameters as described in the paper and report the top-1 accuracy on ImageNet-1K. We first repeat results from the paper for row-major and REOrder with a row-major prior, and then provide results for (1), (2), and (3).

• Row major: 60.80
• Row major with REOrder (row prior): 61.40
• (1) We pick a random permutation ahead of training and then use that for the entire learning process: 53.10
• (2) Take the best learned ordering from (b), fix that, and use that to train from the beginning: 53.07
• (3) Random permutation for every batch during training: 50.23

All together, these results add to the narrative of our paper: simply picking random patch orderings is not sufficient (and therefore, any one random ordering will also perform poorly). The structured exploration provided by REOrder is necessary to improve performance. A novel outcome of this exploration is that simply using a learned patch ordering is not sufficient for gains in performance; the guided exploration process introduced by REOrder is necessary for improved performance.

Paper clarity: We will re-organize the paper for experimental clarity. To answer your specific questions, $M_{TXL}$ is the length of the memory for T-XL, which was 128 in our experiments). We use ARM as our Mamba implementation of choice and mention this on L150 in the Preliminaries section. However, we will re-label our figures to be explicit about ARM. For all experiments, we use a patch size of 16. For visualization purposes only, Figure 6 in the Appendix shows a patch size of 32. We will be sure to address the clarity comments you highlighted in our revision.

Figure related comments: In general, our reported numbers (and Figure 2) reports percentage points. We will clarify this explicitly. We report results using multiple training runs. There are instances in which our figures have broken y-axes, such as Figure 3. These are issues caused due to rounding and will be rectified!

Top-1 Accuracy on ImageNet-1K: Our goal with this work was to isolate the effects of patch ordering. To do so, we simplified the training recipe to its basics and kept it static across all runs: 100 epochs, $1\times10^{-4}$ learning rate, no augmentations, and a cosine learning rate decay. For this rebuttal, we increased the training length to 300 epochs. For T-XL, this improved performance to ~62% top-1 accuracy on ImageNet-1K. For T-XL with REOrder (row-major prior), this improved performance to ~63.5%. The ordering in our results from epoch 100 to epoch 300 did not change, nor did the relative magnitude of the results.

The original ViT paper achieved 77.91% top-1 accuracy on ImageNet-1K, but their training and evaluation setup is drastically different, which explains their higher performance. Dosovitskiy et. al. first pre-train on ImageNet at 224px resolution for 300 epochs, and then fine-tune on ImageNet at 384px resolution. In comparison, our work is trained fully-supervised from random initialization for 100 epochs (300 epochs in the above experiment). The gap in performance to the ViT paper is then largely explained by pre-training and high-resolution inference. Training for longer maintains the ordering of our results.