Visual Scratchpads: Enabling Global Reasoning in Vision

Meaning of globality and globality degree. We first need to clarify the meaning of globality. The definition of globality is independent of the model. Globality is the minimum number of tokens to get non-trivial information about the output. This is stronger than just saying that the function requires a global view of the input, it requires a global view to even get a non-trivial partial information. In the ImageNet example, seeing whiskers in one patch increases the probability for both cat and walrus classes, however, one cannot be certain yet. Hence, the whiskers patch already provides weak learning (i.e., significant information on the label) but not strong learning. This is enough for the task to be of low globality. Our tasks are specifically designed to be hard (e.g., spurious correlations are removed). For instance, it is highly unlikely that one can get some signal about the maze by just looking at one patch (if there were some shortcuts like that, models would have been to weakly learn the task without scratchpads).

Consider a task learned by Transformers such that we have input tokens $X\_1, \\ldots, X\_n$ and the output is $Y$ (e.g. the label). In this case, the globality of the task (given by the distribution of input and output) is well-defined as the minimum number of tokens $k$ such that there are $k$ tokens, $X\_{i\_1}, \\ldots, X\_{i\_k}$, that along with the histogram of the tokens ($\\hat P\_X$) provide non-trivial information on the output. I.e., I(X\_{i\_1}, \\ldots, X\_{i\_k}, \\hat P\_X; Y) \= n^{-O\_n(1)}.
We clarify that the globality degree is indeed defined for the distribution of (X, Y) and not for specific samples (similar to how learning is defined for the distribution of inputs and outputs; while each example may be predicted correctly or not). We note that it has been conjectured and proven in certain cases in $Abbe et al. 2024$ that if globality is not constant (in $n$) then the task is not learnable in $poly(n)$ time by a $poly(n)$-size Transformer.

Now, in the visual domain, the number of tokens (number of patches) is often constant in contrast to text. (Although one could for instance consider mazes of increasing size and resolution). When the number of patches, $n$, is constant, one cannot define terms asymptotics $poly(n), n^{-O\_n(1)}$ properly. Nevertheless, from the asymptotic results in the general case, we know that as tasks get more global, larger models and more training iterations are required to learn them. We believe in practice, one can easily compare the globality of tasks by just inspecting them, or by running some simple experiments to see how masking patches of the images affects learning, as we compared ImageNet to our cycles task.
This is why in Definition 1, we add a threshold $\\alpha$ to our definition. For a fixed $\\alpha$, one can always compare the globality degrees of two tasks.

Are there parts of the definitions of globality and tasks that are still unclear?

Role of supervision. So there are two comparisons here.
(1) Comparing the no scratchpad model to the model with a single-frame scratchpad, the models are almost identical and their expressive power is the same. The improvements are indeed coming from having more supervision as the task is broken into subtasks that are easier (have lower globalities). The comparison methodology here is similar to the original scratchpad paper $Nye et al. 2021$ .
(2) Comparing the single-frame model with the inductive model; here it is not only the amount of supervision that matters but the way the supervision is fed in, i.e., the inductive format of the scratchpad. We have already explored this in Appendix E.4. Here we consider a simple model that predicts all frames in parallel, in contrast to our inductive model. We show that this model has better in-distribution performance compared to the single-frame model which shows the role of more granular supervision. Nevertheless, we show that the performance of the inductive scratchpad on both in-distribution and especially OOD samples is significantly higher, which shows the importance of making the supervision in a more compositional manner.

For some tasks, the intermediate frames may not be available. We note that this is not a new challenge specific to our work and it applies to all scratchpad methodologies $Nye et al. 2021$ . For example, when people supervise math questions with both detailed solution steps and the final answer, one can wonder what happens when the solution steps are not available. In general, one should try to supervise the scratchpad at least for a portion of the data. Note that scratchpad data often exists naturally, e.g., if one takes videos from the web.

We thank the reviewer for bringing these works to our attention.

Yang et al. (2024) and Bai et al. (2023) consider video as a unified interface for diverse tasks, leveraging frame-by-frame generation for broader real-world decision-making tasks like robotics and self-driving cars​. They also show some initial reasoning capabilities of such models. There are indeed similarities on the modeling front, as the frames of the scratchpad in our case can be viewed as frames of a video too.
Our contributions are however different from these prior works.

1. We propose (high-globality) image classification tasks that are impossible to weakly learn for the current models (even by large models trained on extensive data) and expose fundamental limitations of vision models.
2. We show that these tasks are only learnable if a visual scratchpad is used. On the scratchpad front, we study different models, single-frame, multi-frame baseline, the inductive model (with different variations), and we show how different modeling assumptions affect in- and out-of-distribution generalization performances. The whole study on the OOD part is new as well.
3. We extend theoretical concepts like globality degree and provide experimental evidence that directly connects task globality to model performance.

In short, the necessity of visual scratchpads for certain image classification tasks and the reason for this requirement has been absent from the prior work. The idea that we need scratchpads in the vision domain has not been proposed in these works.

The dataset of Cherian et al. also considers vision (alongside language) tasks that require reasoning, including some graph tasks. However, we note that, with having globality in mind, our tasks are designed in a way that there is no shortcut or spurious correlation that the models can use. Thus, the globality of our tasks is high. For our tasks, large models trained with extensive data cannot even achieve weak learning. However, in the work of Cherian et al. we can see that models achieve modest performance in different task categories. Further, our tasks all have adjustable difficulties. Also, our tasks are simple image classification while Cherian et al. is more similar to a VQA dataset.

OOD on bigger mazes. We agree that the idea of handling bigger mazes is interesting. There are two possible directions of approaching this: (1) One idea is to keep the number of patches constant and resize larger mazes (higher resolution images). In this case, smaller mazes will be shown larger (zoomed-in) and vice versa. Since we are training these models from scratch, it is highly unlikely that the model achieves good performance on differently sized mazes. However, if one uses pre-trained models and also tries OODs like (train sizes: 10, 12, 14, 16 and test: 18, 20) one could potentially end up with robust enough models capable of generalizing. (2) Not resizing and instead increasing the number of tokens (similar to Transformers for text). Disregarding the fact that this would increase the FLOPs of the model, making comparisons unreliable, one also has to handle the length generalization similar to text by using suitable positional embeddings, etc. We believe both of these directions are out of the scope of this work. In this work, we are showing that the inductive method is capable of generalizing to harder problems of a given size. Our approach captures what would happen with real-world datasets. Say, if we had a geometry dataset, the resolution and image size would have been almost similar throughout the dataset, however, difficulties would vary. This is what we have done for OOD experiments as well.

Regardless of these design choices, we do not believe these analyses to be relevant for the paper, as we want to isolate algorithmic reasoning, and therefore we design a setting where the atomic building blocks do not change much as that could introduce additional distribution shifts that could act as confounding factors.

Why pretraining is helpful in Figure 4? Pretraining is essentially providing a better initialization compared to the model from scratch, which may help optimization. When initialized randomly, the gradient of the neural network is essentially uncorrelated with the high globality target, which makes learning hard. However this may change when we use a pre-trained model and the weights are not random (i.e., pre-training provides a better initialization.) Moreover, much as the prior we add to the model by using BFS algorithms, pre-training provides human priors from the internet-sourced data the model was trained on. For instance, the fact that adjacent patches are likely more related than patches that are far away from each other, or the concept of a “graph” or a “maze”.

Effective depth, alternative explanations and locality of steps

We believe there to be no meaningful connection between the number of reasoning steps and the required minimum depth of Transformers; we further explain this. First, note that the Cycles task of size 24 is solved with 6 steps of BFS (each small cycle is 12 nodes, expand 2 adjacent nodes at a time). Similarly, the Strings task of size 20 is solvable in 5 steps. Given Figure 5,

1. For cycles(24), without scratchpad, ViT-huge cannot solve the task while having depth 32. However, when using a single-frame scratchpad, it can solve it. Note that the expressive power of these two models is the same. So it’s the supervision of the model that makes it learn the task and not the depth of the model.
2. For strings(20), even ViT-huge cannot solve the task, despite having 6x layers compared to the number of reasoning steps.

To gain more context on this, note that $Abbe et al. 2024$ has proved formally that no Transformer with polynomial size (irrespective of depth) can learn the cycle task in the text domain in poly-time. No properties of effective depth will change this outcome.

If the reviewer meant something different from what we argue here, we invite them to clarify. We are happy to improve our explanations in the paper if more feedback is provided.

Connection to Existing Work

We appreciate the reviewer's suggestions for related work. Below, we clarify our contributions and how they align with prior research:

1. While CoT has been explored for textual tasks (e.g., “Why Think Step by Step?”), our work adapts this to visual tasks, requiring unique considerations for 2D spatial data and pixel-wise reasoning.
2. Existing works like ViperGPT and Visual Programming focus on compositional reasoning or symbolic tasks, whereas we introduce interpretable benchmarks that test fundamental reasoning abilities from scratch. ViperGPT and Visual Programming papers use language models while we work in a vision-only setting.

We will incorporate these references into the revised version to contextualize our contributions better.

Benchmark

Our benchmarks are designed to address a gap in the field: existing literature often tests visual tasks relying on local features, leaving high-globality tasks unexplored. Key contributions of our benchmarks:

1. High-globality visual tasks that challenge modern vision models.
2. A platform to evaluate scratchpad techniques for visual reasoning.
3. Insights into inductive reasoning and OOD generalization in the visual setting

We are not aware of any other visual benchmark that is not learnable for large Transformers with a large training set. Note that the complexities of the tasks are adjustable in our paper.

Globality degree estimation

Estimating globality analytically for tasks like BFS is challenging due to dependencies on image resolution and patch size. Instead, we empirically demonstrate the high globality of our tasks. Models trained on ImageNet succeed with 90% masking, while even 30% masking makes Cycles unlearnable (Figure 3).
This empirical approach avoids computing mutual information estimates while clearly showing the globality of the tasks. (Note that one has to compute the mutual information of every $k$ token with the label, given the definition. This results in a combinatorially large number of mutual information terms to be estimated.)

Response to questions

1. The masking experiment (Figure 3) demonstrates that ImageNet is a low-globality task, as it is solvable even with significant patch masking. In contrast, our tasks fail with masking, underscoring their high globality.
2. Our work shows (a) the limitations of current vision models on high-globality tasks, (b) the necessity of visual scratchpads for learning such tasks in-distribution, (c) the importance of inductive scratchpads for OOD generalization.

Please let us know if any part is still unclear. We believe we have answered most of your concerns and we hope you consider increasing your score accordingly.

Globality vs. experiments

We kindly ask the reviewer to clarify the statement “My main concern is that the experiments don't really support the main story of the paper. They seem related at a glance, but after thinking about it some more I don't think they actually are.”.

We believe that the experiments not only support the claims of the paper but they are clearly aligned with it. So we think there is a misunderstanding about the paper or about the reviewer’s comment. We add clarifications below. Hopefully, this clarifies the alignment, otherwise please provide more precise feedback. We are open to discuss and improve the writing of the paper if it happens to lead to misunderstandings.

Clarification of globality and relation to the experiments:
First of all, we did not define the globality as the “proportion of image patches required to predict the label”. It is more subtle: the number of patches required to get non-trivial information about the label. As a result, high globality does not just say that the task requires many patches to be predicted, but that one in fact needs many patches to even get some non-trivial information about the label. For instance, “counting” is not high globality, despite requiring many patches to fully predict.

Below we explain why we believe the experiments align with the globality story:

• No scratchpad models struggle with high-globality tasks because the problem remains highly global and the model does not get to learn anything meaningful. Our experiments are in agreement with this.
• Single-frame scratchpads are shown to break the globality barrier and our experiments in fact show that the model can suddenly learn in these settings. It is expected that extra supervision helps with the learning, but our story explains exactly why (with the globality measure) and also shows that fairly weak scratchpads can already break the globality (with the staircase phenomenon detailed both analytically and experimentally).
• Inductive scratchpads break globality like the single-frame scratchpad, again in agreement with our story and the globality measure, but have also an additional property: they are more dynamic and can length-generalize on the cycle length. I.e., they offer both in-distribution and OOD generalization in contrast to the single-frame scratchpad that only gets in-distribution generalization.

The relationship between globality and the experiments is clear: tasks with high globality are hard for models to learn and generalize, necessitating the use of visual scratchpads to reduce learning complexity. Further, we specify that hardness precisely: A model of polynomial size would require exponential sample complexity (and thus training time) to learn a non-constant globality degree task.

On transformers and global attention: Transformers, despite their global attention mechanisms, do not inherently solve high-globality tasks. This is in fact an important point: it is not sufficient to have global attention to successfully learn with high globality. For example:

• In the absence of scratchpads, models fail to learn high-globality tasks like Cycles and Strings
• The Strings task, with as few as 5 steps, cannot be solved by deep models like ViT-H, which have 32 layers making it clear that even a large number of global attention operations is insufficient. Note that this is true even when using the single-frame scratchpad.
• The same observation has been made in the text domain. For example, it’s been shown that learning parities (a symmetric global function) using Transformers is hard. $1,2$

$1$ Theoretical limitations of self-attention in neural sequence models, Michael Hahn, 2020

$2$ Why are sensitive functions hard for transformers, Michael Hahn and Mark Rofin, 2024

TL;DR

Thanks to the authors for the technical clarifications about the experiments -- I found the discussion helpful, and I largely accept the technical points about the experimental results and globality degree. Therefore I would like to check with the authors about what they see as the main claim these experiments should support.

I currently still feel that the experiments are not representative enough of computer vision tasks to substantiate the claim this is a “study of locality and globality in computer vision”. I appreciate the chance to clarify the statement, and I also appreciate the authors’ very constructive discussion. I hope the suggestions about alternative benchmarks below provide a similarly constructive discussion to the authors and reviewers.

Checking understanding about the claims of the paper: bringing [Abbe et al. 2024] to a vision setting

We believe that the experiments not only support the claims of the paper but they are clearly aligned with it. So we think there is a misunderstanding about the paper

My understanding is that the contribution of this paper is in applying the results of [Abbe et al. 2024] but now in a vision context. While [Abbe et al. 2024] introduced the theory, proofs, and showed supporting experiments; those supporting experiments are in a textual setting. This submission is mainly experimental, and shows experiments in the image domain.

In particular, the experiments aim to show that for vision, too, scratchpads provides additional process supervision and learning curriculum that makes the optimization tractable for high-globality-degree tasks.

Do the authors feel that characterization is a fair one?

Why I believe the experiments need to be representative of a real-world vision setting

[Abbe et al. 2024] contributed both the definition of globality degree and the formal proofs. Therefore, I believe that the additional contribution here is in extending that work to the image domain. The paper defines some of the contribuitons as “Enabling Global Reasoning in Vision” (title, L0) and as a “study of locality and globality in computer vision” (L95). To make this study more useful, I believe the study's experiments should in settings that are representative of typical computer vision.

Why the current experiments are not representative of a vision setting

As the authors say, they are not aware of any "visual benchmark that is not learnable for large Transformers with a large training set", hence the need for the new benchmark. I appreciate that the authors showed results using vision transformers (ViTs with inputs of shape B x H x W x C), but the benchmark images themselves are very unlike natural images used for typical computer vision

• This benchmark seems like an image analogy of the benchmark of some of the experiments used in [Abbe et al 24], which “learned BFS on text tokens”. Instead of using text and a transformer decoder, these experiments use images + a transformer encoder (ViT). Essentially, these networks must “learn BFS from image patches”, and these the images are different visual depictions of graphs: node/edges, strings, mazes. While the depiction of the task is in a 2D tensors, these images are missing elements of motion blur, heavy occlusion, etc that are characteristic of real-world settings.

The current benchmark doesn’t feel like a problem that requires vision or is inherently visual. But I believe there is probably some vision task requires inductive reasoning?

Vision settings that might be amenable to scratchpads

My recommendation is that the paper would be stronger with experiments that use natural images. Perhaps there are other types of image tasks that do have a high globality degree / require inductive reasoning reasoning on images.

For example, the authors could consider tasks that are currently solved with RL from pixels, that have sparse reward (1/0). For example: low-level behaviors like manipulation (e.g. Maniskill) or high-level ones like in Behavior-1K.

Training models with sparse rewards is difficult because the signal in the search problem is very hard. The type of process supervision suggested here would reduce the problem to behavior cloning, and the specific architectural technique proposed here would be a novel approach.

Novelty of scratchpads: The concept of visual scratchpads has some resemblance to the work by Hsu et al. (2023), and we will cite it in the revised manuscript. However, there are crucial differences in scope and implementation:

• Domain: Hsu et al. focus on augmenting language model reasoning through external tools that generate visual diagrams for text-based tasks. In contrast, our work centers on vision-only tasks requiring global reasoning, such as solving mazes and analyzing graph connectivity (we target hard visual tasks in contrast to this work.)
• Scratchpad implementation: In Hsu et al., the visual diagram is created externally by tools. Conversely, in our framework, the model itself generates the scratchpad, breaking down global problems into intermediate visual reasoning steps inside the model. This includes the use of autoregressive, multi-frame predictions, a novel approach not explored in Hsu et al.
• Task globality: The tasks we introduce (Cycles, Strings, and Mazes) are different in nature, as they are high-globality visual benchmarks, unlike the textual reasoning tasks in Hsu et al. To the best of our knowledge, there are no similar visual benchmarks in the literature.

Implementation details (linear layer for scratchpad): The linear layer operates on each patch representation, not the CLS token. This ensures that the model retains the spatial resolution of the input image. As shown in Figure 7 and Appendix F, the quality of the generated scratchpads is sufficient to solve the proposed tasks. Also, while our framework uses a simple linear layer for its output, it is not constrained to this choice. More sophisticated generation mechanisms can be integrated into our method if necessary. However, we opted for simplicity to avoid introducing confounding factors.

Error propagation: Error propagation can be a concern. To mitigate this, we adopt alternated training (described in Section 3.2 and Appendix C.1), where the model is exposed to both perfect training frames and generated frames. This strategy is effective in reducing the gap between training and inference performance, as experimentally validated in the ablation study in Section 4.3.

Baselines: Our focus is on analyzing and understanding global vision-only reasoning tasks from scratch, not benchmarking against state-of-the-art Vision-Language Models (VLMs). However, we agree that it is interesting to check what the current state-of-the-art is capable of. Preliminary tests with VLMs like LLaVA and InstructBLIP (via image input and text prompting) show that these models struggle with our high-globality tasks. Interestingly, the answer seems to be more correlated with the given prompt rather than with the image. Moreover, even more complex state-of-the-art systems like ChatGPT, Claude and Gemini are unable to tell solvable vs non-solvable mazes apart. The reviewer can test that on their own by taking a screenshot of the figures in the paper and pasting them in their chatbot of choice.

Compute cost: While visual scratchpads introduce some overhead at inference (due to multiple forward passes for multi-frame scratchpads), this is offset by the ability to use smaller models. For instance, the inductive scratchpad achieves comparable performance with fewer parameters (see Figure 5). Additionally, the compute cost is adaptive, scaling with task complexity, which is a desirable property, rather than a weakness.

Response to Questions:

Novelty discussion: The novelty of our paper is twofold:

• We propose a theoretical framework to analyse the globality / locality of tasks and we validate it empirically.
• We apply the scratchpad concept to visual tasks requiring global reasoning. We introduce both single-frame and multi-frame (inductive) scratchpads and demonstrate their effectiveness on novel, challenging datasets. We also analyse in domain and out-of-domain performance of each method.

Implementation details:

• As mentioned, the scratchpad prediction head uses a linear layer applied to patch embeddings. The spatial resolution of the output is a function of the patch resolution and sequence length. Therefore the fact that the prediction head is not very powerful is not an issue, as it only needs to map a patch representation to a few pixels (16x16). Moreover, empirically we found this not to be problematic, examples of generated scratchpads, illustrating their quality, are provided in Figure 7 and Appendix F.

Please let us know if any part is still unclear. We believe we have answered most of your concerns and we hope you consider increasing your score accordingly.