SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jishnu Ray Chowdhury Cornelia Caragea
jraych2@uic.edu cornelia@uic.edu
Beam Tree Recursive Neural Network (BT-RvNN) was recently proposed as an extension of Gumbel Tree RvNN and it was shown to achieve state-of-the-art length generalization performance in ListOps while maintaining comparable performance on other tasks. However, although better than previous approaches in terms of memory usage, BT-RvNN can be still exorbitantly expensive. In this paper, we identify the main bottleneck in BT-RvNN's memory usage to be the entanglement of the scorer function and the recursive cell function. We propose strategies to remove this bottleneck and further simplify its memory usage. Overall, our strategies not only reduce the memory usage of BT-RvNN by 10 - 16 times but also create a new state-of-the-art in ListOps while maintaining similar performance in other tasks. In addition, we also propose a strategy to utilize the induced latent-tree node representations produced by BT-RvNN to turn BT-RvNN from a sentence encoder of the form  $f: \mathbb{R}^{n \times d} \to \mathbb{R}^d$  into a token contextualizer of the form  $f: \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d}$ . Thus, our proposals not only open up a path for further scalability of RvNNs but also standardize a way to use BT-RvNNs as another building block in the deep learning toolkit that can be easily stacked or interfaced with other popular models such as Transformers and Structured State Space models. Our code is available at the link: https://github.com/JRC1995/BeamRecursionFamily.
Recursive Neural Networks (RvNNs) [63] in their most general form can be thought of as a repeated application of some arbitrary neural function (the recursive cell) combined with some arbitrary halting criterion. The halting criterion can be dynamic (dependent on input) or static (independent of the input). From this viewpoint, nearly any neural network encoder in the deep learning family can be seen as a special instance of an RvNN. For example, Universal Transformers [13] repeat a Transformer [85] layer block as a recursive cell and adaptively halt by tracking the halting probability in each layer using some neural function [25, 3]. Deep Equilibrium Models (DEQ) [2] implicitly "repeat" a recursive cell function using some root-finding method which is equivalent to using the convergence of hidden states dynamics as the halting criterion. As Bai et al. [2] also showed, any traditional Transformer - i.e. stacked layer blocks of Transformers with non-shared weights can be also equivalently reformulated as a recursive repetition of a big sparse Transformer block with the halting criterion being some preset static upperbound (some layer depth set as a hyperparameter).
A broad class of RvNNs can also be viewed as a repeated application of a Graph Neural Network (GNN) [68, 91] - allowing iterative message passing for some arbitrary depth (determined by the halting criterion). Transformer layers can be seen as implementing a fully connected (all-to-all) graph with sequence tokens as nodes and attention-based edge weights. In natural language processing, we often encounter the use of stacks of GNNs (weight shared or not) to encourage message-passing through underlying linguistic structures such as dependency parses, constituency parses, discourse structures, abstract meaning representations, and the like [80, 39, 47, 59, 11, 90, 94].
In many cases, such models are implemented with some static fixed number of layers or iterations. However, learning certain tasks in a length-generalizing fashion without prior knowledge of the sequence length distribution should require a dynamic halting function. For example, consider a simple arithmetic task:  $7 \times (8 + 3)$ . It is first necessary to process  $8 + 3$  before doing the multiplication in the next step. A task like this can require structure-sensitive sequential processing where the sequential steps depend on the input length and nested operators. Although the above example is simple enough, we can have examples with arbitrary nested operators within arbitrarily complex parenthetical structures. Datasets like ListOps [57] and logical inference [7], among others, serve as a sanity check for any model's ability to perform similar structure-sensitive tasks. Natural language tasks, on the other hand, can sometimes hide systematic failures because of the distributional dominance of confounding statistical factors [54]. While, in some algorithmic tasks, Universal Transformer [13] can perform better than vanilla Transformers, both can still struggle in length generalization in logical inference or ListOps [71, 84]. This brings us to explore another alternative family of models—Tree-RvNNs with a stronger inductive bias than Universal Transformer but still possessing a form of dynamic halt that we discuss below.
Tree Recursive Neural Networks: Tree-RvNNs [79, 63, 77] (more particularly, constituency Tree RvNNs) treat an input sequence of vectors as the terminal nodes of some underlying latent tree. From this perspective, Tree-RvNNs sequentially build up the tree nodes in a bottom-up manner. Each node (terminal or non-terminal) in the tree will represent some span within the sequence (for text sequences, these would be words, phrases, clauses, and the like). The root node will represent the whole input sequence and thus can be used as an encoded representation of the whole input (for sentences, it can be treated as a sentence encoding). The nodes are generally in the form of vectors.
The non-terminal nodes are built bottom-up through the repeated use of some recursive cell composition function, say rec, starting from height 1 (height 0 being the terminal nodes) all the way up to the root. When building the representation of some node  $p$  at height  $t$ , the Tree-RvNN uses the rec function to take all the immediate child nodes of  $p$  (from the previous heights 0 to  $t - 1$ ) as input arguments and outputs the representation of  $p$ . Typically, we only consider binarized trees so the number of children as arguments will be a constant (two) for rec. Thus, it can be represented in the form  $\text{rec} : \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}^d$ . Generally, we also constrain our consideration to only projective tree structures. This implies that the model enforces only temporally contiguous node pairs be considered as candidate siblings for any parent node. Finally, when the root representation is built the computation terminates. Therefore, reaching the root is the halting criterion. Since the tree structure is dependent on the input and the halting is dependent on the tree structure, the halting mechanism here is dynamic and thus can adapt with input complexity.
This approach has the potential to model mereological (part-whole) structures at different hierarchical scales in a systematic and dynamic manner. As an example, let us consider that we have an input such as  $(4\times 3) + (7 + (8\times 9))$ . Then, in the ideal case, the Tree-RvNN order of operation would potentially yield: rec(rec(rec(4,  $\times$  )3), +), rec(rec(7, +), rec(rec(8,  $\times$  ), 9))). As we can see, Tree-RvNNs have the potential to model sensitive order of operations considering input hierarchies which can be botched by unconstrained attention without careful inductive biases [12, 10, 29] or extensive pre-training. We present an extended related work review in Appendix D.
Motivation for Tree-RvNNs: Although, at first glance Tree-RvNNs as formalized above may appear to have too strong of an inductive bias, there are a few reasons to be motivated to empirically explore the vicinity of these models:
1. Familiar Recurrent Neural Networks (RNNs) [19, 33] are special cases of Tree-RvNNs that follow an enforced chain-structured tree. So at the very least Tree-RvNNs have more flexibility than RNNs which have served as a fairly strong general model in the pre-Transformer era; and even now, optimized linear variants of RNNs [53] are making a comeback out-competing Transformers in long range datasets [26, 27, 60].
This is done mainly to simplify the search space but assuming projective structures is also standard fare in Natural Language Processing (NLP) [41, 96].
${}^{2}$  “temporal” refers to the original sequential order. The original order is preserved in every iteration.
${}^{3}$  Similar inductive biases can be favorable in computer vision as well [76, 30, 74]
4Overall, Tree-RvNNs, as described above, can be also considered as a special case of DAG-GNN [83] where similar principles can apply. Exploring bottom-up tree node building through a more flexible space of DAG-structures can be a direction to consider in the future.
2. Often Tree-RvNNs are the core modules behind models that have been shown to productively length-generalize or compositionally generalize [71, 10, 46, 32, 5, 45, 52] in settings where other family of models struggle (at least without extensive pre-training).
3. The recursive cell function of Tree-RvNNs can still allow them to go outside their explicit structural constraints by effectively organizing information in their hidden states if necessary. The projective tree-structure provides only a rough contour for information flow through Tree-RvNNs which Tree-RvNNs can learn to go around just as an RNN can [7]. So in practice some of these constraints can be less concerning than one may think.
Issues: Unfortunately, the flexibility of Tree-RvNNs over RNNs comes at a cost. While in RNNs we could just follow on the same chain-structure tree (left-to-right or right-to-left) for any input, in Tree-RvNNs we have to also consider some way to dynamically induce a tree-structure to follow. This can be done externally—that is, we can rely on human inputs or some external parser. However, neither of them is always ideal. Here, we are focusing on complete Tree-RvNN setups that do their own parsing without any ground-truth tree information anywhere in the pipeline. Going this direction makes the implementation of Tree-RvNNs more challenging because it is hard to induce discrete tree structures through backpropagation. Several methods have been proposed to either induce discrete structures [31, 9, 66, 61], or approximate them through some continuous mechanism [10, 95, 72, 71]. Nevertheless, all these methods have their trade offs - some do not actually work in structured-sensitive tasks [57, 9, 72], others need to resort to reinforcement learning and an array of optimizing techniques [31] or instead use highly sophisticated architectures that can become practically too expensive in space or time or both [51, 71, 10]. Moreover, many of the above models [10, 31, 9] have been framed and used mainly as a sentence encoder of the form  $f: \mathbb{R}^{n \times d} \to \mathbb{R}^d$  ( $n$  being the sequence size). This formulation as sentence encoder limits their applicability as a deep learning building block - for example, we cannot use them as an intermediate block and send their output to another module of their kind or some other kind like Transformers because the output of a sentence encoder will just be a single vector.
Our Contributions: Believing that simplicity is a virtue, we direct our attention to Gumbel-Tree (GT) RvNN [9] which uses a simple easy-first parsing technique [23] to automatically greedily parse tree structures and compose sentence representations according to them (we provide a more extended discussion on related works in Appendix). Despite its simplicity, GT-RvNN relies on Straight-Through Estimation (STE) [4] which induces biased gradient, and has been shown empirically to fail in structure-sensitive tasks [57]. Yet, a recent approach - Beam Tree RvNN (BT-RvNN) [66] - promisingly shows that simply using beam search instead of the greedy approach succeeds quite well in structure-sensitive tasks like ListOps [57] and Logical Inference [7] without needing STE. Nevertheless, BT-RvNN can still have exhorbitant memory usage. Furthermore, so far it has been only tested as a sentence encoder. We take a step towards addressing these issues in this paper:
1. We identify a critical memory bottleneck in both Gumbel-Tree RvNN and Beam-Tree RvNN and propose strategies to fix this bottleneck (see §3). Our strategies reduce the peak memory usage of Beam-Tree RvNNs by 10-16 times in certain stress tests (see Table 4).
2. We propose a strategy to utilize the intermediate tree nodes (the span representations) to provide top-down signals to the original terminal representations using a parent attention mechanism. This allows a way to contextualize token representations using RvNNs enabling us to go beyond sentence-encoding (see §4).
3. We show that the proposed efficient variant of BT-RvNN incurs marginal accuracy loss if at all compared to the original—and in some cases even outperforms the original by a large margin (in ListOps).
2 Existing Framework
Here we first describe the existing framework used in Choi et al. [9] and Ray Chowdhury and Caragea [66]. In the next section, we identify its weakness in terms of memory consumption and resolve it.
Task Structure: As in related prior works [9, 10, 71], we start with our focus (although we will expand - see §4) on exploring the use of RvNNs as sentence encoders of the form:  $f: \mathbb{R}^{n \times d} \to \mathbb{R}^d$ . Given a sequence of  $n$  vectors of size  $d$  as input (of the form  $\mathbb{R}^{n \times d}$ ),  $f$  compresses it into a single vector (of the form  $\mathbb{R}^d$ ). Sentence encoders can be used for sentence-pair comparison or classification.
Components: The core components of this framework are: a scoring function  $scorer : \mathbb{R}^d \to \mathbb{R}$  and a recursive cell  $rec : \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}^d$ . The scorer is implemented as  $scorer(v) = W_v v$  where  $W_v \in \mathbb{R}^{1 \times d}$ . The  $rec(child_l, child_r)$  function is implemented as 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 {1}
$$
$$
p = L N (\sigma (l) \odot \text {c h i l d} _ {l} + \sigma (r) \odot \text {c h i l d} _ {r} + \sigma (g) \odot h) \tag {2}
$$
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$  . Overall, this is the Gated Recursive Cell (GRC) that was originally introduced by Shen et al. [71] and has been consistently shown to be superior [71, 10, 66] to earlier RNN cells like Long Short Term Memory Networks (LSTM) [33, 80].
Note that these models generally also apply an initial transformation layer to the terminal nodes before starting up the RvNN. Similar to [71, 10, 66], we apply a single linear transform followed by a layer normalization as the initial transformation.
Greedy Search Tree Recursive Neural Networks: Here we describe the implementation of Gumbel-Tree RvNN [9]. Assume that we have a sequence of hidden states of the form  $(h_1^t, h_2^t, \ldots, h_n^t)$  in some intermediate recursion layer  $t$ . For the recursive step in that layer, first all possible parent node candidates are computed as:
$$
p _ {1} ^ {t} = \operatorname {r e c} \left(h _ {1} ^ {t}, h _ {2} ^ {t}\right), p _ {2} ^ {t} = \operatorname {r e c} \left(h _ {2} ^ {t}, h _ {3} ^ {t}\right), \dots , p _ {n - 1} ^ {t} = \operatorname {r e c} \left(h _ {n - 1} ^ {t}, h _ {n} ^ {t}\right) \tag {3}
$$
Second, each parent candidate node  $p_i^t$  is scored as  $e_i^t = scorer(p_i^t)$ . Next, the index of the top score is selected as  $j = \arg\max(e_{1:n-1}^t)$ . Finally, now the update rule for the next recursion can be expressed as:
$$
h _ {i} ^ {t + 1} = \left\{ \begin{array}{l l} h _ {i} ^ {t} & i <   j \\ \operatorname {r e c} \left(h _ {i} ^ {t}, h _ {i + 1} ^ {t}\right) & i = j \\ h _ {i + 1} ^ {t} & i > j \end{array} \right. \tag {4}
$$
Note that this is essentially a greedy tree search process where in each turn all the locally available choices (parent candidates) are scored and the maximum scoring choice is selected. In each iteration, the sequence size is decreased by 1. In the final step only one representation will be remaining (the root node). At this point, however, the tree parsing procedure is not differentiable because it purely relies on an argmax. In practice, reinforcement learning [31], STE with gumbel softmax [9], or techniques like SPIGOT [61] have been used to replace argmax. Below we discuss another alternative and our main focus.
Beam Search Tree Recursive Neural Networks (BT-RvNN): Seeing the above algorithm as a greedy process provides a natural extension through beam search as done in Ray Chowdhury and Caragea [66]. BT-RvNN replaces the argmax in Gumbel-Tree RvNN with a stochastic Top- $K$  [42] that stochastically extracts both the  $K$  highest scoring parent candidates and the  $K$  corresponding log-softmax scores. The process collects all parent node compositions and also accumulates (by addition) log-softmax scores for each selected choices in corresponding beams. With this, the end result is  $K$  beams of accumulated scores and  $K$  beams of root node representations. The final representation is a softmax-based marginalization of the  $K$  root nodes:  $\sum_{i} \frac{\exp(s_i)}{\sum_{j} \exp(s_j)} \cdot b_i$  where  $b_i$  is the root node representation of beam  $i$  and  $s_i$  is the accumulated (added) log-softmax scores at every iteration for beam  $i$ . Doing this enabled BT-RvNN to improve greatly [66] over greedy Gumbel-Tree recursion [9]. However, BT-RvNN can also substantially increase the memory usage, which makes it inconvenient to use.
3 Bottleneck and Solution
In this section, we identify a major bottleneck in the memory usage that exists in the above framework (both for greedy-search and beam-search) that can be adjusted for.
Figure 1: Visualization of the contrast between the existing framework (left) and the proposed one (right).  $H_{1}, H_{2}, H_{3}$  are the input representations in the iteration. The possible contiguous pairs of them are candidate child pairs for nodes to be built in this iteration. On the left side, we see each pair is in parallel fed to the recursive cells to create their corresponding candidate parent representations. Then they are scored and one parent  $(P_{1})$  is selected. On the right side (our approach), each child pair candidate is directly scored. The faded colored bars in  $H_{1}, H_{2}, H_{3}$  represent sliced away vector values. The scoring function then selects one child pair. Then only that specific selected child pair is composed using the recursive cell to create the parent representation  $(P_{1})$  not wasting unnecessary compute by applying the recursive cell for other non-selected child pairs.
Bottleneck: The main bottleneck in the above existing framework is Eqn. 3. That is, the framework runs a rather heavy recursive cell (concretely, the GRC function in Eqn. 2) in parallel for every item in the sequence and for every iteration. In contrast, RNNs could use the same GRC function but for one position at a time sequentially - taking very little memory. While Transformers also use similarly big feedforward networks like GRC in parallel for all hidden states - they have fixed a number of layers - whereas BT-RvNNs may have to recurse for hundreds of layers depending on the input making this bottleneck more worrisome. However, we think this bottleneck can be highly mitigated. Below, we present our proposals to fix this bottleneck.
3.1 Efficient Beam Tree Recursive Neural Network (EBT-RvNN)
Here, we describe our new model EBT-RvNN. It extends BT-RvNN by incorporating the fixes below. We present the contrast between the previous method and our current method visually in Figure 1.
Fix 1 (new scorer): At any iteration  $t$ , we start only with some sequence  $(h_1^t, \ldots, h_n^t)$ . In the existing framework, starting from this the function to compute the score for any child-pair will look like:
$$
e _ {i} = \operatorname {s c o r e r} \circ \operatorname {r e c} \left(h _ {i}, h _ {i + 1}\right) \tag {5}
$$
This is, in fact, the only reason for which we need to apply rec to all positions (in Eqn. 3) at this iteration because that is the only way to get the corresponding score; that is, currently the score computation entangles the recursive cell rec and the scorer. However, there is no clear reason to do this. Instead we can just replace scorer  $\circ$  rec (in Eqn. 5) with a single new scorer function  $(s c o r e r_{n e w} : \mathbb{R}^{2 \times d} \to \mathbb{R})$  that directly interacts with the concatenation of  $(h_i, h_{i+1})$  without the rec as an intermediate step - and thus disentangling it from the scorer. We use a parameter-light 2-layered MLP to replace scorer  $\circ$  rec:
$$
e _ {i} = \operatorname {s c o r e r} _ {\text {n e w}} \left(h _ {i}, h _ {i + 1}\right) = \operatorname {G e L U} \left(\left[ h _ {i}; h _ {i + 1} \right] W _ {1} ^ {s} + b _ {1} ^ {s}\right) W _ {2} ^ {s} + b _ {2} ^ {s} \tag {6}
$$
Here,  $W_1^s \in \mathbb{R}^{2 \cdot d \times d_s}$ ,  $W_1^s 2 \in \mathbb{R}^{d_s \times 1}$ ,  $b_1^s \in \mathbb{R}^{d_s}$ ,  $b_2^s \in \mathbb{R}$ . Since lightness of the scorer_new function is critical for lower memory usage (this has to be run in parallel for all contiguous pairs) we set  $d_s$  to be small ( $d_s = 64$ ). In this formulation, the rec function will be called only for the selected contiguous pairs (siblings) in Eqn. 4.
Fix 2 (slicing): While we already took a major step above in making BT-RvNN more efficient, we can still go further. It is unclear whether the full hidden state vector size  $d$  is necessary for parsing decisions. Parsing typically hangs on more coarse-grained abstract information - for example, when doing arithmetic while precise numerical information needs to be stored in the hidden states for future
computation, the exact numerical information is not relevant for parsing - only the class of being a number should suffice. Thus, we assume that we can project the inputs into a low-dimensional space for scoring. One way to do that is to use a linear layer. However, parallel matrix multiplications on the full hidden state can be still costly when done for each hidden state in the sequence in every recursion. So, instead, we can allow the initial transformation or the RvNN itself to implicitly learn to organize parsing-related information in some sub-region of the vector. We can treat only the first  $\min(d_s, d)$  (out of  $d$ ) as relevant for parsing decisions. Then we can simply slice the first  $\min(d_s, d)$  out before sending the candidate child pairs to the scoring function. Thus, the score computation can be presented as below:
$$
e _ {i} = \operatorname {s c o r e r} _ {\text {n e w}} \left(h _ {i} [ 0: \min  \left(d _ {s}, d\right) ], h _ {i + 1} [ 0: \min  \left(d _ {s}, d\right) ]\right) \tag {7}
$$
So, now  $W_1^s \in \mathbb{R}^{2 \cdot \min(d_s, d)} \times d_s$ . As before we keep  $d_s$  small ( $d_s = 64$ ). Now, the total hidden state size ( $d$ ) can be kept as high as needed to preserve overall representation capacity without making the computation scale as much with increasing  $d$ . The model can now just learn through gradient descent to organize parsing relevant features in the first  $\min(d_s, d)$  values of the hidden states because only through them will gradient signals related to parsing scores propagate. Note that unlike [31], we are not running a different rec function for the parsing decisions. The parsing decision in our case still depends on the output of the same single recursive cell but from the previous iterations (if any).