Sparse Universal Transformer

TL;DR

The Sparse Universal Transformer (SUT) improves parameter efficiency and reduces computation in universal transformers by using sparse mixture of experts and a dynamic halting mechanism, maintaining strong performance and generalization.

Contribution

Introduces SUT, combining SMoE and a novel halting mechanism to reduce computation while preserving the benefits of Universal Transformers.

Findings

Achieves similar performance with half the parameters and computation.

Demonstrates strong generalization on formal language tasks.

Enables 50% inference computation reduction with minimal performance loss.

Abstract

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.