SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The RIR strategy enforces an outer balanced-tree structural restriction on the search space that is generally unsuited for structure-sensitive tasks (and merely done for efficiency). However, the existence of data samples with sequence length less than or equal to chunk size leads to some non-RIR training instances which helps teach the inner recursion to learn to length-generalize. Thus, testing without RIR inference (i.e., running the inner recursion model in the full input) can still lead to length generalization and sometimes be even better due to the lack of balanced tree enforcement.
Table 6: Mean accuracy and standard deviation on SST5 [89] and IMDB [61]. Count. represents counterfactual test split from Kaushik et al. [49] and Cont. represent contrast test split from Gardner et al. [25] We bold the best results. Our models were run 3 times on different seeds. Subscript represents standard deviation. As an example,  $90_{1} = 90 \pm 0.1$
Model	SST5	IMDB
IID	IID	Cont.	Count.
CRvNN	51.7511	91.4712	77.8015	85.383.5
OM	52.302.7	91.690.5	76.985.8	83.687.8
BT-GRC	52.324.7	91.291.2	75.0729	82.8623
BT-GRC OS	51.927.2	90.869.3	75.6821	84.7711
EBT-GRC	52.221	91.471.2	76.1617	84.2912
BBT-GRC	50.983.4	90.571.1	76.377.5	85.381.35
RIR-GRC	50.652.3	90.602.5	77.5330	85.0414
RIR-EBT-GRC	51.017.9	90.864.6	74.2514	82.929.5
Thus, during inference, we have a choice of removing the RIR structure from RIR-EBT-GRC and employing it just like EBT-GRC. We refer to this inference mode as "non-RIR inference". If we make no changes to RIR-EBT-GRC during inference, we refer to it as "RIR inference".
Cost of non-RIR inference: Non-RIR inference can be multiple times slower than RIR inference. Moreover, in most tasks (including LRA), RIR inference works as well as non-RIR inference. So in most cases, RIR inference is a "win-win" choice (several times faster and with similar accuracy).
Gain of non-RIR inference: When a model is working in a structure-sensitive context (such as ListOps-O task), it may find it difficult to learn to length generalize as well through an RIR format. Here, non-RIR inference works much better as seen in the ablation (Table 2)). If RIR training fails to make the model learn the task well, then even for structure-sensitive tasks the difference between the accuracy of RIR inference and non-RIR inference can collapse (it happens for ListOps LRA). RIR inference in structure-sensitive tasks still performs better than MEGA or S4D.
Generally, we can take the choice of using RIR inference as a hyperparameter - the validation metric in the first 5 epochs clues us in which inference strategy to use. For simplicity, we do not tune this choice and by default only use non-RIR inference except for some LRA tasks (Text LRA and Retrieval LRA) where non-RIR inference becomes too slow. For instance, in Text LRA, RIR inference takes around  $\sim 7$  minutes whereas non-RIR inference takes around  $\sim 6$  hours. We find non-RIR inference to mainly make a substantial difference in accuracy in case of ListOps-O and Logical Inference.
We show additional results of our models in sentiment classification (SST5, IMDB) in Table 6 and in natural language inference (MNLI, SNLI) and paraphrase detection in Table 7. Like prior models [82, 13, 77], the natural language results here are mixed. RIR-EBT-GRC does not work well on sentiment classification (it may require more careful tuning), but otherwise keeps up well in sentence-pair matching task - generally even outperforming EBT-GRC. BBT-GRC sometimes does surprisingly well in some OOD sets or stress sets - like the counterfactual set in IMDB and the Break test set from SNLI. We also present an extended comparison in LRA with other competitive models in Table 8. While in the paper we mainly focus on text-based tasks, in theory, RIR-EBT-GRC can be also used for vision-processing tasks. We generally found that we can get greater than  $60\%$  in Sequential CIFAR (LRA Image task). Although that is much lower than SSM-related or Long Convolution models, it still outperforms any pure Transformer-based models. However, in our initial experiments, RIR-EBT-GRC fails in pathfinder tasks. There is more room to explore this area and perhaps look for better ways of creating synergy between SSM-based models and RIR-EBT-GRC. We did not find any notable difference in convergence speed (number of epochs taken) between RIR-EBT-GRC and EBT-GRC.
(a) Time-Accuracy Trade-off
(b) Time-Memory Trade-off
(c) OOD-Accuracy-LRA-Accuracy Trade-off
Figure 2: Scatterplot with Pareto Frontier for different trade offs. We set accuracy to 0 for LRA if the model is too expensive to practically run eg. OM, CRvNN, BT-GRC, EBT-GRC, BT-GRC OS. S4D is S4D, GRC is a recurrent model with GRC cell function, MEG is MEGA, BBT is Balanced Binary Tree GRC, RRG is RIR-GRC, OM is Ordered Memory, CR is CRvNN, BT is BT-GRC, BTO is BT-GRC OS, EBT is EBT-GRC, RRE is RIR-EBT-GRC.
Table 7: Mean accuracy and standard deviation on SNLI [5], QQP, and MNLI [96]. Hard represents the SNLI test set from Gururangan et al. [34], Break represents the SNLI test set from Glockner et al. [27]. Count. represents the counterfactual test set from Kaushik et al. [49].  $\mathrm{PAWS}_{QQP}$  and  $\mathrm{PAWS}_{WIKI}$  are adversarial test sets from Zhang et al. [102], MM is the mismatch dev set from MNLI, ConjNLI is the development set from Saha et al. [80], NegM, NegMM, LenM, LenMM are Negation Match, Negation Mismatch, Length Match, Length Mismatch stress test sets from Naik et al. [69] respectively. Our models were run 3 times on different seeds. Subscript represents standard deviation. As an example,  $90_{1} = 90\pm 0.1$
Models	SNLI Training	QQP Training
IID	Hard	Break	Count.	IID	\( PAWS_{QQP} \)	\( PAWS_{Wiki} \)
(Sequence Encoder Models)
CRvNN	85.32	70.64	55.317	59.86	84.83	34.87	46.66
OM	85.52	70.63	67.49	59.92	84.60	38.17	45.68
BT-GRC	84.91	705	5114	594	84.75	36.917	46.412
BT-GRC OS	84.91	70.36	53.2910	58.63	84.22	37.18	46.36
EBT-GRC	84.74	69.98	55.620	58.11	84.32	36.95	47.55
BBT-GRC	84.44	69.23	65.38	58.14	83.64	42.921	46.712
RIR-GRC	85.04	69.62.9	57.237	59.512	84.24.4	37.310	45.49.8
RIR-EBT-GRC	85.11	70.47	5529	59.614	83.95	37.242	46.85
Models	MNLI Training
Match	MM	ConjNLI	NegM	NegMM	LenM	LenMM
(Sequence Encoder Models)
CRvNN	72.24	72.65	41.710	52.86	53.84.2	6244	63.347
OM	72.53	732	41.74	50.97	51.713	56.533	57.0631
BT-GRC OS	71.71	71.92	41.29	53.22	54.25	65.613	66.79
EBT-GRC	72.12	721	40.930	52.3323	53.2822	64.9210	66.410
BBT-GRC	71.12	71.44	40.912	51.56	52.68	60.615	61.814
RIR-GRC	71.54.8	71.62.4	40.89.9	50.918	51.618	55.219	55.820
RIR-EBT-GRC	71.82	72.34	42.412	53.119	53.325	65.75	67.39
F Pareto Frontier Analysis
We show the plots of three Pareto-frontiers focusing on three trade-offs in Figure 2. In Figure 2(a), we focus on the trade-off between ListOps-O OOD length generalization accuracy (which also correlates with Logical Inference OOD accuracy for most parts) and empirical time cost. In Figure 2(a), we focus on the trade-off between ListOps-O OOD length generalization accuracy and empirical memory cost. In Figure 2(c), we focus on the trade-off between ListOps-O OOD length generalization accuracy and ListOps LRA IID accuracy. We use the empirical time/memory cost from Table 1. In all three subfigures, RIR-EBT-GRC (RRE in the figures) remains as a Pareto-efficient solution that maintains a highly competitive trade-off. S4D, BBT-GRC, and RIR-GRC can win on the time cost and memory cost compared to RIR-EBT-GRC, but with a sharp degradation of OOD performance in structure-sensitive tasks (logical inference, ListOps). While OM, CRvNN, BT-GRC, BT-GRC OS, and EBT-GRC can outperform RIR-EBT-GRC (to an extent) on OOD length generalization accuracy in ListOps and logical inference, they come with much more exorbitant time/memory cost.
G Extended Related Works
RIR-related Approaches: Many approaches have used chunking for hierarchical processing [90, 98, 59, 81, 15, 79, 36, 103, 97, 83, 72, 10, 9, 43, 60] to efficiently tackle long-range inputs. However, they generally neither go all the way in using a full balanced tree recursion (with dynamic layer depth) nor concentrate on remaining competitive on structure-sensitive tasks. One exception is Sliced RNN [101] which introduces a model similar to RIR-GRC but with an older recurrent cell. Some
Table 8: Accuracy on LRA [92]. The results in the top block are copied from the respective papers. For our models, we report the mean of 2 runs. Subscript represents standard deviation. As an example,  $90_{1} = 90 \pm 0.1$
Model	LRA
ListOps	Text	Retrieval
Transformer [60]	37.11	65.21	79.14
Local Attention [92]	15.82	52.98	53.39
Linear Transformer [92]	16.13	65.90	53.09
Reformer [92]	37.27	56.10	53.40
Sparse Transformer [92]	17.07	63.58	59.59
Skinhorn Transformer [92]	33.67	61.20	53.83
BigBird [92]	36.05	64.02	59.29
Performer [92]	18.01	65.40	53.82
Linformer [92]	35.70	53.94	52.27
Longformer [92]	35.63	62.85	56.89
Transformer TLB [18]	37.05	81.88	76.91
H-Transformer-1D [106]	49.53	78.69	63.99
S4-LegS [32]	59.60.7	86.821.3	90.901.5
S4D-Inv [30]	60.183.5	87.342	91.09.1
SGConv [55]	61.45	89.20	91.11
ChordMixer [50]	60.12	88.82	90.17
S5 [86]	62.15	89.31	91.40
Liquid S4 [37]	62.75	89.02	91.20
Long Conv, Exp + Squash [24]	62.2	89.6	91.3
MEGA [60]	63.14	90.43	91.25
LRU [73]	60.28	89.41	89.91
S4D-Inv (reproduction)	59.30	88.180.4	91.17.6
BBT-GRC	55.151	85.690.7	90.781.6
RIR-GRC	20.185.7	86.8941.3	88.2410
RIR-EBT-GRC	59.04	87.971.6	88.995.2
-S4D	59.531	88.131.3	89.743.9
approaches take a block-recurrent [14, 44, 46, 18, 8] approach where some  $k$ -sized chunk is used in a recurrent setup wherein each chunk is processed by a Transformer or some other neural module. These approaches are similar to the RIR framework where the outer recursion is made of some  $k$ -ary chain-tree structured recurrence (instead of balanced-tree recursion), and the inner 'recursion' is implemented for example by a Transformer. A concurrent work Chi et al. [11] also uses a similar structure as RIR but with a Transformer stack replacing the inner inner RvNN.
RvNN: Pollack [76], Goller and Kuchler [29] first introduced Recursive Neural Networks (RvNNs) for bottom-up representation build up based on directed acyclic graphs or trees. Socher et al. [87, 88, 89] extended RvNNs for NLP, and [104, 91, 53, 105] adapted RvNNs with LSTM cells [40]. Le and Zuidema [54], Maillard et al. [64] proposed a chart-based method for simulating bottom-up RvNNs based on CYK algorithm. Drozdov et al. [19, 20], Hu et al. [41, 42] built upon it. Shi et al. [84], Munkhdalai and Yu [68] explored heuristics-tree-based RvNNs including balanced trees. Multiple works explored memory-augmented RNNs for simulating RvNNs [7, 100, 63, 82, 21, 22]. Arabshahi et al. [1] proposed a technique for augmenting recursive cells within RvNNs with a stack-like memory mechanism.
There are other related models [65] that explore structure-sensitive tasks like mathematical equation verification.
H RIR for Language Modeling
Language modeling is today crucial for pre-training and scaling models for competitive performance in several realistic tasks. Nevertheless, directly using RIR-EBT-GRC for language modeling is not straightforward. We discuss some speculative directions we can go forward in re-framing RIR-EBT-GRC for language modeling.
Masked Language Modelingr: In a concurrent work, we propose a method to create contextualized token representations for EBT-GRC or BT-GRC [78] through parent attention (making terminal nodes attend to its direct and indirect parent nodes based on the recursively built-up latent tree). The same principles can be used for RIR-EBT-GRC as well. After that RIR-EBT-GRC can be used for masked language modeling as standard model setups like BERT [17]. It can be also stacked with other types of models like Transformers if needed. Such a setup can also be used as an encoder in a Seq2Seq (Encoder-Decoder) architecture which can be pre-trained with relevant language modeling policies.
Sentence Encoding in a Loop: There is a simple method for using any sentence encoder (including RIR-EBT-GRC) for autoregressive language modeling. At any timestep  $t$ , we can take the sequence of past tokens  $s_{1:t-1}$  as the input and get a sentence encoding - say,  $h_t$  (representing the compression of  $s_{1:t-1}$ ). Then we can predict the next token from  $h_t$  as  $s_t$  by projecting  $h_t$  into a probability distribution over a vocabulary. However, this is a very costly method, because we have to run the whole recursion over the whole generated sequence in every timestep. It would lack the memory reuse capacities of pure recurrent models (using chain-like tree structures), and it would lack the parallelism of Transformers trained with teacher-forcing. Nevertheless, it is a method that can be used given enough resources.
Block-Recurrent Sentence Encoding in a Loop: One way to address the cost of the above approach is to instead take a block-recurrent approach [44]. Combined with RIR-EBT-GRC, this will be similar to taking a three-level nested recursion where the outermost recursion is recurrence and RIR-EBT-GRC will be the inner two-level recursion. In this case, in a timestep RIR-EBT-GRC receives as input not the whole past sequence, but part of the past sequence with the earlier parts recurrently compressed into some finite memory blocks from the outermost recurrence. Thus, the inner recursions (RIR-EBT-GRC) with the "sentence encoding in loop" setup can run with a bounded input size. There could be other creative ways to reduce the overall cost. For example, we can attempt to run the sentence-encoder loop periodically after some intervals intermixing fast (for example recurrently generating from the last sentence encoding instead of re-encoding the whole sequence) and slow generation (re-encoding the whole input sequence).
Parameter Scaling: To be a competitive language model, we need some ways to scale the parameter size of RIR-EBT-GRC. This is challenging for recursive models. The challenge is that the simplest way to increase the parameters is by increasing the hidden state sizes or overall increasing the computation of the recursive cell. Particularly, the increase has to be significant if we want to compete with parameters close to multi-layered models. However, because the recursive cell has to be repeatedly applied for an arbitrary number of times, the overall computation becomes very expensive. Similar challenges also occur for Universal Transformer [71]. One way out of this problem is to stack more layers of RIR-EBT-GRC or alternate models like Transformers or SSMs after creating token-contextualized representations using parent attention [78]. This can also lead to the utilization of a mixture of inductive biases. We can then keep the parameters in the recursive cells themselves to a more manageable level. Another direction to look for can be a modular or mixture-of-expert-based approach. This can allow increasing parameters while keeping computation more manageable via sparse activation. For example, we can attempt to recursively select and activate only a sparse region of parameters (some sparsely selected modules), to keep computation expense controlled while increasing the total parameter budget.
Even after all these, the overall setup would probably not be as fast as Transformer training. But speed or even mere surface capacity to handle big contexts is not everything - and can come at other costs [56, 58].
I Beam Alignment Intuition
In Figure 3, we present a visualization of the String-it solution and in Figure 4, present a visualization of the Beam Alignment solution. In the visualizations, we consider a simple scenario where  $\lfloor \frac{n}{k} \rfloor =$
Figure 3: Visualization of String-it Solution. Beam Sequence x + Beam Sequence y indicates concatenation.
$2, b = 3$  ( $n =$  input sequence length for the recursion,  $k =$  chunk size,  $b =$  beam size). The String-it solution (Figure 3) is to simply concatenate the beam sequences (and add the corresponding scores) from all lists that occur in the same position. For example, in Figure 3, we concatenate the first beam (Beam Sequence 1) from the first list (chunk 1 results) with the first beam (Beam Sequence 4) from the second list (chunk 2 results). Similarly, we concatenate the second beam from the first list with the second beam from the second list, and so on.
However, we want to create "high-scoring" beam combinations with more probability. The String-it solution does not care for that. For example, "Beam Sequence 1 + Beam Sequence 5" would be the second highest-scoring beam combination, but it cannot be ever selected by String-it since Beam Sequence 1 is in the first position of the first list and Beam Sequence 5 is in the second position of the second list - that is, they are not in the same position.
Thus, we propose Beam Alignment where we take a stochastic approach towards the ideal of biasing the preservation of high-scoring combinations. For this approach, instead of immediately applying the String-it solution, we make the beams in each list stochastically 'compete' for their positions. Essentially, we want high-scoring beams to be more likely to occupy more positions in the beam lists from each chunk. This is done by simply sampling a beam per list position (for each of the  $b$  positions) according to the normalized beam score distribution. The result is that the beam lists will be now filled with mostly the higher-scoring beams (See the results after 'Sample' in Figure 4). Next, if we simply apply the String-it solution at this point, it automatically leads to high-scoring combinations, as we wanted, because of the prior sampling. Now there is a possibility of combinations like "beam sequence  $1 +$  beam sequence  $5$ " to arise because the sampling step allows the same beam to occupy multiple positions.
Overall, as we can see, in the simulation of Beam Alignment (Figure 4) the resulting combined beams tend to have higher scores than in the case of the direct application of String-it (Figure 3).
Figure 4: Visualization of Beam-Align Solution. Beam Sequence x + Beam Sequence y indicates concatenation.
J Architecture details
Gated Recursive Cell (GRC): For the Recursive Cell that is the  $R$  function (used within the inner RvNN), we use the GELU-variant of Gated Recursive Cell (GRC) [13] (originally a ReLU variant was proposed in [82]) unless mentioned otherwise. The cell function is formalized below:
$$
\left[ \begin{array}{l} l, \\ r, \\ g, \\ h \end{array} \right] = \operatorname {G e L U} \left(\left[ \begin{array}{l} c h i l d _ {l}; \\ c h i l d _ {r} \end{array} \right] W _ {1} ^ {r e c} + b _ {1}\right) W _ {2} ^ {r e c} + b _ {2} \tag {10}
$$
$$
p = L N (\sigma (l) \odot c h i l d _ {l} + \sigma (r) \odot c h i l d _ {r} + \sigma (g) \odot h) \tag {11}
$$
Here  $-\sigma$  is sigmoid;  $[\cdot ]$  represents concatenation;  $p$  is the parent node representation built from the children childl and childr,  $W_{1}^{rec}\in \mathbb{R}^{2\cdot d\times d_{cell}}$  .  $b_{1}\in \mathbb{R}^{d_{cell}}$  .  $W_{2}^{rec}\in \mathbb{R}^{d_{cell}\times 4\cdot d}$  .  $b_{2}\in \mathbb{R}^{d}$  , and  $l,r,g,h\in \mathbb{R}^d$  .  $LN$  is layer normalization.  $d_{cell}$  is generally set as  $4\times d$
Scorer: Similar to [78], for the scorer we used a two layer MLP as:
$$
e _ {i} = \operatorname {s c o r e r} (x, y) = W _ {2} ^ {s} \operatorname {G e L U} \left(W _ {1} ^ {s} [ x; y ] + b _ {1} ^ {s}\right) + b _ {2} ^ {s} \tag {12}
$$
Here,  $W_{1}^{s}\in \mathbb{R}^{d_{s}\times 2\cdot d_{s}}$ $W_{1}^{s}\in \mathbb{R}^{1\times d_{s}}$ $b_{1}^{s}\in \mathbb{R}^{d_{s}}$ $b_{2}^{s}\in \mathbb{R}$
Sentence Encoder Models: For sentence encoding, the architectural framework we use is the same Ray Chowdhury and Caragea [77]. We use a Siamese dual-encoder setup for sentence-pair tasks as prior works [77].
S4D Pre-Chunk Layer: For the RIR-based models, before starting the RIR-loop, we use a bidirectional S4D [30] layer for pre-chunk processing. Let us represent the basic single layer State space component (without non-linearity) of S4D as  $\mathrm{S4D}_f:\mathbb{R}^{n\times d}\to \mathbb{R}^{n\times d}$  (forward S4D) and  $\mathrm{S4D}_b:\mathbb{R}^{n\times d}\to \mathbb{R}^{n\times d}$  (backward S4D). We express the overall pre-chunk processing layer using S4D below:
$$
x _ {f} = \operatorname {G E L U} \left(\mathrm {S 4 D} _ {f} (x)\right) \tag {13}
$$
$$
x _ {b} = \operatorname {G E L U} \left(\mathrm {S} 4 \mathrm {D} _ {b} (x. \text {f l i p ()})\right). \text {f l i p ()} \tag {14}
$$
$$
x _ {c a t} = \left[ x _ {f}; x _ {b} \right] \tag {15}
$$
$$
o u t = \operatorname {s i g m o i d} \left(x _ {c a t} W _ {u}\right) \odot \left(x _ {c a t} W _ {v}\right) + x \tag {16}
$$
$$
i n p = L N \left(o u t W _ {i n i t}\right) \tag {17}
$$
Here  $W_{u}, W_{v} \in \mathbb{R}^{2d \times d}$ ,  $W_{init} \in \mathbb{R}^{d \times d}$ , and LN is Layer Normalization. Eqn. 16 is the application of gated linear unit [16] and a residual like Gu et al. [30] for feature mixing, and Eqn. 17 is the initial leaf transformation layer that is used in prior works [82, 13, 77].
S4D: When using pure (multi-layered S4D without RIR-models) S4D models, we use similar settings as Gu et al. [30]. We found it more beneficial to use the fused kernel for bidirectionality within their codebase when using S4D (setting 'bidirectional=True' in their S4 function argument).
K Hyperparameter details
For prior sentence encoder models, we use the same hyperparameters as [77] for all the datasets.  $d_{s}$  for EBT-GRC is set as 64. For RIR-EBT-GRC, RIR-inference was only used for Text LRA and Retrieval LRA for computational efficiency. We always use RIR-inference for RIR-GRC because we found to be better by validation accuracy. We use chunk size as 30 for all RIR-models in all tasks, except Retrieval LRA where we set chunk size as 5 for computational efficiency (otherwise it gets slow). For RIR-EBT-GRC, we use a beam size of 7 for all tasks except Retrieval LRA where we use a beam size of 5. Every other hyperparameters are unchanged from BT-GRC for RIR-models
<sup>8</sup>https://github.com/HazyResearch/state-spaces
or BBT-GRC for the earlier tasks. For the new tasks, we share the hyperparameters for MNLI with that of SNLI, QQP, and Retrieval LRA (except that for Retrieval LRA we use a smaller hidden state size - 128 for parameter control and no pre-trained embeddings). We share the hyperparameters for IMDB and Text LRA (except that for Text LRA, we use smaller hidden state size - 128; again, for parameter control, and no pre-trained embedding). We share the hyperparameters of ListOps-O, ListOpsMix, Logical Inference, and ListOps LRA for all models. We initialize S4D (whether when using the pure S4D model or S4D for pre-chunk processing) in S4D-Inv mode and use billinear discretization. For LRA tasks, we use the same hyperparameters as Gu et al. [30] for S4D. We use the same hyperparameters as used for ListOps LRA Ma et al. [60] for MEGA as run in ListOps-O and Logical Inference. Note that all the natural language tasks (except LRA-related ones) are trained with frozen 840B Glove Embeddings [75] same as Ray Chowdhury and Caragea [77]. All models were trained on a single Nvidia RTX A6000.