Leveraging Relaxed Equilibrium by Lazy Transition for Sequence Modeling (2022)

Abstract: In sequence modeling, certain tokens are usually less ambiguous than others, and representations of these tokens require fewer refinements for disambiguation. However, given the nature of attention-based models like Transformer and UT (universal transformer), all tokens are equally processed towards depth. Inspired by the equilibrium phenomenon, we present a lazy transition, a mechanism to adjust the significance of iterative refinements for each token representation. Our lazy transition is deployed on top of UT to build LT (lazy transformer), where all tokens are processed unequally towards depth. Eventually, LT is encouraged to oscillate around a relaxed equilibrium. Our experiments show that LT outperforms baseline models on several tasks of machine translation, pre-training, Learning to Execute, and LAMBADA.

## introduction attention-based models like transformer @xcite underly two core concepts: layer and refinement. with layer-stacked structure, the model entirely relies on attention mechanisms @xcite to refine token representations (e.g.,vectors) layer-by-layer from context-informed representations by using residual connections @xcite . however, layer is not necessary. ut (universal transformer) @xcite and its variants @xcite ) iteratively run single-layer but wide transformer for sequence modeling that token representations are refined at each step by attending to the sequence representation of the previous step, showing higher performance than layer-stacked transformer does on nmt @xcite , pre-training @xcite , language modeling @xcite , and other tasks of varying complexity @xcite with the same number of parameters. the computational bound of recurrence is not the number of tokens in the sequence (like rnn) or the number of stacked layers in the model but is the maximum number of refinements made to token representations, e.g., a pre-defined maximum number of iteration steps. we follow this line but further ponder refinement. concretely, the model refines token representations iteratively that tokens are equally processed towards depth for disambiguation. essentially, different token representations are refined at a step equally and concurrently. thus, regardless of the ambiguous state, a step shows the same importance for all the tokens. however, certain tokens are less ambiguous than others in sequence modeling, especially in nmt. it raises two questions: 1) do deep representations have to be learned for lessambiguous tokens; 2) do tokens must consider the contributions or importance of a step equally? meanwhile, our work derives its motivation from the combination of ut and act (adaptive computation time) @xcite that can dynamically estimate the importance of iteration steps for understanding different sentences in the babi task @xcite . however, we consider the importance of every iteration step for all the tokens in the same sequence, sharing a similar motivation with depth-adaptive transformer @xcite . on the other hand, these methods learn a halting probability for each step without considering the correspondence and interdependence throughout the iteration and the model's convergence. we attempt a mechanism like the reset gate in gru @xcite that the model is allowed to dynamically forget some information based on the correspondence and interdependence, where in our case, this information is the further refinement at the current step for the corresponding token representation. in this way, we weaken the importance of an iteration step for less-ambiguous token representations to prevent over-refining but retain the significance for more-ambiguous token representations to avoid under-refining, refining token representations at the same step unequally. to this end, we present a lazy transition that performs between two consecutive steps when running ut iteratively. it forms iterative refinements @xcite for token representations with a residual structure: our contributions are: 1) we present a lazy transition that is placed on top of ut to build lt (lazy transformer). our lazy transition dynamically forms a refinement path for a token at each step by pondering the step importance and the local equilibrium. eventually, lt is encouraged to oscillate around a relaxed equilibrium in sequence modeling; 2) we provide an empirical study to quantitatively analyze how the relaxed equilibrium impacts nmt in performance, deep models, and zero-shot inferring; 3) we show our model can consistently improve the performance of pre-training and two tasks of varying complexity, where standard transformer fails; 4) our empirical study shows that stable and smooth refinements at the early iteration steps are significant, which results in a strong equilibrium and a stable oscillation. adaptive computation while a token-wise recurrent model like ut can theoretically run infinitely, the model is commonly trained by setting a maximum number of iteration steps. however, the number of inference steps is flexible. @xcite first introduce act (adaptive computation time ) to compute a scalar halting probability predicted by the model at each step for each token in rnn. then, @xcite adapt act for ut to halt the iteration before reaching the maximum number of training steps 1 in inferring. furthermore, @xcite present depthadaptive decoding for tokens to model the distribution of exiting with the probability of computing each layer/step and then emitting token predictions. our lazy transition follows this line somewhat that (1 -σ t i i ) in eq.1 dynamically approximates the computation time for tokens. it reflects on two characteristics: 1) (1 -σ ## approach we present lt (lazy transformer), a token-wise recurrent model. lt only consists of a ut block and our lazy transition. in training, we iteratively run lt but constrain the computational bound to control the training process by setting a maximum number of iteration steps. however, iteration steps are dynamic in inferring. ## experiment we divide our empirical studies and experiments into two genres: 1) we experiment with nmt (our main task) to confirm the effectiveness of our methods for large-scale sequence modeling and further quantitatively justify our hypotheses and assumptions; 2) we attempt pre-training tasks and two somewhat rare but challenging tasks: the learning to execute task @xcite and the lambada @xcite task, to observe the performance on tasks of varying complexity. all the links of datasets, libraries, scripts, and tools marked with are listed in appendix h. we open source code on github. training our code is implemented on tensorflow 2.6 (abadi et al., 2016) with 4 nvidia ti-tan xp 12g gpu. we implement our model based on the codebase of official ut from tensor2tensor foot_4 and official cka . we use the default setting: universal_transformer_base from tensor2tensor. concretely, we use adam optimizer (kingma and ba, 2015) with parameters β 1 = 0.9,β 2 = 0.997 @xmath0, and a dynamic learning rate with @xmath1, 7e -4 ]) is employed. we set dropout regularization with a drop rate @xmath2.1 and label smoothing with @xmath3.1 @xcite . for data feeding efficiency, each batch of similar-length sequences are padded to the same length and may have a different number of elements in each batch. reimplementation and reconfiguration we reimplement some models on our machine with the same batch size. we compare the reimplemented results to the reported results on the same test set to ensure the difference is less than 5% (or 1 in bleu). then, we can confirm the reimplementation and reconfiguration. ## conclusion in this work, we place our lazy transition on top of ut to build lt. our lazy transition leverages the relaxed equilibrium for sequence modeling and provides step identifications that the model ponders the importance of every step for different tokens throughout the iteration. our main experiment shows that lt can achieve strong performance on translation tasks, facilitate the training of deep models, and tackle the challenge of zero-shot inferring. our method retains the recurrent inductive bias learned by its ut component, which is confirmed by our secondary experiments. lt tends to focus on token-specific characteristics at the early steps and then turns to find sequence-specific ones at the late steps, especially in deep settings. meanwhile, stable and smooth behaviors in the early iteration are significant. although there are some practical limitations, as we mentioned in this paper, we believe our lazy transition is a novel perspective to reconsider the models based on iterative refinements in sequence modeling.