SlowGuess/NeurIPS_2023_OCR huggingface dataset:
jraych2@uic.edu
cornelia@uic.edu
Binary Balanced Tree Recursive Neural Networks (BBT-RvNNs) enforce sequence composition according to a preset balanced binary tree structure. Thus, their nonlinear recursion depth (which is the tree depth) is just  $\log_2n$  ( $n$  being the sequence length). Such logarithmic scaling makes BBT-RvNNs efficient and scalable on long sequence tasks such as Long Range Arena (LRA). However, such computational efficiency comes at a cost because BBT-RvNNs cannot solve simple arithmetic tasks like ListOps. On the flip side, RvNN models (e.g., Beam Tree RvNN) that do succeed on ListOps (and other structure-sensitive tasks like formal logical inference) are generally several times more expensive (in time and space) than even Recurrent Neural Networks. In this paper, we introduce a novel framework — Recursion in Recursion (RIR) to strike a balance between the two sides - getting some of the benefits from both worlds. In RIR, we use a form of two-level nested recursion - where the outer recursion is a  $k$ -ary balanced tree model with another recursive model (inner recursion) implementing its cell function. For the inner recursion, we choose Beam Tree RvNNs. To adjust Beam Tree RvNNs within RIR we also propose a novel strategy of beam alignment. Overall, this entails that the total recursive depth in RIR is upper-bounded by  $k\log_k n$ . Our best RIR-based model is the first model that demonstrates high ( $\geq 90\%$ ) length-generalization performance on ListOps while at the same time being scalable enough to be trainable on long sequence inputs from LRA (it can reduce the memory usage of the original Beam Tree RvNN by hundreds of times). Moreover, in terms of accuracy in the LRA language tasks, it performs competitively with Structured State Space Models (SSMs) without any special initialization - outperforming Transformers by a large margin. On the other hand, while SSMs can marginally outperform RIR on LRA, they (SSMs) fail to length-generalize on ListOps. Our code is available at: https://github.com/JRC1995/BeamRecursionFamily/.
Non-linear Recurrent Neural Networks (RNNs) [23, 40] have the potential to enable powerful nonlinear sequential dynamics that dynamically adapt with the input length. However, this ability that makes them powerful also causes several problems. First, RNNs are slow because of their non-linear sequential operations that scale linearly with the input sequence length. Second, backpropagation through arbitrarily large non-linear recurrence depth (depending on the input length) can lead to vanishing/exploding gradients [39, 3] - or can have stability issues overall [67]. Potentially, because of the combination of these factors, it was found that fast convolution-based models [99, 48, 51, 16, 57, 26, 2, 47] with fixed non-linear depth could often do just as well as an RNN. Transformers [94] are built upon this trajectory replacing bounded kernel window in convolution with an attention
mechanism that can dynamically compute interaction weights for input tokens at any arbitrary distance.
However, despite the tremendous success of Transformers in all manner of tasks, they tend to struggle in hierarchical structure-sensitive tasks that require dynamic adaptation to input complexity and modeling of long-range dependencies [93, 82, 35, 92]. Recently, a family of models closely connected to linear RNNs [31, 73, 86, 66, 33, 30] have been emerging. Most of these models greatly outperform Transformers in long-range datasets [31, 86]. Instead of adding non-linearity in every recurrent step, they simply stack a fixed number of linear recurrences with non-linearities in between. However, in doing so these models lose the ability to adapt the non-linearity depth to input. This can make it harder to tackle tasks requiring adaptation to unseen lengths of nested non-linear operations.
In this paper, we take a step back and look again at non-linear RNNs. Non-linear RNNs, more generally, can be seen as a special case of (non-linear) Tree Recursive Neural Networks (Tree-RvNNs) [76, 87, 89, 91]. Tree-RvNNs can process sequences in an arbitrary tree-structured order. Non-linear RNNs are a special case of Tree-RvNNs following a chain-structured tree. However, such a chain-structured tree is an extreme case - requiring the maximum tree traversal height. On the other side of the extreme are balanced tree models. For instance, the height of a  $k$ -ary balanced tree with  $n$  terminal nodes (the input length for our case) would be  $\log_k(n)$ . Balanced Tree-RvNNs [68, 84] can then logarithmically scale with the sequence length  $(n)$ . To an extent, this can greatly mitigate the original issues of non-linear RNNs while still preserving some degree of ability to dynamically adapt with input length. Fascinatingly, we find that a simple model like a Balanced Binary Tree RvNN (BBT-RvNN) with a modern recursive cell [82, 13] can come close to the performance of the recent strong linear RNNs or Structured State Space Models (SSMs) without much tuning or any specialized initialization on Long Range Arena (LRA) text datasets [92] - thereby, greatly outperforming Transformers in general.
Nevertheless, while BBT-RvNNs can be quite fast and scalable, they are incapable of excelling (or even passing) simple structure-sensitive tasks like ListOps<sup>1</sup> [70] or formal logical inference [6] in a length-generalizable manner. Transformers were already shown to underperform in these tasks [82, 93], and we find that S4D [30] (a representative of linear RNN/SSM model family) and MEGA [60] (a hybrid of Transformer and SSM) also struggle in these tasks. Looking back at Tree-RvNNs, there are several models [82, 13, 77] that dynamically learn the tree structure from the input (instead of using some enforced heuristics) and determine the order of operations accordingly. They have been able to perform well in structure-sensitive tasks. Thus, BBT-RvNNs enjoy high speed and computational efficiency but at the cost of systematic issues in modeling task structures as revealed in ListOps or logical inference tasks, whereas in sharp contrast, the more powerful RvNNs [82, 13, 77] enjoy their strong length-generalizing accuracy in ListOps, logical inference, and reasonable performance in natural language domain but at the cost of high computational resource and inability to effectively train on datasets like LRA.
To address the above challenges, we propose a new framework called Recursion in Recursion (RIR) in which we have a nested level of recursion. We use a balanced  $k$ -ary Tree-ordered recursion in the outer loop and we use a strong RvNN (particularly we use a Beam Tree RvNN-based model [77]) in the inner loop to process the  $k$  arguments sent in from the outer loop. The inner RvNN can scale linearly with input size but now that size is bounded by a fixed constant  $k$ . Thus, the total non-linear depth will be bounded by a linear factor of  $k \log_k(n)$  (the total outer loop being  $\log_k(n)$  and the inner loop being  $k$ ). RIR makes Beam Tree RvNN (BT-RvNN) much faster and trainable on LRA but at the cost of some accuracy loss. On the flip side, RIR can be seen as a slower version of more basic balanced tree models but at the gain of length-generalizability in ListOps and general accuracy gain in several other datasets. In addition, we identify an issue related to beam alignment in incorporating BT-RvNN within RIR and propose a strategy to circumnavigate it.
Tree Recursive Neural Networks (Tree-RvNN): Here we focus on constituency Tree-RvNNs [89, 91] that build up higher-order constituent representations of a given input in a bottom-up manner. Particularly, given an input sequence of vectors, Tree-RvNNs treat them as terminal nodes of a latent
<sup>1</sup>An example sample of ListOps - Input: [SM[SM[SM[MAX 56]2]0]5086]; Output: 7. SM is modulo-10 summation.
tree to be built up. The non-terminal nodes of the tree are filled sequentially by composing their immediate child nodes using some recursive cell function  $-R: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}^d$  (we consider binary trees here so  $R$  will always have only 2 children as arguments). Given some input sequence of vectors  $7 + 9 \times 5 - 2$  (assume each of the symbols correspond to some vector  $\in \mathbb{R}^d$ ), an example of Tree-RvNN-like processing can be expressed as:
$$
R (R (R (7, +), R (R (9, \times), 5)), R (-, 2)) \tag {1}
$$
Let us say  $R$  is considered as a single non-linear neural layer. In that case, the total non-linear depth of this processing can be interpreted as the maximum nested  $R$  operations which is 4 in Eqn. 1. RNNs can be now seen as a special case where the order of operation is always like this:
$$
R (R (R (R (R (R (R (h _ {0}, 7), +), 9), \times), 5), -), 2) \tag {2}
$$
Here,  $h_0$  is the initial hidden state for RNNs. As we can see, because of their chain-like structure RNNs will have the maximum possible recursion depth within the Tree-RvNN framework. In this case, the depth is 7. Another special case is BBT-RvNN (Balanced Binary Tree RvNN) which follows an enforced binary-balanced tree-based order of operations and can be expressed as:
$$
R (R (R (7, +), R (9, \times)), R (R (5, -), 2)) \tag {3}
$$
Here the maximum non-linearity depth is only 3 (i.e.,  $\sim \log_2(n)$ ; where  $n = 7$ ). With longer sequences the differences between RNNs and BBT-RvNNs will grow bigger because BBT-RvNNs can scale logarithmically to the input length.
Greedy Search Tree RvNN: If our Tree-RvNN does not employ any fixed heuristic tree structure (as in RNNs or BBT-RvNNs) and if we are not relying on an external system or user to provide tree structures, we have to decide on some policy to automatically induce the latent tree structure from raw sequences. Greedy search Tree RvNNs provide us with one way of doing that. The basic idea of Greedy Search Tree RvNN can be found in multiple works [12, 88] and is related to easy-first parsing [28]. A general framework for greedy search parsing is as follows. Assume we are at some recursion step  $t$ . Let the input sequence of node representations in that step be  $(h_1^t, h_2^t, \ldots, h_K^t)$  (where  $h_i^t \in \mathbb{R}^d$ ). Assume we have the recursive binary cell  $R$  as before and a neural scorer function of the form scorer:  $\mathbb{R}^{2 \times d_s} \to \mathbb{R}$  (where  $d_s < d$ ; the role of  $d_s$  will be clarified). Next, we consider all possible node pairs that can be built into a parent node. Note that for simplicity we maintain a projective tree constraint as is often standard [12, 38, 13] - so only contiguous node pairs  $(h_i^t, h_{i+1}^t)$  are considered as candidates for being composed into a higher-level parent node (whose children will be those nodes). The scorer assigns a scalar score to every candidate node pair as:  $e_i = scorer(h_i^t[0:d_s], h_{i+1}^t[0:d_s])$ . The slicing above is done for efficiency reasons (we set  $d_s = 64$  and generally keep it smaller than  $d$ , i.e.,  $d_s < d$ ) following [78]. Let  $e_{1:k}$  represent the sequence of all scores. Assume  $j$  is the position of the first node in the maximum scoring pair, i.e.,  $j = \text{argmax}(e_{1:k})$ . Now the update rule for the input to the next recursion will be:
$$
h _ {i} ^ {t + 1} = \left\{ \begin{array}{l l} h _ {i} ^ {t} & i <   j \\ R \left(h _ {i} ^ {t}, h _ {i + 1} ^ {t}\right) & i = j \\ h _ {i + 1} ^ {t} & i > j \end{array} \right. \tag {4}
$$
The iterative application of this rule can achieve both tree-parsing (through scorer) and sequence composition (by application of  $R$ ) simultaneously. Finally, we will have the root representation left which can be used as a sentence encoding (the intermediate non-terminal nodes can be also used if needed). Note that the framework we describe already incorporates the efficiency fixes introduced in [78]. Works [12, 77] prior to [78] made the scorer's input to be  $R(h_{i}^{t}, h_{i+1}^{t})$  which requires running expensive recursive cells for all possible node pair siblings (not just the chosen one) in parallel and they use no slicing. A problem with directly using the above framework is that the argmax for selecting  $j$  is non-differentiable. So it becomes impossible to train the scorer. Various techniques have been used to address this problem - e.g., with autoencoder loss [88], gumbel softmax [62, 45] with Straight-through Estimation (STE) [4], or with SPIGOT [74]. Although not all of them have been exhaustively explored, STE gumbel softmax [12] which is one of the more popular approaches has been found to fail in structure-sensitive tasks [70, 38]. However, viewing the above family of approaches as a greedy search approach - opens up another way that is using simple beam search replacing argmax with a top- $k$  operator for beam search [77]. We discuss this below.
Beam Search Tree RvNN: Extending Greedy Search Tree RvNNs with beam search allows for keeping multiple hypotheses or beams. Each beam would be one possible induced tree structure with its node representations and a corresponding beam score (addition of the log-softmax of  $e_j^t$  for each iteration  $t$  in that beam). Setting up a beam-score-based softmax-based marginalization over the root representations of each beam in the end allows us to create a sentence vector representation:  $\sum_{i} \frac{\exp(s_i)}{\sum_{j} \exp(s_j)} \cdot b_i$  where  $b_i$  is the root representation of beam  $i$  and  $s_i$  is the beam score. The final softmax can allow competition in the top  $B$  selected beams (where  $B$  is the beam width) which can allow meaningful gradient signals to reach the scorer. Ray Chowdhury and Caragea [77] showed that this is enough to make the former framework succeed in ListOps and other logical inference tasks. We use the efficient variant Beam Tree RvNN (BT-RvNN), i.e., Efficient Beam Tree RvNN (EBT-RvNN) which is BT-RvNN built on the specific greedy framework described above rather than the prior frameworks [12] that we contrasted. Both BT-RvNN and EBT-RvNN use stochastic Top-K [52]. More details on BT-RvNN can be found in [77] and more details on EBT-RvNN can be found in [78].
Balanced Tree RvNNs: Above, we focused on Binary (2-ary) Tree RvNNs. Here we consider a more general  $k$ -ary Tree RvNN. The non-terminal nodes of  $k$ -ary Tree RvNN can have at most  $k$  children. Since we defined the recursive cell  $R$  to have two parameters, we now introduce a generalized form of the recursive cell  $-RK: \mathbb{R}^{k \times d_h} \to \mathbb{R}^{d_h}$ .  $RK$  takes a sequence of  $k$  vectors as input and outputs a single vector. Essentially  $RK$ , by itself, operates like a sentence encoder. Next, we introduce a balanced tree RvNN model that uses a  $k$ -ary recursive cell.
We call the balanced tree variant of a  $k$ -ary Tree RvNN as Balanced  $k$ -ary Tree-RvNNs (BKT-RvNNs). Balanced Binary Tree RvNNs (BBT-RvNNs) is a special case of BKT-RvNN when  $k = 2$ . In each recursion, BKT-RvNNs make non-overlapping chunks of size  $k$  (so we can also say  $k$  is the "chunk size" henceforth) and then compress each chunk into a single vector using  $RK$ . This is done in every recursion until there is only one vector left - which becomes the root node representation, i.e., the final sentence encoding. Let us say, we have a sequence  $(r_{1:n} \in \mathbb{R}^{n \times d_h})$  of the form below:
$$
r _ {1: n} = r _ {1}, r _ {2}, \dots , r _ {n} \tag {5}
$$
Assume that  $n$  is divisible by  $k$  - if not, we can always use padding to extend the sequence length to make it divisible. We can then create a chunked sequence  $(c_{1:\lfloor \frac{n}{k}\rfloor}\in \mathbb{R}^{\lfloor \frac{n}{k}\rfloor \times k\times d_h})$  as below:
$$
c _ {1} = \left(r _ {1}, r _ {2}, \dots , r _ {k}\right), \quad c _ {2} = \left(r _ {k + 1}, \dots , r _ {2 k}\right), \dots , \quad c _ {\lfloor \frac {n}{k} \rfloor} = \left(r _ {\left(\lfloor \frac {n}{k} \rfloor - 1\right) \cdot k + 1}, r _ {\left(\lfloor \frac {n}{k} \rfloor - 1\right) \cdot k + 2}, \dots , r _ {n}\right) \tag {6}
$$
Each chunk  $(c_{i})$  can be then processed independently (and simultaneously) by  $RK:\mathbb{R}^{k\times d_h}\to \mathbb{R}^{d_h}$  to compress the  $k$ -sized chunks into single vectors:
$$
p _ {1} = R K \left(c _ {1}\right), p _ {2} = R K \left(c _ {2}\right), \dots , p _ {\lfloor \frac {n}{k} \rfloor} = R K \left(c _ {\lfloor \frac {n}{k} \rfloor}\right) \tag {7}
$$
Thus the sequence of chunks,  $c_{1: \lfloor \frac{n}{k} \rfloor} \in \mathbb{R}^{\lfloor \frac{n}{k} \rfloor \times k \times d_h}$  turns into a sequence of vectors  $p_{1: \lfloor \frac{n}{k} \rfloor} \in \mathbb{R}^{\lfloor \frac{n}{k} \rfloor \times d_h}$ .  $p_{1: \lfloor \frac{n}{k} \rfloor}$  then becomes the input to the next recursion. We terminate the recursion when only one item is left in the sequence. Since in every recursion the sequence size is cut by  $k$  times, the total recursions taken (until a sequence of size 1 is left) would be  $\lceil \log_k n \rceil$  assuming  $n$  as the original sequence length.
3 Recursion in Recursion
Here, we introduce our framework - Recursion in Recursion (RIR). RIR is essentially an implementation of BKT-RvNN. The question to ask when using BKT-RvNN is what to use to implement  $RK$ . In case of RIR, we use another Tree-RvNN to implement  $RK$ . This results in a two-level nested recursion as discussed below. We present a visualization in Figure 1 and explain the two levels below.
Outer Recursion: The outer recursion (outer RvNN) is the BKT-RvNN setup as discussed.
Inner Recursion: In RIR, the inner recursion is another RvNN (inner RvNN). In principle, we can use nearly anything for the inner RvNN or the  $RK$  function in general - even a Transformer [94] or
Figure 1: Visualization of the Recursion In Recursion (RIR) framework. The sequence  $(x_{1}, x_{2}, \ldots, x_{9})$  is the input sequence of vectors. The RvNN block indicates the inner RvNN.  $R$  is the recursive cell within the inner RvNN. The inner RvNN block is working within a balanced  $k$ -ary Tree like structure (here  $k = 3$ ) - effectively serving as the recursive cell within BKT-RvNN.
an RNN. Using an RNN makes the whole model similar to Sliced RNN [101]. Our main RIR-based model uses EBT-RvNN as the inner RvNN. We motivate this choice below.
Motivation for EBT-RvNN:
1. BT-RvNN performs quite well on logical inference and ListOps while maintaining reasonable natural language performance - a feature lacking in plain RNNs, Gumbel-Tree variants, or other simpler RvNNs [77]. EBT-RvNN is a more promising and efficient variant of BT-RvNN [78].
2. EBT-RvNN as an inner RvNN (after some adjustments to be discussed in the next section) offers us a clean interpretation of the RIR model as an approximation of EBT-RvNN searching in an additionally constrained space of tree structures (constrained by the outer balanced tree recursion). While other ListOps competent models can be also forcibly used as an inner RvNN [13, 82], there will be less of an interpretation of their modeling choices under RIR.
Time Complexity Analysis: The time complexity of the outer RvNN (i.e. BKT-RvNN) in RIR (if the inner RvNN were  $\Theta(1)$ ) is  $\Theta(log_k(n))$  following what we already discussed. Let us say that the time complexity of the inner RvNN is some  $\Theta(f(n))$  if  $n$  is the input sequence length. Given these factors, the overall time complexity of an RIR model will be  $\Theta(\log_k(n)f(min(k,n)))$ . While inner RvNNs like EBT-RvNN can be slow by itself, within the RIR framework its recursive depth is now bound by the chunk size  $k$  instead of the sequence length  $n$ . The chunk size  $k$  is a hyperparameter that is preset as a constant before training. It remains constant during training no matter how big the sequence length  $n$  becomes. As such although the inner RvNN can be expensive it now becomes bound by constant factors during training. This makes the RIR model with EBT-RvNN as inner RvNN much more efficient than EBT-RvNN itself. The same is also true for other possible implementations of an inner RvNN within the RIR framework.
However, the inner recursion with EBT-RvNN can be still relatively slow and add large overheads. Thus, the bigger the  $k$ , the slower the overall model will be. On the other hand, lower  $k$  pushes the model closer to BBT-RvNN and makes it faster. When  $k = 2$ , the model would just collapse into BBT-RvNN, and when  $k = \infty$ , the model reduces into just EBT-RvNN.
3.1 Beam Alignment
Context: There are a few challenges to naively integrating EBT-RvNN in the RIR framework. First, note that we initially defined  $RK$  as of the form  $RK: \mathbb{R}^{k \times d_h} \to \mathbb{R}^{d_h}$ . If we keep it that way then we have to marginalize the beams at the end of every inner recursive loop to get a single vector output instead of beams of vectors. This can make it harder to extract discrete structures from the
beam search process if needed for interpretability or for sending top-down attention signals for token contextualization (as used in [78] with parent attention). Also, it is in conflict with the second motivation for using EBT-RvNN. Moreover, keeping the beams active instead of marginalizing them every so often can allow for more robustness and recovery from error [20]. Now, if the beams are kept active, there are two main changes.
1. The input sequence in every outer recursion would belong to  $\mathbb{R}^{b \times n \times d_h}$  where  $b$  is the beam size. The output sequence after the outer recursion should be in  $\mathbb{R}^{b \times \left\lfloor \frac{n}{k} \right\rfloor \times d_h}$ . In addition, we need to also maintain an input of beam scores  $\in \mathbb{R}^b$  and an output of the same form.
2. Each chunk that is fed to  $RK$  should be now of the form  $\mathbb{R}^{b \times k \times d_h}$ . In addition, the beam scores  $\in \mathbb{R}^b$  are also fed to  $RK$ . The output of  $RK$  should be in the form of a tuple:  $(BS, BR)$  where  $BS \in \mathbb{R}^b$  represents the list of beam scores for each beam and  $BR \in \mathbb{R}^{b \times 1 \times d_h}$  is the list of corresponding chunk-level sentence encodings² for each beam. Overall  $RK$  is to be redefined as:  $RK : \mathbb{R}^b \times \mathbb{R}^{b \times k \times d_h} \to \mathbb{R}^b \times \mathbb{R}^{b \times 1 \times d_h}$ . The first input argument represents the input beam scores, the second represents the beam of chunks.
Problem: Now, the output of each  $RK$  is a sequence  $\in \mathbb{R}^{b \times 1 \times d_h}$  and beam scores  $\in \mathbb{R}^b$ . Since there would be  $\left\lfloor \frac{n}{k} \right\rfloor$  chunks (if the input sequence length to that level of outer recursion is  $n$ ), there will be simultaneous  $\left\lfloor \frac{n}{k} \right\rfloor$  sequences  $\in \mathbb{R}^{b \times 1 \times d_h}$  and similarly,  $\left\lfloor \frac{n}{k} \right\rfloor$  lists of beam scores  $\in \mathbb{R}^b$ . The challenge is to answer how we get a single sequence  $\in \mathbb{R}^{b \times \left\lfloor \frac{n}{k} \right\rfloor \times d_h}$  and a single list of beam scores  $\in \mathbb{R}^b$  from this circumstance to serve as inputs to the next outer recursion.
For simplicity, let us say  $\lfloor \frac{n}{k} \rfloor = 2$  and as such, we have two chunk outputs:  $(BS_{1}, BR_{1})$  and  $(BS_{2}, BR_{2})$ . The question that is now raised is: How do we combine the different beams of the different output tuples for the input to the next outer recursion?
String-it Solution: One answer can be to just string together (concatenate) the corresponding beams in  $BR_{1}$  and  $BR_{2}$  as:  $BR_{new} = \text{concatenate}([BR_{1}, BR_{2}], dim = 1)$ . In other words, we concatenate  $BR_{1}[i]$  with  $BR_{2}[i]$  for all  $i \in \{0, 1, \ldots, b - 1\}$ . Here,  $BR_{new} \in \mathbb{R}^{b \times 2 \times d}$ . The beam scores can be summed up for accumulation:  $BS_{new} = BS_{1} + BS_{2}$ .