  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.
