Introduction
Recent theoretical work has pointed out that finite-depth Transformers have an issue of expressibility that will result in failure to generalize (Hahn, 2020 ; Hao et al., 2022 ; Merrill et al., 2022 ; Liu et al., 2022 ) . Delétang et al. ( 2022 ) ran several neural architectures on a suite of different synthetic languages generated from different levels of the Chomsky hierarchy and empirically confirmed these results, showing that VTs have difficulty generalizing to Regular languages.
Universal Transformers (UTs; Dehghani et al. 2018 ) are Transformers that share parameters at every layer of the architecture. Csordás et al. ( 2021 ) performed several compositional generalization experiments on VTs and UTs along with absolute and relative position embeddings, and showed that UTs with relative positional embeddings performed better on these tasks.

However, the task of scaling UTs is challenging due to its computation complexity (Kaplan et al., 2020 ; Tay et al., 2022 ; Takase and Kiyono, 2021 ) .
Consider a VT with P parameters for each layer and L layers.
Evaluating such a VT has computation complexity associated with the model size LP .
A size-equivalent UT would have a UT block with LP parameters and computation complexity of approximately LP to run the block once.
To run such a UT for equivalent L layers would incur a complexity of L^{2}P .
This increased computation complexity directly translates to increased training and inference time.
According to Takase and Kiyono ( 2021 ) , UT requires two times the training time and far more GPU memory than VT in WMT English-German translation task.

Sparsely activated neural networks were introduced to reduce the computation complexity of large models.
These networks activate parts of the network conditioned on the input, computing only parts of the model, thereby disentangling the number of parameters from the computation complexity.
This method allows for drastically increasing the number of parameters without proportionally increasing the computation complexity. Shazeer et al. ( 2017 ) introduced Sparse Mixture of Experts (SMoE), using the top- k operator to allow for sparse computation of experts.
This allows for replacing the FeedForword (FFD) layer in the Transformer with an ensemble of E_{\text{ffd}} FFDs, but only k FFDs (where k<E ) would have to be evaluated, conditioned on the input. Zhang et al. ( 2022 ) then introduced the Mixture of Attention heads (MoA), which allows Transformers to replace its Multihead Attention (MHA) layer with an ensemble of E_{\text{att}} attention heads and only activates k heads condition on the input, further sparsifying the model.

This paper introduces the Sparse Universal Transformer (SUT), which applies the above sparse computation techniques to UT.
Additionally, we replace the per-position halting mechanism in UT with a new stick-breaking formulation that has a probabilistic interpretation, allowing us to introduce an Adaptive Computation Time (ACT; Graves 2016 ) penalty to minimize layer use.
It also provides an easy way to adjust the trade-off between the amount of computation and model performance during inference, further improving the efficiency of the SUT at inference time.

To demonstrate effective scaling, we perform experiments on WMT’14 English to German translation, showing that an SUT can achieve better performance for the same parameter count, while incurring less computation cost than an equivalent dense UT.
Since the UT setting is a specific case of SUT, we show on the Compositional Freebase Questions (CFQ; Keysers et al. 2019 ) tasks that UTs have better compositional generalization properties, improving upon CFQ results from Csordás et al. ( 2021 ) .
Using the Logical Inference task (Bowman et al., 2015 ) , we analyse the behaviour of our UT on length and compositional generalization.
Finally, we show that the halting mechanism can be used to achieve further efficiency during inference time, and study the trade-off between efficiency and performance.