SlowGuess/NeurIPS_2023_OCR huggingface dataset:
RvNNs can also be simulated by stack-augmented recurrent neural networks (RNNs) to an extent (similar to how pushdown automata can model context-free grammar [69, 40]). There are multiple works on stack-augmented RNNs [8, 93, 50]. Ordered Memory [71] is one of the more modern such examples. More recently, DuSell and Chiang [17, 18] explored non-deterministic stack augmented RNNs and Del'etang et al. [14] explored other expressive models. Wu [90] presented a survey of latent structure models.
Choi et al. [9] proposed a greedy search strategy based on easy-first algorithm [23, 48] for auto-parsing structures for recursion utilizing STE gumbel softmax for gradient signals. Peng et al. [61] extended the framework with SPIGOT and Havrylov et al. [31] extended it with reinforcement learning (RL). Ray Chowdhury and Caragea [66] extended it with beam search and soft top-k. Chowdhury and Caragea [10], Zhang et al. [95] introduced different forms of soft-recursion.
Top-down Signal: Similar to us, Teng and Zhang [82] explored bidirectional signal propagation (bottom-up and top-down). However, they sent top-down signal in a sequential manner which can be expensive - either it can get slow without parallelization or memory-wise expensive with parallelization of contextualization of nodes in the same height. Our approach in EBT-GAU also has some kinship with BP-Transformer [92]. BP-Transformer allows message passing between a fixed subset of parent nodes and terminal nodes created using a heuristics-based balanced binary tree. Chart-based models can also create sequence contextualized representations [15, 16] but they can be quite expensive by default [66] needing their own separate techniques [34, 35].
Algorithm 1 Efficient Beam Tree Cell (without slicing)
Input: data  $X = [x_{1},x_{2},\dots,x_{n}],k$  (beam size)   
Beam  $X\gets [X]$    
BeamScores  $\leftarrow [0]$    
while True do   
if len(BeamX[0])  $= = 1$  then   
Beam  $X\gets$  [beam[0] for beam in BeamX] break   
end if   
if len(BeamX[0])  $= = 2$  then   
Beam  $X\gets$  [cell(beam[0],beam[1]) for beam in BeamX] break   
end if   
NewBeam  $X\gets []$    
NewBeamScores  $\leftarrow []$    
for Beam,BeamScore in zip(BeamX, BeamScores) do Scores  $\leftarrow$  log o softmax([scorer(beam[i],beam[i+1]) for parent in Parents]) Indices  $\leftarrow$  topk(Scores,k)   
for i in range(K) do newBeam  $\leftarrow$  deepcopy(Beam)   
newBeam[Indices[i]]  $\leftarrow$  cell(Beam[Indices[i]], Beam[Indices[i] + 1]) Delete newBeam[Indices[i] + 1] NewBeamX.append(newBeam)   
newScore  $\leftarrow$  BeamScore + Scores[indices[i]]   
newBeamScores.append(newScore)   
end for   
end for   
Indices  $\leftarrow$  topk(newBeamScores,k)   
BeamScores  $\leftarrow$  [newBeamScores[i] for i in Indices]   
Beam  $X\gets$  [newBeamX[i] for i in Indices]   
end while   
BeamScores  $\leftarrow$  Softmax(BeamScores)   
Return sum([score * X for score, X in zip(BeanScores, BeamX)])
Transformers + RvNNs: There have been several approaches to incorporating RvNN-like inductive biases to Transformers. For instance, Universal Transformer [13] introduced weight-sharing and dynamic halt to Transformers. Csordás et al. [12] extended on universal transformer with geometric attention for locality bias and gating. Shen et al. [74] built on weight-shared transformers with high layer depth and group self-attention. Wang et al. [87], Nguyen et al. [58], Shen et al. [73] added hierarchical structural biases to self-attention. Fei et al. [20] biased pre-trained Transformers to have constituent information in intermediate representations. Hu et al. [34] used Transformer as binary recursive cells in chart-based encoders.
E Pseudocode
We present the pseudocode of EBT-RvNN in Algorithm 1. Note that the algorithms are written as they are for the sake of illustration: in practice, many of the nested loops are made parallel through batched operations.
F Architecture details
F.1 Sentence Encoder Models
For the sentence encoder models the architectural framework we use is the same siamese dual-encoder setup as Ray Chowdhury and Caragea [66].
F.2 Sentence Interaction Models
GAU-Block: Our specific implementation of a GAU-block [36] is detailed below. Our GAU-Block can be defined as GAUBlock  $(x,p,G)$ . The function arguments are of the following forms:  $x \in \mathbb{R}^{n \times d}$ ,  $p \in \mathbb{R}^{l \times d}$  and  $G \in \{0,1\}^{n \times l}$ .  $x$  accepts the main sequence of vectors that is to serve as attention queries;  $p$  accepts either the sequence of intermediate node representations created from our RvNN (for parent attention) or it accepts the same input as  $x$  (for usual cases);  $p$  serves as keys and values for attention;  $G$  accepts either the adjacency matrix in case of parent attention (where  $G_{ij} = 1$  iff  $p_j$  is a parent of  $x_i$  else  $G_{ij} = 0$ ), otherwise, it accepts just the usual attention mask; either way,  $G$  serves as an attention mask.
$$
x ^ {\prime} = L N \left(x W _ {i n i t} + b _ {i n i t}\right); p ^ {\prime} = L N \left(p W _ {i n i t} + b _ {i n i t}\right) \tag {8}
$$
$$
u = \operatorname {S i L U} \left(x ^ {\prime} W _ {u} + b _ {u}\right); v = \operatorname {S i L U} \left(p ^ {\prime} W _ {v} + b _ {v}\right) \tag {9}
$$
$$
q = z _ {q} \odot \operatorname {S i L U} \left(x ^ {\prime} W _ {z} + b _ {z}\right) + z b _ {q}; k = z _ {k} \odot \operatorname {S i L U} \left(p ^ {\prime} W _ {z} + b _ {z}\right) + z b _ {k} \tag {10}
$$
$$
A = \operatorname {S o f t m a x} \left(\frac {q k ^ {T} + p o s}{\sqrt {2 d}}, \text {m a s k} = G\right) \tag {11}
$$
$$
v ^ {\prime} = A v \tag {12}
$$
$$
o = \left(u \odot v ^ {\prime}\right) W _ {o} + b _ {o} \tag {13}
$$
$$
g = \text {S i g m o i d} ([ o; x ] W _ {\text {g a t e}} + b _ {\text {g a t e}}) \tag {14}
$$
$$
o u t = g \odot o + (1 - g) \odot x \tag {15}
$$
Here,  $W_{init} \in \mathbb{R}^{d \times d}$ ;  $W_z \in \mathbb{R}^{d \times d_h}$ ,  $W_u, W_v \in \mathbb{R}^{d \times 2d}$ ,  $b_{init}, b_z, b_o \in \mathbb{R}^d$ ;  $z_q, zb_q, z_k, zb_k \in \mathbb{R}^{d_h}$ ;  $b_u, b_v \in \mathbb{R}^{2d}$ ,  $W_o, W_{gate} \in \mathbb{R}^{2d \times d}$ . [;] represents concatenation.
$LN$  is layer normalization.  $pos$  is calculated using the technique of Raffel et al. [64] using relative tree height distance for parent attention, or relative positional distance for usual cases.
GAU Sequence Interaction Setup: Let GAUStack represent some arbitrary number of compositions of GAUBlocks (multilayered GAU block). GAUStack has the same function arguments as GAUBlock. Given two sequences  $(x_{1}, x_{2})$  and their corresponding attention masks  $(M_{1}, M_{2})$  as inputs where  $x_{1} \in \mathbb{R}^{n_{1} \times d}$ ,  $x_{2} \in \mathbb{R}^{n_{2} \times d}$ ,  $M_{1} \in \{0, 1\}^{n_{1} \times n_{1}}$ ,  $M_{1} \in \{0, 1\}^{n_{2} \times n_{2}}$ , the GAU setup can be expressed as:
$$
i n p = \left[ C L S + s e g _ {1}; x _ {1} + s e g _ {1}; S E P; C L S + s e g _ {2}, x _ {2} + s e g _ {2} \right] \tag {16}
$$
$$
r = \operatorname {G A U S t a c k} (x = i n p, p = i n p, G = f \left(M _ {1}; M _ {2}\right)) \tag {17}
$$
$$
\alpha = \operatorname {S o f t m a x} \left(\operatorname {G E L U} \left(r W _ {1} + b _ {1}\right) W _ {2} + b _ {2}\right) \tag {18}
$$
$$
c l s ^ {\prime} = \sum_ {i} \alpha_ {i} r \tag {19}
$$
$$
l o g i t s = \operatorname {G E L U} \left(c l s ^ {\prime} W _ {1} ^ {\text {l o g i t s}} + b _ {1} ^ {\text {l o g i t s}}\right) W _ {2} ^ {\text {l o g i t s}} + b _ {2} ^ {\text {l o g i t s}} \tag {20}
$$
Here,  $CLS,SEP,seg_{1},seg_{2}\in \mathbb{R}^{1\times d}$  are randomly initialized trainable vectors;  $seg_{1},seg_{2}$  are segment embeddings.  $W_{1}\in \mathbb{R}^{d\times d}$ ,  $W_{2}\in \mathbb{R}^{d\times 1}$ ;  $b_{1},b_{2},b_{1}^{\mathrm{logits}}\in \mathbb{R}^{d}$ ;  $b_{2}^{\mathrm{logits}}\in \mathbb{R}^{c}$ ;  $W_{1}^{\mathrm{logits}}\in \mathbb{R}^{d\times d}$ ,  $W_{2}^{\mathrm{logits}}\in \mathbb{R}^{d\times c}$ .  $c$  is the number of classes for the task.  $f$  is a function that takes the attention masks as input and concatenates them while adjusting for the special tokens (CLS, SEP).
EGT-GAU Sequence Interaction Setup: EGT-GAU starts from the same input as above. Let us also assume we have the EGT-GRC  $(x)$  module which takes a sequence of vectors  $x\in \mathbb{R}^{n\times d}$  as the input to recursively process and outputs  $(cls,p,G)$  where  $cls\in \mathbb{R}^{1\times d}$  is the root representation,  $p\in \mathbb{R}^{l\times d}$  is the sequence of non-terminal representations from the tree, and  $G\in \{0,1\}^{n\times l}$  is the
adjacency matrix for parent attention (i.e.,  $G_{ij} = 1$  iff  $p_j$  is a parent of  $x_i$ , else  $G_{ij} = 0$ ). Technically, tree height information is also extracted for relative position but we do not express that explicitly for the sake of brevity. With these elements, EGT-GAU can be expressed as below:
$$
\operatorname {c l s} _ {1}, p _ {1}, G _ {1} = \operatorname {E G T - G R C} \left(x = x _ {1}\right); \operatorname {c l s} _ {2}, p _ {2}, G _ {2} = \operatorname {E G T - G R C} \left(x = x _ {2}\right) \tag {21}
$$
$$
x _ {1} ^ {\prime} = \operatorname {G A U S t a c k} _ {1} (x = x _ {1}, p = p _ {1}, G = G _ {1}); \quad x _ {2} ^ {\prime} = \operatorname {G A U S t a c k} _ {1} (x = x _ {2}, p = p _ {2}, G = G _ {2}) \tag {22}
$$
$$
c l s _ {1} ^ {\prime} = \operatorname {G E L U} \left(c l s _ {1} W _ {1} ^ {c l s} + b _ {1} ^ {c l s}\right) W _ {2} ^ {c l s} + b _ {2} ^ {c l s}; c l s _ {2} ^ {\prime} = \operatorname {G E L U} \left(c l s _ {2} W _ {1} ^ {c l s} + b _ {1} ^ {c l s}\right) W _ {2} ^ {c l s} + b _ {2} ^ {c l s} \tag {23}
$$
$$
i n p = \left[ c l s _ {1} ^ {\prime} + s e g _ {1}; x _ {1} ^ {\prime} + s e g _ {1}; S E P; c l s _ {2} ^ {\prime} + s e g _ {2}, x _ {2} ^ {\prime} + s e g _ {2} \right] \tag {24}
$$
$$
r = \operatorname {G A U S t a c k} _ {2} (x = i n p, p = i n p, G = f \left(M _ {1}, M _ {2}\right)) \tag {25}
$$
Everything else after eqn. 25 is the same as eqn. 18 to 20.  $SEP,seg_{1},seg_{2}\in \mathbb{R}^{1\times d};seg_{1},seg_{2}$  are segment embeddings as before.  $W_{1}^{cls},W_{2}^{cls}\in \mathbb{R}^{d\times d};b_{1}^{cls},b_{2}^{cls}\in \mathbb{R}^{d}$
EBT-GAU Sequence Interaction Setup: This setup is similar to that of EGT-GAU but with a few changes. EBT-GAU uses EBT-GRC as a module instead of EGT-GRC. EBT-GAU returns outputs of the form  $(cls, bp, bG, s)$  where  $cls \in \mathbb{R}^{1 \times d}$  is the beam-score-weighted-averaged root representation,  $bp \in \mathbb{R}^{b \times l \times d}$  are the beams (beam size  $b$ ) of sequences of non-terminal representations from the tree,  $bG \in \{0, 1\}^{b \times n \times l}$  are the beams of adjacency matrices for parent attention, and  $s \in \mathbb{R}^b$  are the softmax-normalized beam scores. Let NGAUStack represent the same function as GAUStack but formalized for batched processing of multiple beams of sequences. With these elements, EBT-GAU can be expressed as:
$$
c l s _ {1}, b p _ {1}, b G _ {1}, s _ {1} = \operatorname {E B T - G R C} (x = x _ {1}); c l s _ {2}, b p _ {2}, b G _ {2}, s _ {2} = \operatorname {E B T - G R C} (x = x _ {2}) \tag {26}
$$
$$
b x _ {1} = \operatorname {r e p e a t} \left(x _ {1}, b\right); \quad b x _ {2} = \operatorname {r e p e a t} \left(x _ {2}, b\right) \tag {27}
$$
$$
b x _ {1} ^ {\prime} = \mathrm {N G A U S t a c k} _ {1} \left(b x _ {1}, b p _ {1}, b G _ {1}\right); \quad b x _ {2} ^ {\prime} = \mathrm {N G A U S t a c k} _ {1} \left(b x _ {2}, b p _ {2}, b G _ {2}\right) \tag {28}
$$
$$
x _ {1} ^ {\prime} = \sum_ {i} s [ i ] \cdot b x _ {1} ^ {\prime} [ i ]; \quad x _ {2} ^ {\prime} = \sum_ {i} s [ i ] \cdot b x _ {2} ^ {\prime} [ i ] \tag {29}
$$
Everything else after eqn. 29 is the same as the equations 23-25 followed by the equations 18 to 20. repeat  $(x,b)$  changes  $x\in \mathbb{R}^{n\times d}$  to  $bx\in \mathbb{R}^{b\times n\times d}$  by batching the same  $x$  for  $b$  times.
G Hyperparameter details
For sentence encoder models, we use the same hyperparameters as [66] (the preprint of the paper is available in the supplementary in anonymized form) for all the datasets. The only new hyperparameter for EBT-GRC is  $d_{s}$  which we set as 64; otherwise the hyperparameters are the same as that of BT-GRC or BT-GRC OS. We discuss the hyperparameters of the sequence interaction models next. For EBT-GAU/EGT-GAU, we used a two-layered weight-shared GAU-Blocks for NGAUStack $_1$ /GAUStack $_1$  and a three-layered weight-shared GAU-Blocks for GAUStack $_2$  (for parameter efficiency and regularization). GAU uses a five-layered GAU-Blocks (weights unshared) for GAUStack so that the parameters are similar to that of EBT-GAU or EGT-GAU. We use a dropout of 0.1 after the multiplication with  $W_{o}$  in each GAUBlock layer and a head size  $d_{h}$  of 128 (similar to Hua et al. [36]). For relative position, we set  $k = 5$  ( $k$  here corresponds the receptive field for relative attention in Shaw et al. [70]) for normal GAUBlocks and  $k = 10$  for parent attention (since parent attention is only applied to higher heights, we do not need to initialize weights for negative relative distances). Other hyperparameters are kept same as the sentence encoder models. The hyperparameters of MNLI, SNLI, and QQP are shared. Note that all the natural language tasks are trained with fixed 840B Glove Embeddings [62] as in Ray Chowdhury and Caragea [66]. All models were trained in a single Nvidia RTX A6000. The code is available in the supplementary.