Efficient Beam Tree Recursion

Beam Tree Recursive Neural Network (BT-RvNN) was recently proposed as a simple 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 not the worst in its kind, BT-RvNN can be still exorbitantly expensive in memory usage. 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} \rightarrow \mathbb{R}^{d}$ into a sequence contextualizer of the form $f:\mathbb{R}^{n \times d} \rightarrow \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.

Paper

Similar papers

Peer review

Reviewer S18L5/10 · confidence 3/52023-07-06

Summary

This paper relates to tree-recursive neural networks, which are a type of network where network modules are wired up in a tree structure that is itself determined by outputs of earlier parts of the network. This recursion poses significant challenges, including speed, memory usage, and potentially other issues (such as the need to use reinforcement learning to train the selection of tree structure). This paper proposes a more lightweight architecture for the portion of the network that determines the topology of the tree, in the context of a family of approaches that have good behavior in terms of gradient estimation. As a result, the method achieves both good memory usage and time complexity. Additionally, the paper proposes the use of tree-recursive networks to produce representations for each token, instead of a single representation that summarizes the entire input.

Strengths

The proposed method performs well in terms of speed, memory usage, and out-of-distribution generalization on ListOps -- it seems to be the overall winner compared to the baseline recursive NN methods.

Weaknesses

One concern regarding Tree-RvNNs is that the efficiency limitations set put this entire class of methods outside of practical usefulness for many tasks. The paper is notably light in comparisons to other families of approaches, such as the standard Transformer. This weakness is most pronounced in the part of the paper dealing with token representations and evaluating on SNLI/QQP/MNLI, which are standard tasks to run Transformers on (unlike ListOps, where the applicability of Tree-RvNNs is readily apparent). The part of the paper that deals with token-level encoding (Section 4) seems underdeveloped, at least compared to the discussion of Tree-RvNNs in the sentence-encoding scenarios. As a result the paper gives a bit of an impression of two distinct ideas ideas being glued together, as opposed to a more coherent direction for the work.

Questions

How does the proposed approach compare to more standard architectures like the Transformer?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

No concerns regarding discussion of limitations.

Reviewer XtZN7/10 · confidence 3/52023-07-09

Summary

This paper makes 2 contributions. The major one is a more efficient Beam Tree Recursive Neural Network (BT-RvNN) that learns a binary tree structure over a sequence of tokens. The second is a cross attention mechanism to incorporate this tree structure into transformers. They increase the efficiency by decreasing the hidden dimension and increasing the parallelizability and remove OneSoft. They verify that they do indeed see memory usage drop and can apply their method to longer sequences. They achieve minimal or no loss in quality with their new efficient method when using it as a sentence encoder. Their model seems to work especially well in the sentence interaction context when they preserve the token representation.

Strengths

Good exposition and necessary background provided. Efficiency results are very strong and the method clearly achieves a convincing win with memory usage. Experiments are well done a reproducible.

Weaknesses

The part about sentence interactions detracted from the clarity of the paper since I think the main contribution is the efficiency wins. The method on occasion does perform worse?

Questions

I don't complete understand section 4. Is the model jointly trained with the transformer or is the model trained separately and frozen and the tree structure is what is consumed? When using EBT-GAU vs EGT-GAU in the sentence interaction models, the difference is in the hidden states h_{i]^{t} of the tree nodes?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Method is still slow.

Reviewer oTCB5/10 · confidence 3/52023-07-21

Summary

This paper presents a comprehensive analysis of Beam Tree Recursive Neural Networks (BT-RvNN), a recently proposed variant of recursive neural networks. The study identifies suboptimal neural network parameterization in BT-RvNNs, which adversely affects running efficiency and results in excessive memory consumption. By proposing several techniques to remove the memory bottleneck, the authors demonstrate that BT-RvNNs can be run with significantly reduced memory while still maintaining task performance. Additionally, the paper presents an attention-augmented approach to expand the score of Tree-RvNNs beyond sentence encoders to encompass sentence contextualization, enriching the potential downstream applications of Tree-RvNNs.

Strengths

- The paper is well structured and offers an extensive background of recursive neural networks. - The proposed generalization from sentence encoders to sentence contextualizers is compelling, showcasing promising empirical performance compared to plain attention modules without tree-RvNNs. This expansion broadens the potential applications of Tree-RvNNs and establishes them as building blocks with specific inductive biases for learning tasks.

Weaknesses

- My main concern lies in the clarity of motivation. If my understanding is correct, the strategies proposed to enhance efficiency seem to be primarily related to improving neural network parameterization, rather than improving the beam tree recursion algorithm, which may have been inferred during the reading of the introduction. As a result, the motivation could be perceived more as an engineering effort with limited technical novelty. - Although the proposed method demonstrates a significant reduction in memory usage, the trade-off between runtime and memory utilization becomes evident in both baselines (GT-GRC and BT-GRC). A more in-depth exploration of this trade-off and its implications would enhance the comprehensiveness.

Questions

1. Why does the simplified network parameterization outperform the original architecture? An in-depth examination of this (e.g., conducting ablation studies of proposed fixes) would contribute to a more comprehensive understanding of the underlying mechanisms. 2. The experimental findings reveal that EGT-GRC exhibits a considerable performance drop compared to the original GT-GRC across diverse tasks, whereas EBT-GRC outperforms BT-GRC on most tasks. Why is this the case? The paper would benefit from providing a thorough discussion of the factors influencing the contrasting performance trends observed in EGT-GRC and EBT-GRC. 3. Typo: L269: missing close parenthesis.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors have adequately addressed the limitations.

Reviewer x8Vp5/10 · confidence 4/52023-07-26

Summary

1. This paper identify a critical memory bottleneck in both Gumbel-Tree RvNN and Beam-Tree RvNN and propose effective strategies to address it. The approaches reduce the peak memory usage of Beam-Tree RvNNs by 10-16 times in certain stress tests (refer to Table 4). 2. This paper introduce a novel strategy that utilizes intermediate tree nodes (span representations) to provide top-down signals to the original terminal representations using a parent attention mechanism. 3. This paper demonstrate that the proposed efficient variant of BT-RvNN incurs minimal accuracy loss compared to the original and, in some cases, even outperforms the original by a significant margin (especially in ListOps).

Strengths

1. This paper identify a critical memory bottleneck in both Gumbel-Tree RvNN and Beam-Tree RvNN and propose effective strategies to address it. The approaches reduce the peak memory usage of Beam-Tree RvNNs by 10-16 times in certain stress tests (refer to Table 4). 2. This paper introduce a novel strategy that utilizes intermediate tree nodes (span representations) to provide top-down signals to the original terminal representations using a parent attention mechanism. 3. This paper demonstrate that the proposed efficient variant of BT-RvNN incurs minimal accuracy loss compared to the original and, in some cases, even outperforms the original by a significant margin (especially in ListOps).

Weaknesses

The paper is written very clearly, and the methods are simple yet effective.

Questions

1. If some baselines based on sequence-to-sequence methods are provided, the motivation for the Tree-based model will become clearer.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

Reviewer oTCB2023-08-16

I thank the authors for the feedback. The provided clarification and new experimental results make the context clearer to me and resolve some of my concerns (especially Weakness 2 and Questions 1 and 2). However, I am still not fully convinced by the motivation. While it is true that the memory efficiency is largely improved, it appears that the advancements primarily come from an improved neural network parameterization (it is not about parameter count but a more suitable parameter allocation). In this regard, the technical contribution seems limited, as it focuses more on designing a neural network that implements the algorithm with less memory, rather than improving the algorithm itself. After carefully reading through the clarification as well as the other reviews, I decided to stay at my initial rating of 5.

Authorsrebuttal2023-08-17

Clarifications

Thank you for the response. Based on the response we would like to be clear and explicit about a few points: 1. **On Parameterization:** You are right that the EBT-RvNN and EGT-RvNN are, in a sense, better parameterization (better allocation of compute) of BT-RvNN and GT-RvNN, respectively. However, contributions in similar spaces appear in NeurIPS [1]. 2. **On the clarity of the motivation:** The motivation for EBT-RvNN simply is memory reduction (as you noted can be inferred from the abstract) while preserving accuracy. Our specific form of parameterization is proposed and demonstrated as an effective means to achieve the motivated goal. 3. **On change in the algorithm (Part 1):** Whether the improvement (improvement in memory reduction is also an improvement) is made on the side of the algorithm or elsewhere depends on how abstract the algorithm we have in mind. This can be a subjective point. Precisely, it depends on whether we want to consider the use of next parent representations as input to the score function in Choi et al. as a mere "implementation detail" or as part of the specific algorithm used in Choi et al. We would admit, however, that we are not changing the more general idea of easy-first parsing algorithm. 4. **On change in the algorithm (Part 2):** If we fix the level of detail to be that in which Algorithm 1 (Greedy Search-based RvNN) and Algorithm 2 (Beam Search based RvNN) are presented in Appendix A of the “Beam Tree Recursive Cell” paper (from supplementary material) [2], then technically the parameterization does involve changes of several lines in the algorithms. For example, the line ```Parents← [cell(childL, childR) for childL, childR in zip(ChildrenL,ChildrenR]``` will be entirely removed, and there will be separate lines for selecting the chosen sibling pairs and apply the cell function to them alone besides changing the scorer. We will add an updated pseudocode for better contrast against Algorithm 1 and 2 as presented in “Beam Tree Recursive Cell”. 5. **Other contributions:** We also want to highlight that, even if the technical contribution of EBT-RvNN may by itself be considered limited, this is not the sole contribution. As you noted (in strengths), we also introduce the parent attention mechanism for creating token contextualized representations which further adds to our overall technical contribution (EBT-RvNN + parent attention). This is not just a change of parameterization but requires further non-trivial technical changes (Appendix 5.2) and additionally paves the way for broader applicability of EBT-RvNN (as you acknowledged). [1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - Dao et al. NeurIPS 2022 [2] Anonymized paper; published in ICML 2023

Reviewer S18L2023-08-19

Thank you for the response. The additional information in the response helps put the results in context. I still think that the lack of discussion of Transformer baselines in the paper itself can make it harder to contextualize for a reader who is not already familiar with the line of work on recursive networks and related architectures.

Authorsrebuttal2023-08-19

Thank you for the feedback. We will add the additional results for better contextualization.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC