Sparse Universal Transformer
Shawn Tan Yikang Shen Zhenfang Chen Aaron Courville Chuang Gan
The Universal Transformer (UT) is a variant of the Transformer that shares
parameters across its layers. Empirical evidence shows that UTs have better
compositional generalization than Vanilla Transformers (VTs) in formal language
tasks. The parameter-sharing also affords it better parameter efficiency than
VTs. Despite its many advantages, scaling UT parameters is much more compute
and memory intensive than scaling up a VT. This paper proposes the Sparse
Universal Transformer (SUT), which leverages Sparse Mixture of Experts (SMoE)
and a new stick-breaking-based dynamic halting mechanism to reduce UT's
computation complexity while retaining its parameter efficiency and
generalization ability. Experiments show that SUT achieves the same performance
as strong baseline models while only using half computation and parameters on
WMT'14 and strong generalization results on formal language tasks (Logical
inference and CFQ). The new halting mechanism also enables around 50\%
reduction in computation during inference with very little performance decrease
on formal language tasks.