SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
p _ {t} = \sigma (\max  _ {-} \operatorname {p o o l} (\operatorname {h a l t} _ {-} \operatorname {c o n v} (\mathbf {L} _ {\mathbf {t} - \mathbf {1}}))) \tag {8}
$$
$$
\hat {\mathbf {y}} _ {a c t} = \text {r e a d o u t} \left(\sum_ {t = 1} ^ {t _ {h a l t}} p _ {t} \cdot \mathbf {L} _ {\mathbf {t}}\right) \tag {9}
$$
Initially, both  $\mathbf{L}_0$  and  $\mathbf{S}_0$  populations are set to a tensor of zeros with the same shape as  $\mathbf{h}_0$ , a  $4d$  tensor of shape (batch_size  $\times$  channels  $\times$  height  $\times$  width).
$L$  and  $S$  update gates  $\mathbf{G}_{\mathbf{t}}^{\mathbf{L}}$  and  $\mathbf{G}_{\mathbf{t}}^{\mathbf{S}}$  (same shape as  $\mathbf{L}$  and  $\mathbf{S}$ ) are computed as functions of the input and current hidden states  $\mathbf{L}_{\mathbf{t}-1}$  and  $\mathbf{S}_{\mathbf{t}-1}$  using 1x1 convolutions  $\mathbf{U}_{*}$ .
$$
\mathbf {G} _ {\mathbf {t}} ^ {\mathbf {L}} = \sigma \left(L N \left(\mathbf {U} _ {\mathbf {L}} * \mathbf {h} _ {\mathbf {0}}\right) + L N \left(\mathbf {U} _ {\mathbf {L} \rightarrow \mathbf {L}} * \mathbf {L} _ {\mathbf {t} - \mathbf {1}}\right)\right) \tag {10}
$$
$$
\mathbf {G} _ {\mathbf {t}} ^ {\mathbf {S}} = \sigma \left(L N \left(\mathbf {U} _ {\mathbf {S}} * \mathbf {h} _ {\mathbf {0}}\right) + L N \left(\mathbf {U} _ {\mathbf {S} \rightarrow \mathbf {S}} * \mathbf {S} _ {\mathbf {t} - \mathbf {1}}\right)\right) \tag {11}
$$
Each of the 4 types of lateral connections are modeled by convolution kernels  $\mathbf{W}_{\mathbf{L}\rightarrow \mathbf{L}},\mathbf{W}_{\mathbf{L}\rightarrow \mathbf{S}},\mathbf{W}_{\mathbf{S}\rightarrow \mathbf{S}}$  , and  $\mathbf{W}_{\mathbf{S}\rightarrow \mathbf{L}}$  respectively of shape  $d\times d\times k\times k$  where  $d$  is the dimensionality of the hidden state and  $k$  represents the kernel spatial size.
$$
\tilde {\mathbf {L}} _ {\mathbf {t}} = \gamma \left(\mathbf {W} _ {\mathbf {L}} * \mathbf {h} _ {\mathbf {0}} + \mathbf {W} _ {\mathbf {L} \rightarrow \mathbf {L}} * \mathbf {L} _ {\mathbf {t} - \mathbf {1}} + \mathbf {W} _ {\mathbf {S} \rightarrow \mathbf {L}} * \mathbf {S} _ {\mathbf {t} - \mathbf {1}}\right) \tag {12}
$$
$$
\tilde {\mathbf {S}} _ {\mathbf {t}} = \gamma \left(\mathbf {W} _ {\mathbf {S}} * \mathbf {h} _ {\mathbf {0}} + \mathbf {W} _ {\mathbf {L} \rightarrow \mathbf {S}} * \mathbf {L} _ {\mathbf {t} - \mathbf {1}} + \mathbf {W} _ {\mathbf {S} \rightarrow \mathbf {S}} * \mathbf {S} _ {\mathbf {t} - \mathbf {1}}\right) \tag {13}
$$
Once the long-range lateral influences are computed and stored in  $\tilde{\mathbf{L}}_{\mathrm{t}}$  and  $\tilde{S}_{\mathrm{t}}$ , these are mixed with the previous hidden states using the gates computed in Eq. 10 and 11. These hidden states are then passed on to subsequent recurrent iterations where even longer-range interactions occur (as time increases).
$$
\mathbf {L} _ {\mathbf {t}} = \kappa \left(L N \left(\mathbf {G} _ {\mathbf {t}} ^ {\mathbf {L}} \odot \tilde {\mathbf {L}} _ {\mathbf {t}} + \left(1 - \mathbf {G} _ {\mathbf {t}} ^ {\mathbf {L}}\right) \odot \mathbf {L} _ {\mathbf {t} - 1}\right)\right) \tag {14}
$$
$$
\mathbf {S} _ {\mathbf {t}} = \kappa \left(L N \left(\mathbf {G} _ {\mathbf {t}} ^ {\mathbf {S}} \odot \tilde {\mathbf {S}} _ {\mathbf {t}} + \left(1 - \mathbf {G} _ {\mathbf {t}} ^ {\mathbf {S}}\right) \odot \mathbf {S} _ {\mathbf {t} - 1}\right)\right) \tag {15}
$$
In the above equations,  $LN()$  stands for Layer Normalization (Ba et al., 2016), and the nonlinearities  $\gamma$  and  $\kappa$  are both set to ReLU. One of the differences between LocRNN and ConvGRU (which we also evaluate) is the presence of the interneuron S population in the former.
5.1 Adaptive RNNs scale their computation as a function of input difficulty
In this section, we report our observations from training AdRNNs on a mixture of difficulty levels from PathFinder and Mazes & evaluating them on a held-out set of images from the same mixture of difficulty levels of PathFinder and Mazes. We created one training dataset for each of these challenges by combining input samples from multiple "easy" difficulty levels. Combining input samples from multiple difficulty levels increases the diversity of computational requirements during training and helps in developing models that do not degenerate to using the same computational steps for all input samples.
As described in Section 3, we created a PathFinder training set of images and corresponding labels by sampling an equal number of images from three difficulty levels: PathFinder-9, PathFinder-14, and PathFinder-18. Similarly, we created a Mazes training set of images and ground-truth segmentation labels by sampling an equal number of images from two difficulty levels: 9x9 mazes and 15x15 mazes.
On held-out test sets that matched the difficulty level of the above-described training data, we tested whether models were able to scale their recurrent computational steps as a function of the input difficulty level. We show the results from this evaluation in Table. 1. We observed that both the sophisticated AdRNNs tested (ConvGRU and LocRNN) were able to generalize to the held-out set. We also observed that variants of horizontal GRU (Linsley et al., 2018) and Linsley et al. (2020) models generalized to the (within-difficulty) held-out set without using ACT. The simpler recurrent network without gating, weight-tied R-ResNet-30, was unable to learn on PathFinder and on Mazes highlighting the importance of specialized operations such as gating and backpropagation through time. The presence of skip connections between R-ResNet-30's recurrent blocks could still not match the expressivity of the specialized RNNs.
If ACT is working as expected, we must observe that AdRNNs that learn the task dynamically use less compute to solve easy examples and more compute (more recurrent iterations) to solve harder examples. To check for this trend, we analyzed the number of steps chosen by the model before halting for each example from the validation sets of PathFinder and Mazes; with varying contour lengths and maze sizes respectively. These results are shown by the cool colors in Figure 3 for PathFinder and in the Supplementary for Mazes. As is clearly observable from the trend in these Figures, examples that we consider as harder (longer contours in PathFinder and larger Mazes) are assigned a higher number of recurrent computation steps by ACT than easy examples. Hence we show that AdRNNs obtained by combining ConvRNNs with ACT training are capable of learning both PathFinder and Mazes in addition to learning to adapt their recurrent computational steps as a function of input example difficulty. To the best of our knowledge, we are the first to show the above result for visual tasks inspired by stimuli used in prior cognitive science research.
PathFinder-Mixed (%)	Mazes-Mixed (%)
ResNet-30	50.41	0.0
R-ResNet-30 (ACT)	49.37	0.0
Linsley et al. (2020)	50.0	78.33
hConvGRU \(t_{inference}=t_{train}\)	89.66	99.69
hConvGRU (stable halting)	89.65	99.54
ConvGRU (ACT)	95.26	98.4
LocRNN (ACT) (ours)	97.13	98.4
Table 1: Accuracies  $(\%)$  ↑ of models on the two visual reasoning tasks. Chance performance is  $50\%$  for pathfinder and  $0\%$  for Mazes.
Figure 3: Distributions of halting steps across samples in each validation set of PathFinder. Computation appears to scale to match difficulties of the datasets for both LocRNN (left) and ConvGRU (right) models. 9-length contours typically halt after 4-6 steps, while 24-length contours can take up to 9 steps. The red bars show the distribution in the extrapolation datasets.
Figure 4: Vanilla RNNs (solid lines) and AdRNNs (dashed) trained on 12 iterations (a) within-difficulty (b) extrapolating to PathFinder-21 (c) extrapolating to PathFinder-24
5.2 Adaptive RNNs generalize to novel difficulty levels by scaling their computation
In typical circumstances where RNNs are expected to solve instances that are more difficult, human intervention in the form of specification of the number of recurrent iterations is commonplace (Schwarzschild et al., 2022). This approach is both laborious and expensive to perform as it requires training and/or evaluation on new test sets with various settings of the number of recurrent timesteps. Additionally, this process needs to be repeated at every occasion of new incoming test data that is potentially of increased difficulty. In the previous section, we showed the ability of AdRNNs in learning to solve PathFinder and Mazes while also demonstrating that they assign computation as a function of input difficulty on a test-set with matched task difficulty level as the training set. Here we are interested in testing whether AdRNNs are able to extrapolate (by using more recurrent computational steps compared to training) in complex reasoning tasks. We constructed two additional datasets each in the PathFinder and Mazes families to test these adaptive models' ability to extrapolate, as mentioned in Section 3. While AdRNNs are trained on the training difficulty levels for a maximum of  $t_{train}$  iterations, we evaluated them on these new (more difficult) datasets at inference with a maximum of  $t_{inference} \geq t_{train}$  recurrent iterations where we refer to AdRNNs as operating in their extrapolation phase. We report the results from this evaluation in Table. 2 and in Fig. 4.
PathFinder-21 (%)	PathFinder-24 (%)	Mazes-19 (%)	Mazes-25 (%)
ResNet-30	50.0	50.0	0.	0.
R-ResNet-30 (ACT)	50.0	50.0	0.	0.
Linsley et al. (2020)	50.0	50.0	2.93	0.01
hConvGRU \(t_{\text{inference}} = t_{\text{train}}\) *	64.21	58.35	16.2 ± 2.66	5.29 ± 0.26
hConvGRU (stable halting)	50.0	50.0	50.26 ± 6.04	21.36 ± 2.82
ConvGRU (ACT)	82.63 ± 4.84	74.14 ± 6.52	75.1 ± 11.96	46.93 ± 4.2
LocRNN (ACT) (ours)	92.89 ± 0.9	85.81 ± 5.57	86.83 ± 2.94	49.99 ± 4.48
Table 2: Accuracies (Mean % ± SEM) ↑ on extrapolation datasets. Chance performance is 50% for pathfinder and 0% for Mazes.
$^*$ hConvGRU training converged only on 1 out of 3 seeds (which were chosen randomly for all models here) on PathFinder, corresponding results above show this converged model's performance.
Adaptive RNNs generalize to novel difficulty levels by scaling their computation. Our evaluation of AdRNNs trained with ACT on novel harder difficulty levels shows that as expected, AdRNNs have learned to use the optimal number of recurrent computational steps required for achieving strong generalization to the novel difficulty levels on a per-image level. On PathFinder, as seen in Figure 3, most instances from PathFinder-21 and PathFinder-24 take 9 steps in LocRNN and up to 7 steps in ConvGRU, higher than the number of steps used for easy training difficulty examples. In Fig 5 we visualize the relationship between mazes' difficulty and the number of recurrent iterations used by ACT (on a per-instance level). We observe that maze difficulty (length of the ground-truth route in pixels) and the number of ACT iterations are strongly positively correlated. Notably, AdRNNs trained with ACT choose to make  $t_{halt}$  during inference on longer mazes greater than  $t_{halt}$  used on the shorter training mazes.
Figure 5: Relationship between halting step and difficulty level of Mazes for the extrapolation evaluation. Here we visualize the relationship between the halting step and solution length of mazes at a per-instance level. We note a strong positive correlation between the difficulty level (length of maze solution segmentation) and halting step used by LocRNN (left) and ConvGRU (right) AdRNNs.
5.3 AdRNNs outperform halting in hConvGRU based on stability of hidden-state dynamics
Using stability in hidden-state space as an alternative way to perform halting during extrapolation, we evaluate hConvGRU on extrapolation. For this method, we halted processing in hConvGRU when the mean absolute difference between two subsequent states  $(||\mathbf{h}_{\mathbf{t}} - \mathbf{h}_{\mathbf{t} - 1}||)$  reduced to less than a tenth of the difference between the first two states (heuristic for stability). This method referred as hConvGRU (stable halting) in Tables 1 and 2 shows differences in generalization across our datasets. AdRNNs perform better than hConvGRU (stable halting) on both datasets; especially, stable halting completely fails to show generalization to PathFinder-21 and -24. We hypothesize this method to work suboptimally for the following reasons:
- Stability-based halting requires defining a hand-engineered heuristic on how much change in the output is considered small enough to halt. In typical segmentation tasks like Mazes, the network output for the initial few timesteps is highly stable in making nonsensical
predictions (predicting all pixels as the negative class) and thus, heuristics need to identify an inflection point in the output trajectory where meaningful predictions start to emerge and stabilize. In the absence of ground truth information, one cannot pick a heuristic that generalizes to unseen data.
- Learnable halting makes less assumptions about the hidden state's properties, and hence doesn't enforce hard constraints such as stability to be satisfied by training. Some, but not all RNNs, have stable hidden states wherein the network response stops changing after reaching an attractor. Linsley et al. (2020) argue that RNNs that are expressive have an intrinsic inability to learn stable hidden states. For RNNs to be stable, their hidden state transformation needs to model a contractive mapping (Miller and Hardt, 2018; Pascanu et al., 2013). That is, the recurrent transition function  $F$  satisfies the following inequality:  $||F(h_{t}) - F(h_{t - 1})||_{p} < \lambda ||h_{t} - h_{t - 1}||_{p}$ . RNNs with stable hidden states that satisfy the above inequality are quite difficult to train on challenging problems in practice. Even when stable models perform comparably to unstable models as in (Miller and Hardt, 2018), the authors show unstable models' advantages such as performance improvements in the short-time horizon and lesser vanishing gradient issues.
The advantage of using recurrent networks for processing static inputs adaptively, particularly in order to zero-shot generalize to new difficulty levels is understudied. In this work, we show that deep convolutional networks with intermediate recurrent blocks (ConvRNNs operating on image features) can be combined with an adaptive computation technique to learn to dynamically process different input examples based on a per-instance difficulty level. We combine ConvRNNs with a learnable halting mechanism that is based on Graves (2016) to produce AdRNNs. We evaluated diverse implementations of recurrence in increasing level of sophistication/complexity, R-ResNet-30 with weight-tying, ConvGRU and hConvGRU with gating, and LocRNN with two separate populations of horizontally connected units (one of which are interneurons) and gating on two challenging visual reasoning problems, PathFinder and Mazes, which are generated at various levels of difficulty. First, we showed that only the specialized RNNs (hConvGRU) and AdRNNs (LocRNN and ConvGRU) are capable of learning PathFinder and Mazes. Second, these AdRNNs trained with ACT are learning to dynamically use less (or more) recurrent computational steps for easy (or hard) PathFinder and Maze problems respectively as discussed in Section 5.1. More interestingly when the difficulty level of the test set was increased relative to training difficulty levels, AdRNNs generalized to these harder instances in a zero-shot manner by allocating more recurrent computation than was ever used during training. They also outperform stability heuristics-based adaptive recurrent networks. Our work empirically shows this hypothesized advantage of using adaptive recurrent processing on static-image tasks for the first time to the best of our knowledge.
In addition to adaptively scaling computation according to the need, the AdRNNs can also solve the problems more efficiently, choosing to stop at earlier times than the non-adaptively trained RNNs (compare halting times in Fig 3 with performance curves in Fig 4) with human arbitrarily chosen training iterations. Thus, just as computation can be scaled to unseen problem difficulties without human intervention, the training iterations required can be (better) discovered automatically.
7 Limitations and Future Work
The current work has been only applied to static images, but recurrent networks are able to process time-varying input; future work will explore our models on video input. In addition, there appears to be a benefit of the LocRNN model over ConvGRU that deserves further study to determine the mechanism behind the benefit of the interneuron population.
For future work, we would like to further explore the benefits of training only on very small problem instances with the ACT mechanism automatically allowing for harder instances. This could greatly reduce latency and energy consumption both during training and inference, an increasingly relevant concern as the size, memory footprint, training time and inference latency of models grow drastically. Finally, we wish to explore how this approach might link to the human developmental literature, similar to earlier work by Elman (1993).
8 Acknowledgments & Funding disclosure
We thank Garrison W. Cottrell, Michael C. Mozer, Michael L. Iuzzolino, Pradeep Shenoy and Saurabh Singh for their feedback and insights on this work. This work is funded by: (1) NSF CRCNS Award No: 2208362, (2) Sony Faculty Innovation Award, and (3) Kavli Symposium Inspired Proposal and support from Social Sciences Computing Facility and Department of Cognitive Science at UC San Diego.
J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
N. Ballas, L. Yao, C. Pal, and A. Courville. Delving deeper into convolutional networks for learning video representations. arXiv preprint arXiv:1511.06432, 2015.
A. Banino, J. Balaguer, and C. Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021.
A. Bansal, A. Schwarzschild, E. Borgnia, Z. Emam, F. Huang, M. Goldblum, and T. Goldstein. End-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking. arXiv preprint arXiv:2202.05826, 2022.
D. Eigen, J. Rolfe, R. Fergus, and Y. LeCun. Understanding deep architectures using a recursive convolutional network. arXiv preprint arXiv:1312.1847, 2013.
J. L. Elman. Learning and development in neural networks: The importance of starting small. Cognition, 48(1):71-99, 1993.
C. Eyzaguirre and A. Soto. Differentiable adaptive computation time for visual reasoning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12817-12825, 2020.
A. Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.08983, 2016.
M. Iuzzolino, M. C. Mozer, and S. Bengio. Improving anytime prediction with parallel cascaded networks and a temporal-difference loss. Advances in Neural Information Processing Systems, 34: 27631-27644, 2021.
P. Jolicoeur and M. Ingleton. Size invariance in curve tracing. Memory & Cognition, 19(1):21-36, 1991.
P. Jolicoeur, S. Ullman, and M. Mackay. Curve tracing: A possible basic operation in the perception of spatial relations. Memory & Cognition, 14(2):129-140, 1986.
J. Kim, M. Ricci, and T. Serre. Not-so-clevr: learning same-different relations strains feedforward neural networks. Interface focus, 8(4):20180011, 2018.
W. Li, V. Piech, and C. D. Gilbert. Contour saliency in primary visual cortex. Neuron, 50(6):951-962, 2006.
Z. Li. A neural model of contour integration in the primary visual cortex. Neural computation, 10(4): 903-940, 1998.
Q. Liao and T. Poggio. Bridging the gaps between residual learning, recurrent neural networks and visual cortex. arXiv preprint arXiv:1604.03640, 2016.
D. Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. In Advances in Neural Information Processing Systems, pages 152-164, 2018.
D. Linsley, A. Karkada Ashok, L. N. Govindarajan, R. Liu, and T. Serre. Stable and expressive recurrent vision models. Advances in Neural Information Processing Systems, 33:10456-10467, 2020.
J. Miller and M. Hardt. Stable recurrent models. arXiv preprint arXiv:1805.10369, 2018.
A. Nayebi, J. Sagastuy-Brena, D. M. Bear, K. Kar, J. Kubilius, S. Ganguli, D. Sussillo, J. J. DiCarlo, and D. L. Yamins. Recurrent connections in the primate ventral visual stream mediate a tradeoff between task performance and network size during core object recognition. bioRxiv, pages 2021-02, 2022.
R. Pascanu, T. Mikolov, and Y. Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318. Pmlr, 2013.
A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer. Automatic differentiation in pytorch. 2017.
P. Pinheiro and R. Collobert. Recurrent convolutional neural networks for scene labeling. In International conference on machine learning, pages 82-90. PMLR, 2014.
P. R. Roelfsema. CORTICAL ALGORITHMS FOR PERCEPTUAL GROUPING. Annual Review of Neuroscience, 29(1):203-227, 7 2006. ISSN 0147-006X. doi: 10.1146/annurev.neuro.29.051605.112939. URL http://www.annualreviews.org/doi/10.1146/annurev.neuro.29.051605.112939.
P. R. Roelfsema, V. A. Lamme, and H. Spekreijse. The implementation of visual routines. Vision research, 40(10-12):1385-1411, 2000.
A. Schwarzschild, E. Borgnia, A. Gupta, F. Huang, U. Vishkin, M. Goldblum, and T. Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. Advances in Neural Information Processing Systems, 34:6695-6706, 2021.
A. Schwarzschild, A. Gupta, M. Goldblum, and T. Goldstein. The uncanny similarity of recurrence and depth. 10th International Conference on Learning Representations, ICLR 2022, 2022.
C. J. Spoerer, T. C. Kietzmann, J. Mehrer, I. Charest, and N. Kriegeskorte. Recurrent neural networks can explain flexible trading of speed and accuracy in biological vision. PLoS computational biology, 16(10):e1008215, 2020.
Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.04006, 2020.
S. Ullman. Visual routines. Cognition, 18(1-3):97-159, 12 1984. ISSN 00100277. doi: 10.1016/0010-0277(84)90023-4.
V. Veerabadran and V. R. de Sa. Learning compact generalizable neural representations supporting perceptual grouping. arXiv preprint arXiv:2006.11716, 2020.
Supplementary
A Training and implementation details
All architectures we evaluated within a task were matched in terms of the number of parameters. The input convolutional layer's kernel size is  $7 \times 7$ . Number of channels used by the model remains unchanged across layers and is determined per-model for matching overall number of trainable parameters across models. For PathFinder, we fix the number of channels to be 32 for LocRNN, 21 for hConvGRU and ConvGRU, and 64 for ResNet-30 with a filter size of 9x9 in the intermediate recurrent layers. For Mazes, we fix the number of kernels (d) to be 128 for LocRNN, ConvGRU, hConvGRU, and 100 for ResNet-30 & R-ResNet-30 and the kernel size is fixed to be 5x5. Our readout layers for classification (PathFinder) contain a global average pooling layer followed by a fully-connected layer with output dimensionality of 1 producing the classification logit. We apply binary cross-entropy loss on the logit to train models on PathFinder. On Mazes, we use a  $1 \times 1$  convolution with 1 output channel to produce a binary segmentation map. We use pixel-wise binary cross-entropy to train models on Mazes. These readout layers are used uniformly for all architectures evaluated.