Tackling the Abstraction and Reasoning Corpus with Vision Transformers: the Importance of 2D Representation, Positions, and Objects

We appreciate the reviewer’s encouraging comments.

Questions

Adding caveat regarding the divergence from the ARC challenge, and that the real-world relevance is yet to be demonstrated.

We appreciate this point and will include this caveat in the updated version to clarify that ViTARC’s contributions are complementary to existing methods and that further testing on broader tasks would better demonstrate its utility (ChangeLog #4.1).

ViTARC’s contributions are intended to complement current ARC solutions and support the challenge as a whole. Since the current SOTA in ARC relies on LLM-based transduction models that handle tasks through supervised input-output transformations [1], integrating the 2D inductive bias from ViTARC could provide an orthogonal benefit. Prior studies indicate that the sequential nature of 1D methods in LLMs can limit ARC performance; for example, because the input grid is processed in raster order, LLMs experience a significant drop in success rates when horizontal movement/filling tasks are rotated 90 degrees [2].

As we will discuss in Weakness #1, the insights from ViTARC may also extend to reasoning tasks that require n x n pixel precision, with potential applications in physical-aware image or video generation in real-world scenarios.

[1] Jack Cole and Mohamed Osman. Dataset-induced meta-learning (and othe tricks): Improving model efficiency on arc. https://lab42.global/community-model-efficiency/ ,2023.

[2] Xu, Y., Li, W., Vaezipoor, P., Sanner, S., & Khalil, E. B. (2024). LLMs and the Abstraction and Reasoning Corpus: Successes, Failures, and the Importance of Object-based Representations. Transactions on Machine Learning Research (TMLR).

Explanation of Task A (6e02f1e3) in Fig. 2

Thank you for pointing this out. Task A in Fig. 2 follows a rule based on color count: if the input grid has two distinct colors, the output contains a grey diagonal from the top-left to the bottom-right. Conversely, if the input grid has three colors, the grey diagonal is from the top-right to the bottom-left. We agree that additional explanations or demonstrations for such tasks would be beneficial and have incorporated these clarifications in the updated version (ChangeLog #4.2).

Weaknesses #1: Straightforwardness of Modifications and Novelty in Existing Literature.

Treating each ARC task as an Abstract Visual Reasoning (AVR) task places us in a specific setup that combines reasoning, variable-sized grids, and a generative approach with pixel-level precision. Unlike standard vision tasks, we can’t use resizing or cropping to handle variable sizes, as ARC’s precision requirements make these methods infeasible.

ARC tasks are also inherently more complex than benchmarks like RPM or odd-one-out (o3), which only require selecting a correct answer instead of generating a complete output. VQA, by contrast, involves vision input with NLP generation. Given these differences, to the best of our knowledge, there is no prior work addressing this same combination of requirements.

Notably, the findings in ViTARC could be extended to support reasoning tasks requiring n x n pixel precision, including emerging discussions on physical reasoning in vision generation, where accurate spatial relationships are essential.

Weakness #2

Addressed in Question #1

ChangeLog 4.1: Updated the conclusion to discuss broader contributions.

ChangeLog 4.2: Added a description for Task A in Figure 2.

We thank the reviewer for the thorough and thoughtful review.

Before addressing your questions in detail, we would like to reiterate the main goal of this work. While vision transformers are being touted as a basic piece in vision and multimodal foundation models, we show that a vanilla vision transformer cannot solve ARC tasks even when trained on 1,000,000 examples. This finding indicates a clear representational deficiency in vision transformers when applied to visual reasoning. We diagnose this deficiency and address it in this work. Ultimately, the architecture and results are not limited to the ARC, as we will develop next. We focus on the ARC as a representative domain for visual reasoning with growing research interest from the community and reasonable computational resource needs that we can meet.

Weaknesses (part 1)

Many of the architectural designs in the proposed model are made for solving ARC specifically.

We acknowledge that the naming of the padding tokens <arc_pad> and border tokens <arc_endxgrid> unintentionally suggests they are tailored specifically for the ARC domain. To clarify this, we will revise the terminology in the paper accordingly (ChangeLog #3.1).

That said, the underlying implementation of 2D padding and border tokens is inherently generic and can be adapted to process any variable-sized images. While ARC demands pixel-level precision (1x1 patches), the approach has the potential to generalize to visual reasoning tasks involving n×n patches, where standard resizing or cropping techniques—typical in conventional vision tasks—are unsuitable.

ViTARC stands out as the first approach to expose the representational deficiencies in current ViT structures under these constraints and to propose necessary adaptations for addressing such challenges. We will discuss the inductive bias issue and the potential to extend this work to other visual reasoning datasets in response to Question #1.

Strong inductive biases including a priori knowledge of what the final grid shape should look like, and also what the initial grid shape looks like.

We appreciate this concern, as the core challenge of ARC indeed involves inferring the output dimensions.

The inductive biases we introduce are minimal and not specific to the ARC domain:

a) The maximum grid size, a common practice in NLP tasks where padding requires a maximum length. Notably, most vision modeling imposes stricter constraints, such as fixed image sizes after resizing or cropping.

b) The assumption that the maximum grid shape is rectangular, a typical prior in image processing.

Importantly, these biases apply only to the maximum grids. We make no assumptions about the actual input or output grid shapes, which remain unbounded—even with the introduction of border tokens and 2D padding. This is especially true since ViTARC treats each pixel as a token.

We hope these clarifications demonstrate that our injected biases are modest and consistent with standard practices in related fields. Regarding the comment, "Giving the model knowledge that one token should represent one block in ARC is a strong prior to be injecting," we are unclear about this interpretation. In ViTARC, each token represents one pixel so we assume the reviewer is referring to a pixel as a “block”? If so, we see the pixel simply as the atomic representation of ARC tasks that is the most prior-free representation we can conceive. If the reviewer was referring to a different meaning of “block”, we would appreciate it if they would consider clarifying their interpretation.

Object-Based Positional Encodings

We agree with the reviewer that, ideally, the model would learn to infer objectness implicitly through attention among pixels. However, as shown in our failure case analysis (Figure 6), the model sometimes requires a stronger signal to handle complex shapes, such as multi-colored or concave shapes. OPE serves as a demonstration of how architectural design can allow for the injection of external knowledge as “positional embeddings” within ViT models.

Moreover, OPE works synergistically with PEmixer. If OPE’s signal is unreliable, the learned vector weighting can adapt to focus more on the original x and y coordinates. We will clarify this more clearly in the updated version (ChangeLog #3.2).

Weaknesses (part 2)

No other baselines than ViT are explored, though there have been many works proposed for ARC & related reasoning tasks.

Our goal of the paper is to reveal limitations in current vision transformer architecture for performing abstract visual reasoning. Therefore, ViTARC is not yet a general ARC solver and cannot be compared directly with other ARC solvers. We aim to extend ViTARC to a general ARC solver in future work.

Furthermore, AVR models currently rely on discriminative architectures rather than generative models. To our knowledge, there is no previous model positioned within the AVR + variable-sized grid + generative model regime. Therefore, we used the vanilla ViT setup as the baseline for this comparison. We would greatly appreciate any suggestions from the reviewer for relevant models or previous work to consider for comparison.

Questions

Can you show that the model works on other complex reasoning tasks?

The diverse reasoning required across ARC tasks, ranging from spatial transformations to abstract rule discovery, already offers valuable insights into Abstract Visual Reasoning (AVR). A major strength of ViTARC lies in its dual capability as both a reasoning and generative model.

While adapting ViTARC to other visual reasoning tasks is a promising future direction, it is beyond the scope of this paper. Notably, if benchmarks like RPM or odd-one-out were reframed as generative tasks, a significant portion would align with subsets of ARC.

Additionally, our findings with ViTARC could also extend to reasoning tasks requiring n x n pixel precision, such as physical reasoning in vision generation, where spatial accuracy is equally critical. We have updated the paper to emphasize these broader contributions (ChangeLog #3.2).

Does the model work without strong priors, e.g., on input and output grid shape?

As previously described, we do not explicitly inject strong priors, such as predefined input or output grid shapes. However, if we were to remove the existing priors—such as 2D positional encoding, border tokens, and 2D padding—the setup would resemble a vanilla ViT more closely, and we expect the model’s performance would align similarly with it.

Additional Baselines Trained on the Same Amount of Data

Addressed in response to weakness 4.

ChangeLog 3.1: Updated the terminology and description of vision tokens to ensure clarity and domain-agnostic applicability.

ChangeLog 3.2: Updated the conclusion to discuss broader contributions.

ChangeLog 3.3: Clarified OPE’s role in handling complex shapes, its synergy with PEmixer, and its novelty as a method for injecting external objectness knowledge.

Weaknesses (part 2)

No Baselines or Comparisons

We thank the reviewer for their concerns regarding a lack of comparison with other ARC models. We address each point below:

• Performance on Private ARC Set: Our goal of the paper is to reveal limitations in current vision transformer architecture for performing abstract visual reasoning. Therefore, ViTARC is not yet a general ARC solver and cannot be applied directly to solve the overall ARC benchmark. We aim to extend ViTARC to a general ARC solver in future work, but resolving these critical learning deficiencies on within-task learning performance were our focus in this paper and a requisite stepping stone to future work on cross-task generalization with ViTs.

• Baseline Comparisons: AVR models currently rely on discriminative architectures rather than generative models. To our knowledge, there is no previous model positioned within the AVR + variable-sized grid + generative model regime. Therefore, we used the vanilla ViT setup as the baseline for this comparison.

• Comparison to Sequence-Based ARC Models Using LLMs: Each pixel in our setup is treated as a token, similar to many LLM-based ARC solver techniques. This design allows the vanilla ViT in our setup to function as a small-scale LLM trained from scratch, serving as a meaningful baseline to showcase the contributions of ViTARC's vision-specific architectural enhancements.

  We acknowledge that comparing ViTARC with pretrained sequence-based LLMs could be an interesting direction for future work. However, such comparisons are beyond the scope of this paper, as their performance gains may stem from scaling laws rather than architectural advancements.

Minor Weaknesses:

Placement of Figure 8: We appreciate the reviewer’s recognition of Figure 8, however, we decided to move it to the Appendix as it only contributes partly as the motivation, and we did not want to further confuse the readers from the main motivation which is the example shown in Figure 4.

Purpose of PEmixer: As noted earlier, the primary purpose of PEmixer is to control the ratio between contextual and positional information. We further opted for a vector PEmixer over a scalar for greater flexibility by allowing the model to learn when certain coordinates are more significant for a given task. For example, in ARC tasks with horizontal transformations, x-coordinates may hold greater importance.

ChangeLog 2.1: Refined Section 5 on 2D-RPE to clarify the slope design.

ChangeLog 2.2: Refined Sections 5 and 5.1 to better illustrate the different positional encodings (PEs) and provide a clearer discussion of the results.

ChangeLog 2.3: Section 4, clarified the role of border tokens in 2D padding and dynamic grids.

Weaknesses (part 1)

Unclear that this is Vision Modeling

We appreciate the reviewer’s concern that some of our adaptations to the ARC domain implemented for ViTARC make the model less "vision-specific." However, we argue that these adaptations are necessary for the ARC domain and do not shift our model away from vision modeling. Specifically:

• ARC tasks demand pixel-level precision, which makes Conv-oriented setups (e.g., local attention or patching) less suitable. By using 1x1 patches, ViTARC ensures to capture every pixel detail. This adaptation doesn’t alter ViTARC’s core as a Vision Transformer as larger patch size can be implemented for tasks with greater tolerance (e.g., n x n pixel precision).

• Unlike typical convolutional vision models where the image size and the number of patches are fixed, ViTARC must handle the generation of variable-sized grids. Since resizing or cropping is infeasible due to ARC's need for pixel-perfect precision, the End-of-Grid (EOS) token is a natural design choice. ViTARC-VT’s performance further demonstrates its importance.

• We note that without positional encodings, a transformer inherently processes and outputs unordered sets. Thus, the 2D positional encodings, including OPE, explored in ViTARC reinforce our approach as being more vision-specific, not less as the reviewer suggests.

We agree with the reviewer that in certain vision tasks, token content is more important than spatial relationships. However, in abstract visual reasoning tasks, such as those in ARC, spatial relationships sometimes hold primary importance. For instance, in tasks where same-color pixels move (e.g., ARC task #025d127b https://kts.github.io/arc-viewer/page1/#025d127b), the meaningful information lies in the positional shifts rather than in the token content itself.

Unclear Significance of Contributions

Thank you for raising concerns about disconnects between the different contributions. We’ll clarify each point below and update our paper accordingly.

• PEmixer vs APE: PEmixer is primarily used to adjust the balance between positional and token information; it’s applied as an additional layer on top of APE.

• RPE vs APE: As the reviewer pointed out, RPE was introduced to enhance spatial awareness that APE alone struggled with, as illustrated in the Figure 6 example. We acknowledge that the final encoding, which combines 2DAPE, 2DRPE, and OPE, may appear complex. We will clarify the progression and motivation behind each component in the updated version (ChangeLog #2.2).

• PEmixer and OPE without RPE worsen performance: As discussed in section 5.1, this decrease in performance is due to OPE occupying space in the concatenated encoding and reducing the dimensionality available for APE. The inclusion of RPE recovers the positional signals at the attention level, thus improving the overall performance.

We appreciate the reviewer’s concern about a lack of global understanding of the problems and the suggestions for restructuring the paper. Our paper introduces many improvements to the ViT architecture motivated by different failure analyses and aims to address very different challenges. Therefore, we structured our paper with the goal that readers can easily follow the progression of our various contributions, which we are grateful the reviewer acknowledges as part of the paper’s strength.

Necessity of Padding Contributions

We thank the reviewer for asking for additional clarification on the need for paddings. We address them in the following, and will update the paper accordingly:

Sequential padding at the end where the padding tokens are ignored by the attention should be the correct one: While we can define 2D coordinates (x,y) for the encoder, where grid size is known in advance, the decoder lacks this spatial reference unless it operates on a fixed-size 2D grid. The <arc_pad> tokens serve as placeholders to support this fixed 2D template and therefore require attention. This setup was critical in enhancing performance from ViT to ViTARC-VT.

Per-row EOS tokens should be enough and <arc endxgrid> are not needed: The need for border tokens arose after introducing 2D padding. Without them, the model would have to count tokens to determine boundaries, the precision of which quickly becomes an issue—especially in ARC tasks with dynamically defined output grid sizes (eg. task C in Figure 2). Additionally, inner grid borders differ from maximum grid boundaries, so separate tokens are used to prevent the model from having to learn this distinction implicitly. We will add these clarifications to the updated version (ChangeLog #2.3).

We thank the reviewer for the thorough and thoughtful review.

Questions:

RoPE vs RPE?

We share this curiosity and agree that RoPE (or even 2DRoPE) could be an interesting future adaptation for ViTARC. However, it wasn’t included in the current experimental setup for the following reasons: a) We prioritized PE methods with more established compatibility since RoPE is not yet a standard approach for ViTs. b) We hypothesized that RoPE’s multiplicative nature doesn’t integrate smoothly with certain components in our design, such as PEmixer, which uses additive interactions. This creates challenges in adjusting the context-to-position ratio without compromising RoPE’s key property of position-difference sensitivity through dot product alignment. We will incorporate this discussion into our paper as future work.

Encoding Directions in RPE, why not use explicit encodings for all four diagonal directions?

Since the 2D image generation occurs in raster order, the pixels in the image can be related using only 2 directions without losing the 2D nature. This is reflected using our current 2-slope approach and we will clarify this in the updated paper (ChangeLog #2.1).

However, we appreciate the reviewer’s suggestion of adding 2 additional slopes, which can potentially enhance precision by providing explicitly 2D representational capacity. We are conducting some experiments to test this four-slope approach and will aim to provide the results in a future update.

How do your architectural changes influence other vision tasks?

One of the main benefits of our proposed architectural changes is its enhanced focus on positional information, which we have found to be crucial for abstract visual reasoning tasks. This finding aligns with our intuition on vision-based tasks: spatial relations in higher dimensions carry more importance than in the 1D sequence setting. While we have not conducted experiments on other vision tasks, we hypothesize that this enhanced focus on positional information will be beneficial, especially in high-precision settings cases. We will incorporate this discussion into our paper as future works.

Starting from a Pre-trained ViT?

We agree that starting from a pre-trained ViT could be beneficial since fine-tuning a pre-trained (non-vision) transformer is currently state-of-the-art [1]. However, we did not conduct experiments using pre-trained ViTs, since the main goal of our study is to uncover and address the fundamental limitations within the current ViT setup which we believe is shown more prominently when studied in the simplest setting: a ViT trained for a single simple AVR task from scratch.

[1] Jack Cole and Mohamed Osman. Dataset-induced meta-learning (and other tricks): Improving model efficiency on arc. https://lab42.global/community-model-efficiency/ ,2023.

Did you try training jointly for all tasks?

We are very interested in running this experiment, but are currently unable to do so given our computational resource limits.

Overfitting, Underfitting, and Data Scaling.

During development, our preliminary experiments suggested that 1M samples are indeed essential for the model size we used. For example, in ARC#007bbfb7 (https://kts.github.io/arc-viewer/page1/#007bbfb7), ViTARC-VT’s performance dropped significantly when using fewer samples: with 1M training samples, the model solved 99% of test instances, but reducing to 100k samples caused test performance to drop to 0.7%. As for additional data, while we haven’t tested beyond 1M samples, we have observed that training for more epochs improves performance. For instance, in ARC#1cf80156, increasing from 1 to 2 epochs improved ViTARC-VT’s performance from 9.6% to 17.5%. This suggests that increasing the training size could likely further enhance performance. We limited our experiments to 1 epoch due to the limitation in our computational resources.

Pixel Prediction Details

For the domain of ARC, each pixel takes only 1 of 9 possible colors. Therefore in ViTARC, each pixel corresponds to a single token in the vocabulary space and is handled as a standard transformer generation task, where an exact match is required.

Follow-Up 2D RPE Experiment Results:

We conducted experiments to evaluate the reviewer’s suggestion of incorporating four diagonal directions into the relative positional encoding (RPE) design. These experiments were performed on a randomly selected subset of 50 tasks from the original 400 training tasks. The setup for the four diagonal directions was based on a 2D Manhattan distance approach, using the following slopes for the four directions:

• Top-right: $1/2$
• Top-left: $1/2^{0.25}$
• Down-right: $1/2^{0.5}$
• Down-left: $1/2^{0.75}$

Additionally, for comparison, we tested an alternative "4-direction" version that used slopes for only the up, down, left, and right directions. In this setup, pixels above and to the right received slope additions for their respective directions, divided by 2 to preserve normalization.

Results Table

Below is the performance summary for the three configurations:

Key Observations:

1. The 4-diagonal-direction RPE slightly outperformed the original 2-direction RPE in mean performance, supporting the reviewer's suggestion that explicitly modeling all four diagonal directions enhances precision.

2. While the performance gaps were not significant, both 2D Manhattan-based designs outperformed the cardinal-direction-based RPE (up, down, left, right). This indicates that incorporating diagonal relative distances provides additional spatial information, contributing positively to performance.

These findings will be incorporated into future revisions and experiments as part of our continued refinement of RPE design.

We thank the reviewer for their detailed review.

Weaknesses

Clarity on Training/Evaluation Protocol and Generalization Ability

Our main goal is to study the fundamental limitations of ViT in the simplest setting: a ViT trained for a single simple AVR task (ie. a single ARC task) from scratch. We observe that ViTs fail to learn on individual ARC tasks even when given 1,000,000 examples. Hence our primary objective in this paper is to understand why ViTs do not learn to generalize in-task and to resolve the sources of these learning deficiencies as we contributed in this paper. Addressing this critical gap on in-task ViT learning deficiencies is a requisite stepping stone for future extensions to cross-task generalization, an interesting direction we aim to explore in future work.

Novelty of Techniques and Limited Contribution

While we acknowledge that 2DAPE has been explored in prior ViT work, the effect of positional encodings has been largely underestimated due to the use of larger patch sizes, which results in fewer patches and reduces the influence of positional context per patch. For instance, the original ViT paper [1] reported minimal performance differences between 1DAPE, 2DAPE, and learned APE under these conditions. In contrast, we are the first to demonstrate that positional encodings become crucial in fine-grained reasoning settings, such as ARC with 1x1 patches, where positional context may dominate patch-level information—especially in tasks involving movement of same-color shapes (e.g. ARC#025d127b). This finding reveals a limitation in current vision reasoning approaches, as demonstrated by ViTARC’s improved performance.

The necessity for enhanced 2D representations and an initial design similar to ViTARC has been independently highlighted by François Chollet, the creator of the ARC, in a recent Twitter post (Chollet, 2024) [2].

Moreover, we believe this insight on the significance of positional encoding in visual reasoning tasks has implications beyond ARC, potentially informing fields such as physical reasoning in vision generation tasks. We have updated the paper to emphasize these broader contributions (ChangeLog #1.1).

[1] Dosovitskiy. et al. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations.

[2] Chollet, F. (2024). Twitter post. Retrieved from https://twitter.com/fchollet/status/1846263128801378616

Questions

Variance in ViTARC-VT Performance (Figure 7)

The large variance in ViTARC-VT’s performance comes from the variety of difficulties across ARC tasks. Some tasks are almost always solved perfectly, while others tend to hover around 70% accuracy, which increases overall variance. Figure 2 (left side) shows this distribution more clearly, where a strong performance on many tasks is balanced by lower scores on a smaller group. While more training could possibly reduce this, our goal here was to compare different architectures under the same conditions rather than fine-tune for each task.

Key Technique for Performance Improvement

Regarding the main driver of improvement from ViT-Vanilla to ViTARC-VT, the key technique is indeed the use of 2D padding rather than BorderTokens. Applying 2DAPE directly is challenging with ARC’s variable grid sizes. While we could adjust 2DAPE to accommodate different input grid sizes in the encoder, the output grid would still be generated relative to unknown 2D positions, making it difficult to map accurately. By using 2D padding, we create a fixed generation schema—essentially a template that maintains consistent spatial alignment—which greatly improves performance. We appreciate the reviewer highlighting this, and we will update the experiment discussion in the paper to emphasize it (ChangeLog #1.2).

Clarification on Equation (12)

Thank you for this feedback. We will update the paper to clarify that $r_{left}$ and $r_{right}$ are fixed slopes following the original ALiBi setup, with adjusted starting values. We will update the paper to include the full equation and additional description for clarity (ChangeLog #1.3).

Tuning of Hyper-parameters α and β

α and β are not manually tuned but are learned parameters. We’ve updated the paper to clarify this point (ChangeLog #1.4).

ChangeLog 1.1: Updated the conclusion to discuss broader contributions

ChangeLog 1.2: Updated Section 4.1 to highlight 2D padding as the key driver of ViTARC-VT's performance boost and explain its role.

ChangeLog 1.3: Clarified in Section 5 that $r_{left}$ and $r_{right}$​ follow a geometric sequence with different starting values, inspired by the ALiBi setup.

ChangeLog 1.4: Updated section 5 PEmixer.

Dear Reviewer,

Now that the authors have posted their rebuttal, please take a moment and check whether your concerns were addressed. At your earliest convenience, please post a response and update your review, at a minimum acknowledging that you have read your rebuttal.

Thank you, --Your AC