SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Vijay Veerabadran\* Srinivas Ravishankar\*
Yuan Tang $^{1\dagger}$  Ritik Raina $^{1\dagger}$  Virginia R. de Sa $^{1,2}$
$^{1}$  Department of Cognitive Science,  $^{2}$  Halicioglu Data Science Institute,
University of California, San Diego, La Jolla, CA 92093
vveeraba@ucsd.edu
Humans solving algorithmic (or) reasoning problems typically exhibit solution times that grow as a function of problem difficulty. Adaptive recurrent neural networks have been shown to exhibit this property for various language-processing tasks. However, little work has been performed to assess whether such adaptive computation can also enable vision models to extrapolate solutions beyond their training distribution's difficulty level, with prior work focusing on very simple tasks. In this study, we investigate a critical functional role of such adaptive processing using recurrent neural networks: to dynamically scale computational resources conditional on input requirements that allow for zero-shot generalization to novel difficulty levels not seen during training using two challenging visual reasoning tasks: PathFinder and Mazes. We combine convolutional recurrent neural networks (ConvRNNs) with a learnable halting mechanism based on (Graves, 2016). We explore various implementations of such adaptive ConvRNNs (AdRNNs) ranging from tying weights across layers to more sophisticated biologically inspired recurrent networks that possess lateral connections and gating. We show that 1) AdRNNs learn to dynamically halt processing early (or late) to solve easier (or harder) problems, 2) these RNNs zero-shot generalize to more difficult problem settings not shown during training by dynamically increasing the number of recurrent iterations at test time. Our study provides modeling evidence supporting the hypothesis that recurrent processing enables the functional advantage of adaptively allocating compute resources conditional on input requirements and hence allowing generalization to harder difficulty levels of a visual reasoning problem without training.
Recurrent Neural Networks (RNNs) have emerged as a powerful tool for solving machine reasoning tasks, demonstrating impressive performance across a variety of tasks that require processing of sequential inputs. While recurrent processing is valuable for time-varying inputs, it can also be useful for static inputs as recurrent networks have the ability to scale computation to varying difficulty levels by varying the number of recurrent iterations. This can be helpful when different problems from the same task family exhibit significant variations in complexity (consider mazes, for example). However, conventional RNNs typically struggle to automatically generalize across different difficulty levels, requiring retraining, fine-tuning, or human intervention to pick the number of recurrent iterations needed. This limitation hampers their practical utility in real-world scenarios, since the span of difficulty levels seen during training is finite.
*Equal contribution
†Work done while at UC San Diego
Human reasoning, in contrast, is characterized by a highly adaptable use of computation - arbitrarily more computation can be used for more difficult tasks. Despite the evidence for such scaling in human computing there is limited work in the literature analyzing Adaptive RNNs (AdRNNs) that, on a sample-by-sample basis, automatically decide when to stop. Graves introduced Adaptive Computation Time (ACT), a mechanism to automatically stop computation, by learning to generate a scalar halting probability (Graves, 2016). However, similar to subsequent work (Banino et al., 2021), it was evaluated on either simple tasks such as parity checking, or language tasks. Few studies examine the training and evaluation of these AdRNNs on visual reasoning tasks, which pose unique challenges due to the redundant and high-dimensional nature of visual data. Furthermore, extrapolating to harder instances within-task has not been sufficiently studied, with Banino et al. (2021) evaluating such zero-shot performance only on the simple parity task.
In this paper, we study the problem of computation scaling in recurrent vision RNNs, with an emphasis on zero-shot extrapolation to harder/larger problems within the same task. The ability to handle unseen difficulty levels without fine-tuning, or human intervention, enables more robust and adaptable computer vision systems, which are crucial for real-world applications. We investigate the effectiveness of AdRNNs on two challenging, publically available, visual reasoning tasks based on curve tracing and route segmentation, namely PathFinder (introduced by Linsley et al. (2018)) and Mazes (introduced by Schwarzschild et al. (2021)).
The contributions of this work are as follows:
- We combine Convolutional RNNs with an adaptive computation method of Graves (2016) producing Adaptive ConvRNNs (AdRNNs) that are capable of learning a downstream task simultaneously while also learning to scale their computational steps as per input image/task requirements.
- We introduce LocRNN, a high performing recurrent architecture inspired by prior computational models of recurrence in biological vision.
- We show that AdRNNs learn to dynamically halt processing early (or late) to solve easier (or harder) problems when the train- and test-difficulty levels are matched on complex visual reasoning problems inspired by stimuli used in prior cognitive science research.
- During test time we introduce a previously unseen harder difficulty level. We evaluate AdRNNs on these new difficulty levels and show that they zero-shot generalize to more difficult problem settings not shown during training by dynamically increasing the number of recurrent iterations at test time well beyond the number of recurrent steps used during training.
Our work is relevant to the visual routines literature introduced by Ullman (1984) and further reviewed elaborately by Roelfsema et al. (2000). The core idea of visual routines that make it relevant to our studied question of task extrapolation is the flexible sequencing of elemental operations resulting in a dynamic computational graph, which makes RNNs a natural approach to solve such tasks.
Prior research has proposed mechanisms for RNNs that learn the amount of recurrent computational steps required. Following are two particularly relevant attempts in prior literature in this area: Graves (2016) developed the Adaptive Computation Time (ACT) method to train RNNs on Natural Language Processing tasks with a sigmoidal halting unit that determines when to halt recurrent processing, and Banino et al. (2021) extend upon this pioneering work by taking a probabilistic approach to halting and introducing a geometric prior-based specification of computational budget. However, neither of these studies applied adaptive computation to vision RNNs (that are significantly trickier to optimize). Also the prior studies did not attempt to examine computation scaling under zero-shot generalization to harder difficulty levels. Additionally, Eyzaguirre and Soto (2020) developed a version of ACT applied to visual reasoning problems, however, their observations on the CLEVR dataset (in which the issues related to object diversity are highlighted by Kim et al. (2018)) do not explore generalization to new difficulty levels which is central to our results.
In terms of generalization from easy to hard vision problems, our work is relevant to (Schwarzschild et al., 2021) and (Bansal et al., 2022) where the authors evaluate the application of recurrent neural networks to generalize from easier to harder problems. Our work extends and differs from their
intriguing study of recurrence in task extrapolation in the following ways: 1) For recurrent networks used in their study, human intervention is required to specify the number of recurrent computational steps during the testing phase. 2) Their work explores sequential task extrapolation in general with abstract problems such as Prefix Sum and solving Chess Puzzles while our work extends it to particularly focus on extrapolation in visual task learning. 3) We present evaluation on the Pathfinder challenge (Linsley et al., 2018), a relatively more large-scale visual reasoning task shown to be very challenging (Tay et al., 2020), the design of which dates back to (Jolicoeur et al., 1986). 4) Schwarzschild et al. (2021) and Bansal et al. (2022) implement only the most straightforward form of recurrence realized by weight-tying, i.e., they only evaluate ResNets with weight sharing across residual blocks (albeit with a novel training scheme in the follow up study (Bansal et al., 2022)). In addition to such recurrent ResNets, we present analyses with highly sophisticated recurrent architectures specialized for recurrent image processing. 5) We introduce LocRNN, a high performing recurrent architecture based on prior computational models of cortical recurrence.
On the design of recurrent architectures, our work is loosely related to (Eigen et al., 2013), (Pinheiro and Collobert, 2014), (Veerabadran and de Sa, 2020) and (Liao and Poggio, 2016) which discuss the role of weight sharing in feedforward networks to produce recurrent processing. We are interested, however, in designing new specialized recurrent architectures that play a role both in human and machine vision. While there exist more such recurrent architectures informed by neuroscience such as Linsley et al. (2018); Nayebi et al. (2022), we find these architectures to be quite difficult to interpret and unstable to train (based on in-difficulty performance evaluation included in the Supplementary). Our proposed LocRNN architecture in comparison is an elegant and easy-to-interpret implementation of recurrence that is stable to train. Our work is also relevant to prior work on modeling speed-accuracy tradeoffs observed during visual perception (Spoerer et al., 2020).
3 Datasets
For evaluating the ability of various models in exhibiting task extrapolation, we curate two challenging visual reasoning tasks, Mazes and PathFinder, with instances at multiple parametric difficulty levels. Both tasks involve the visual routines of marking and curve tracing (Ullman, 1984). These datasets are inspired by prior visual psychophysics research where such tasks were used abundantly to estimate the cognitive and neural underpinnings of sequential visual processing, such as incremental grouping, structure and preference of lateral connections, etc. (Jolicoeur and Ingleton, 1991; Ullman, 1984; Li et al., 2006; Roelfsema, 2006). In the following subsections we describe the specifics of our tasks.
3.1 PathFinder challenge - curve tracing
Figure 1: Representative examples from PathFinder and Mazes datasets. Left: (a) Positive PathFinder-9 (b) Negative PathFinder-18 (c) Positive PathFinder-24; Right: (a)  $9 \times 9$  mazes, (b)  $15 \times 15$  mazes, (a)  $19 \times 19$  mazes, (a)  $25 \times 25$  mazes
Task description: In the PathFinder task introduced by Linsley et al. (2018), models are trained to identify whether two circular disks in an input stimulus form the two ends of a locally connected path made up of small "segments". Each image consists of two main long connected paths  $P_0$  and  $P_1$  made up of locally aligned segments as well as shorter distractor paths. Each image also contains two circular disks which are placed at two of the 4 possible endpoints of  $P_0$  and  $P_1$ . Images that contain a disk on both ends of the same path are classified as positive, and those containing a disk on endpoints of different paths are classified as negative. Examples are shown in Figure 1.
Difficulty levels: Pathfinder is designed at different difficulty levels parameterized by the length (number of segments) of the paths  $P_0$  and  $P_1$  mentioned above. The easiest version uses paths that are 9 segments long (PathFinder-9), while the medium and hard versions contain paths that are
14 (PathFinder-14) and 18 (PathFinder-18) segments long respectively (see example images in the Appendix). This dataset consists of a total of 800,000 RGB images at each difficulty level, each with a spatial resolution of  $160 \times 160$  pixels. There are an equal number of positive and negative instances at each difficulty level. We use 700,000 images for training and 100,000 images as the test set from each difficulty level. We combined these datasets to create a more challenging dataset with varying levels of difficulty, which we refer to as PathFinder-Mixed. To evaluate the zero-shot difficulty extrapolation on PathFinder in Sec. 5.2 we generated 100,000 images each with contour lengths 21 and 24 respectively; we call these PathFinder-21 and PathFinder-24.
Evaluation criteria: Since this is a classification challenge, we use accuracy, i.e. % correct on test-images as the evaluation metric to rank model performance on PathFinder. Model architectures receive an input image and process it via a stack of standard convolution/recurrent-convolution layers followed by a classification two-class readout. Since this is a binary classification challenge with balanced classes, chance performance is  $50\%$  for random predictions.
3.2 Mazes challenge - route segmentation
Task description: Human beings are adept at solving mazes, a task that requires application of a similar serial grouping operation like PathFinder in order to discover connectivity from a starting point to the final destination of the maze amidst blocking walls. For evaluating model performance on solving mazes of varying difficulty, we use the publicly available version of the Mazes challenge developed by Schwarzschild et al. (2021). They implemented this Mazes challenge as a binary segmentation problem where models take  $N \times N$  images of square-shaped mazes as input with three channels (RGB) with the start position, end position, permissible regions and impermissible regions marked in the image. The output produced by models is a binary segmentation of the route discovered by the model from the start to the end position.
**Difficulty levels:** The Mazes challenge has been designed at several difficulty levels, each difficulty level is parameterized by the size of the square grid that the maze is designed into. We use maze datasets of grid sizes  $9 \times 9$  and  $15 \times 15$  for training. Each dataset consists of 50,000 training images and 10,000 test images that are guaranteed to not overlap. The spatial resolution of  $9 \times 9$  maze images is  $24 \times 24$  pixels and that of  $15 \times 15$  maze images is  $36 \times 36$  pixels. Similar to PathFinder-Mixed, we combine  $9 \times 9$  and  $15 \times 15$  grid sizes to form Mazes-Mixed. We also constructed larger Mazes with grid sizes  $19 \times 19$  (with spatial resolution  $44 \times 44$  pixels) and  $25 \times 25$  (with spatial resolution  $56 \times 56$  pixels) to evaluate a model's ability to extrapolate to difficulties not seen during training. As above with PathFinder's extrapolation evaluation, these novel difficulty Mazes were only used during testing and were not in any form used in the training or hyperparameter optimization process.
Evaluation criteria: Mazes is a segmentation challenge and hence, one could potentially consider partially correct routes during evaluation (for example with average of per-pixel accuracy). However this is less strict than giving each image a single binary score reflecting if all pixels are labeled correctly. Evaluation criteria for mazes is hence the total % of test-set mazes completely accurately solved at a given difficulty level.
4 Model architectures and training
4.1 Implementations of adaptive computations evaluated on task extrapolation
In this section, we describe our choice of recurrent architecture designs that we use to study the behavior of adaptive computation and task extrapolation in deep learning. We process images (denoted as  $\mathbf{X} \in \mathbb{R}^{c,h,w}$ ) in three stages that are common to all models we evaluate in this study. These three stages are as follows: (1) an input convolution layer (denoted as  $\mathrm{input}(.)$ ) that operates on the image directly.
$$
\mathbf {h} _ {\mathbf {0}} = \operatorname {R e L U} (\operatorname {i n p u t} (\mathbf {X})) \tag {1}
$$
The output from this stage  $\mathbf{h_0}\in \mathbb{R}^{d,h,w}$  is fed as input to the following recurrent block in (2).
$$
\mathbf {h} _ {\mathbf {t}} = \mathbf {r} \left(\mathbf {h} _ {\mathbf {t} - \mathbf {1}}, \mathbf {h} _ {\mathbf {0}}\right), t \in [ 1, t _ {\text {t r a i n}} ] \tag {2}
$$
where  $\mathbf{r}(.)$  is the recurrent block.  $t_{train}$  is a training hyperparameter indicating the maximum number of timesteps for unrolling  $\mathbf{r}(.)$  during training. This block consists of a sequence of convolution
layers applied in an iterative manner for any arbitrary number of timesteps. This is the sub-part of our model that is capable of performing adaptive computations. (The specific architecture of these convolution blocks constitute different implementations of recurrent operations described further below.) (3) a readout layer containing a block of convolution and pooling operations that produce the desired output from our network.
$$
\hat {\mathbf {y}} _ {\mathbf {t}} = \operatorname {r e a d o u t} \left(\mathbf {h} _ {\mathbf {t}}\right), t \in [ 1, t _ {\text {t r a i n}} ] \tag {3}
$$
While we keep the input and readout layer architectures the same for all models we evaluate, their intermediate recurrent blocks have different architectures (corresponding to their respective recurrent cells). We explore three different implementations of recurrent computations which are used in the second, recurrent block of models. Our first choice as the intermediate feature processing block consists of a residual network with weight tying across all layers; we refer to this model as R-ResNet30. Next, we study the performance of the following specialized convolutional recurrent units from prior work: horizontal convolutional GRU (hConvGRU) and its stable variant (Linsley et al., 2018, 2020) and a convolutional Gated Recurrent Unit (ConvGRU) (Ballas et al., 2015) with LayerNorm (ConvGRU does not converge in the absence of LayerNorm). Third, we design a novel recurrent cell based on prior computational models of cortical recurrent processing (Li et al., 2006); this model equips the biologically inspired design choices of long-range lateral interactions, gating, and a separate population of interneurons. This model is referred to as LocRNN and is described in the following Sec. 4.3. Importantly all models are matched for trainable parameters.
4.2 Combining ConvRNNs with Adaptive Computation Time (ACT)
The central theme of our work is to show that RNNs can flexibly adapt (or scale) their computation according to input requirements. We achieve this ability by combining ConvRNNs with an adaptive computation mechanism based on Graves (2016) called Adaptive Computation Time (ACT). A difference between our work and ACT is that our visual reasoning task involves static inputs (i.e. sequence of length 1) whereas Graves (2016) deals with variable-length sequences and learns adaptive processing of each token. Owing to this difference, our halting mechanism is similar to ACT applied to a 1-token input sequence.
The key idea of ACT is to introduce a separate "halting mechanism" that learns to control the number of recurrent computation steps dynamically, conditioned on each input example's processing. In addition to producing the next recurrent state, an RNN equipped with ACT also produces a scalar value called the "halting score" for each computation step. In addition to the next hidden state computed using Equation 2, we generate a scalar halting score  $p_t$  at each step using a learnt convolution layer (halt_conv(.)) that is shared across timesteps:
$$
p _ {t} = \sigma (\max  _ {-} \operatorname {p o o l} (\operatorname {h a l t} _ {-} \operatorname {c o n v} (\mathbf {h} _ {\mathbf {t} - \mathbf {1}}))) \tag {4}
$$
The RNN treats the cumulative sum of the halting scores up to timestep  $t$  ( $P_t$  described below) as a quantity used to determine whether processing is terminated at that timestep.
$$
P _ {t} = \sum_ {t ^ {\prime} = 1} ^ {t} p _ {t ^ {\prime}}
$$
The RNN keeps track of the accumulated halting scores and checks if the cumulative sum of the halting scores at each step reaches a predefined threshold  $(1 - \epsilon)$ . If the threshold is reached, the computation stops (i.e.,  $p_t = 0 \forall t > t_{halt}$ ) where the cumulative halting score threshold is reached at timestep  $t_{halt} = \min \{t : P_t >= (1 - \epsilon)\}$ . The adaptive hidden state is then computed as a weighted average of the hidden states up to  $t_{halt}$ , scaled by the halting scores at each time step. The readout is then applied to this final adaptive hidden state to produce the ACT task prediction,  $\hat{\mathbf{y}}_{act}$ :
$$
\mathbf {h} _ {\text {a c t}} = \sum_ {t = 1} ^ {t _ {\text {h a l t}}} p _ {t} \cdot \mathbf {h} _ {\mathbf {t}} \tag {5}
$$
$$
\hat {\mathbf {y}} _ {a c t} = \operatorname {r e a d o u t} \left(\mathbf {h} _ {a c t}\right) \tag {6}
$$
The abovementioned ConvRNN combined with ACT is trained to optimize both the downstream task loss  $(||\mathbf{y} - \hat{\mathbf{y}}_{act}||_p)$  and an auxiliary 'ponder cost' corresponding to the halting mechanism
that encourages models to use the fewest number of recurrent step to process an input example. Ponder cost is computed as the cumulative halting score until timestep  $t_{halt} - 1$ , maximizing this term encourages (or minimizing its negative as shown below) completing the task as quickly (with as few recurrent steps) as possible. The following final objective function is optimized to train the full model containing ConvRNNs with ACT where  $\tau$  is a hyperparameter.
$$
\mathcal {L} = \sum_ {i = 0} ^ {i = | | \mathcal {D} | |} \frac {1}{| | \mathcal {D} | |} \| \mathbf {y} ^ {i} - \hat {\mathbf {y}} _ {a c t} ^ {i} \| _ {p} - \tau \sum_ {t = 1} ^ {t _ {h a l t} ^ {i} - 1} p _ {t} ^ {i} \tag {7}
$$
4.3 Formulation of LocRNN
We note that prior work has explored the development of recurrent architectures tailored for vision. Here, we introduce a similar but highly expressive recurrent architecture designed based on a computational model of iterative contour processing in primate vision from Li (1998). This model is an ODE-based computational model of interactions between cortical columns in primate area V1 mediated by "lateral connections" local to a cortical area.
Figure 2: Contrasting the architectures of LocRNN (left) and ConvGRU (bottom-right).  $\mathbf{h}_0$  is shown as a magenta arrow and does not change across timesteps. As illustrated also in this figure, LocRNN's interneuron activations  $(\mathbf{S}_{\mathbf{t}})$  are not passed to the next layer. ConvGRU on the other hand uses a single uniform neural population. (Top-right) shows a visualization of LocRNN's  $\mathbf{L}_{\mathbf{t}}$  activations as they perform contour tracing to solve an input image from PathFinder-14 (PF-14) (visualization of all timesteps available in Supplementary Fig. 8).
By discretizing the continuous-form ODE dynamics of processing units from Li (1998), we arrived at an interpretable and powerful set of dynamics for "LocRNN". As in ConvGRU, the effective receptive field of LocRNN output neurons increases linearly with recurrent timesteps.
Unlike ConvGRU, the LocRNN's hidden state is composed of two neural populations, L and S; the activity in these populations are referred to as  $\mathbf{L}_{\mathbf{t}}$  and  $\mathbf{S}_{\mathbf{t}}$  at timestep  $t$  respectively. These two populations are motivated by the  $x$  (excitatory) and  $y$  (inhibitory) neurons in Li (1998). However we found that restricting the signs of their weights (to reflect exclusively excitatory and inhibitory connections) led to less stable behavior. On the other hand, retaining the interneuron property (local computation not projecting out for downstream processing) of the S population performed better than using a single uniform population. The following equations illustrate the working of ACT and how the readout is applied in LocRNN.