Sparse Universal Transformer (2023)

Abstract: The Universal Transformer (UT) is a variant of the Transformer that shares parameters across its layers and is Turing-complete under certain assumptions. Empirical evidence also 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, most state-of-the-art NLP systems use VTs as their backbone model instead of UTs. This is mainly because scaling UT parameters is 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) to reduce UT’s computation complexity while retaining its parameter efficiency and generalization ability. Experiments show that SUT combines the best of both worlds, achieving strong generalization results on formal language tasks (Logical inference and CFQ) and impressive parameter and computation efficiency on standard natural language benchmarks like WMT’14.

## introduction recent theoretical work has pointed out that finitedepth transformers have an issue of expressibility that will result in failure to generalize @xcite @xcite . @xcite 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; @xcite are transformers that share parameters at every layer of the architecture. ## background & related work overcoming vt limitations with @xcite and @xcite find that vanilla transformers learn shortcuts for tasks that require multistep compositional operations, and fail to generalize on larger instances of the problem that require more steps. theoretical results have also shown that vanilla transformers have limitations in what they can compute that support these findings @xcite @xcite . universal transformers @xcite are transformers with tied weights across all layers, and an additional halting mechanism to decide when to stop. in an ideal scenario of infinite layers (now that all layers have the same parameters) the ut, like the neural gpu @xcite , is turing-complete, which overcomes many of the abovementioned issues. in practice, even with limited depth, uts have exhibited properties that afford them better performance in compositional generalization tasks @xcite . uts allow operations learned in the transformer during training to be depth-order invariant. if some operations during training are learned to be performed in a certain order, during test time, the ut could generalize to an unseen order of operations. challenges with scaling the ut despite these compositional abilities, performance tends to decrease on real-world tasks when using uts. al-bert @xcite improved parameter efficiency by sharing parameters across layers. this was motivated by an observation that transformers tend to learn to perform similar operations in the layers, and that sharing these parameters would reduce this redundancy foot_0 . however, the authors observe a dip in performance when sharing parameters, contrary to @xcite . could the issue be one of model capacity? experiments with albert show that scaling up al-bert can outperform the bert baseline, even on real-world tasks @xcite . @xcite also show that a shared-parameter transformer has better scaling properties in terms of parameter-to-performance, but poorer properties in terms of computation-to-performance, since parameter count causes the computation to increase. tay et al. (2022) scale up different sequence models, and remark on difficulties with scaling up uts, limiting the experiments they can perform on ut. @xcite outline several strategies of scaling up shared-parameter transformers to deal with these issues by using different parametersharing schemes. our experiments show that smoe techniques can be applied successfully to the ut to scale it up, achieving the ut's parameter efficiency while not incurring the same computation costs. we also perform experiments that support the compositional generalization claims of prior work, and provide better baselines for those tasks. ## method like ut, we reuse the same sut block for every layer of the transformer. within each sut block, we use smoes to achieve sparsity for the feedforward network (ffd) and attention heads separately. we use the mutual information maximization loss proposed in @xcite and modified for unsupervised tasks in @xcite . finally, we propose a stick-breaking process formulation of dynamic halting, which affects how the attention mechanism works in the sut, and the adaptive computation time (act) auxiliary loss we use to minimize the number of layers used. we replace the feed-forward layer (ffd) in the transformer block with a mixture of ffds. each mixture of ffd can be described with a 3-tuple, (e, k, d): e experts, k for the number of experts to be used in the top-k operation, and d for the dimension of the hidden layer for each ffd expert. for the attention layer, we use the mixture of multihead attention (momha) proposed by @xcite and @xcite . each momha module can be described by a 5-tuple, (e, k, h, d, w ), with e representing the number of experts, k representing the parameter k in a top-k operation, h representing the number of attention heads per expert, and d for the dimensions per head. like in moa, momha maintains only a single set of h key-value projections shared among the experts, while there are e query projections of h heads each. w represents the relative position embedding window size, parameterizing 2w + 1 embeddings for w positions before and after the present position. figure 3 (left) shows the schematic of a sut block. ## experiments first, we show that we can scale the ut with sut on the wmt'14 english-german @xcite translation task. we then ran experiments on compositional freebase questions (cfq; @xcite to test for compositional generalization model valid loss bleu #params sut base 2.192 29.2 66m -mim loss 2.221 28.9 66m -momha 2.232 28.7 66m -act loss 2.217 29.0 66m -halting 2.219 29.1 65m properties. to further analyze the behaviour of the model under compositional generalization settings, we test our model on the logical inference task from (bowman et al., 2015). all experiments were implemented within the fairseq framework (ott et al., 2019) foot_1 . ## conclusion we show that it is possible to scale up the ut via suts, and suts outperforms models of the same capacity in the wmt'14 english-to-german translation task. the recursive nature of both uts and suts allows for better inductive biases, which we have demonstrated in synthetic tasks like cfq and logical inference. vts have been shown to be poor at these compositional generalization tasks without additional domain knowledge. the stick-breaking dynamic halting mechanism also allows post-training adjustment of computation cost, which is a boon for deployment at scale. limitations while the experiments in this paper show the desirable generalization properties of uts, there are some aspects of compositional generalization that suts do not solve. importantly, while we demonstrate scaling uts up via smoes, further experiments on larger settings are needed to ascertain viability in large scale systems. other issues may also crop up in the further scaling of suts, but we believe there is ample literature to draw on to finding solutions for these problems.