[论文解读] Compositional Generalization with Tree Stack Memory Units

Forough Arabshahi, Zhichu Lu+3 more|arXiv (Cornell University)|Nov 5, 2019

Topic Modeling参考文献 35被引用 5

一句话总结

本文提出树形堆栈记忆单元（Tree-SMU），一种递归神经网络架构，通过在每个节点中集成可微分堆栈记忆单元，提升组合泛化能力。通过支持软压入和弹出操作，结合子节点状态与堆栈信息，Tree-SMU 能够捕捉长距离依赖关系并保持结构顺序，相较于 Transformer 和 Tree-LSTM，在数学推理基准上实现了最先进水平的零样本泛化表现。

ABSTRACT

We study compositional generalization, viz., the problem of zero-shot generalization to novel compositions of concepts in a domain. Standard neural networks fail to a large extent on compositional learning. We propose Tree Stack Memory Units (Tree-SMU) to enable strong compositional generalization. Tree-SMU is a recursive neural network with Stack Memory Units (\SMU s), a novel memory augmented neural network whose memory has a differentiable stack structure. Each SMU in the tree architecture learns to read from its stack and to write to it by combining the stacks and states of its children through gating. The stack helps capture long-range dependencies in the problem domain, thereby enabling compositional generalization. Additionally, the stack also preserves the ordering of each node's descendants, thereby retaining locality on the tree. We demonstrate strong empirical results on two mathematical reasoning benchmarks. We use four compositionality tests to assess the generalization performance of Tree-SMU and show that it enables accurate compositional generalization compared to strong baselines such as Transformers and Tree-LSTMs.

研究动机与目标

• 为解决标准神经网络在零样本组合泛化中的失败问题，尤其是在分布偏移情况下的表现。
• 克服递归神经网络（如 Tree-LSTM）中误差传播和长距离依赖捕捉能力有限的问题。
• 设计一种内存增强的递归架构，支持对新组合的系统性泛化。
• 在多个组合性测试中评估泛化性能，包括深度外推和语义等价性测试。
• 证明内部记忆集成可同时提升零样本泛化能力和样本效率。

提出的方法

• Tree-SMU 是一种递归神经网络，其中每个节点均包含一个具有可微分堆栈结构的堆栈记忆单元（SMU）。
• 每个 SMU 通过门控机制控制的软压入和弹出门，学习从堆栈读取和写入信息。
• 父节点的堆栈通过门控机制组合其子节点的堆栈和隐藏状态形成。
• 软压入和弹出操作使模型能够存储和检索来自后代的信息，实现对非局部依赖关系的间接访问。
• 堆栈结构保留了后代的顺序，维持局部性并支持层次化组合。
• 该架构将记忆集成在单元内部，避免了先前工作中外部记忆模块失效的问题。

实验结果

研究问题

• RQ1 具有内部堆栈记忆的递归神经网络是否能在新组合上实现优于标准架构的零样本泛化？
• RQ2 堆栈记忆机制是否能提升对训练中未见的更深或更浅组合的泛化能力？
• RQ3 模型能否为结构不同但功能相同的表达式学习到语义等价的表示？
• RQ4 堆栈记忆的集成如何影响递归网络中的样本效率和错误纠正能力？
• RQ5 所提出的架构是否在组合推理任务上优于强基线模型（如 Transformer 和 Tree-LSTM）？

主要发现

• 在局部性测试中，Tree-SMU 在测试集上达到 98.86% 的准确率，显著优于 Tree-LSTM（91.58%）和 Tree-RNN（85.36%）。
• 在组合性测试中，Tree-SMU 达到 83.29% 的准确率，超过 Tree-LSTM（83.06%）和 Tree-RNN（83.06%），且对分布偏移更具鲁棒性。
• 在深度外推测试中，Tree-SMU 达到 79.57% 的准确率，显著优于 Tree-LSTM（77.58%）和 Tree-RNN（75.00%）。
• Tree-SMU 在全部四项组合性测试中表现一致，展现出强大的零样本泛化能力。
• 该模型表现出改进的样本效率，在训练集（深度 5–13）上达到 95.04% 的训练准确率，表明其能从有限数据中更高效学习。
• 消融实验确认，将堆栈增强 RNN 简单扩展至树形网络结构会失败，验证了集成 SMU 设计的必要性。