Deep Equilibrium Non-Autoregressive Sequence Learning (2023)

Abstract: In this work, we argue that non-autoregressive (NAR) sequence generative models can equivalently be regarded as an iterative refinement process towards the target sequence, implying an underlying dynamical system of NAR model: z = f (z, x) → y. In such a way, the optimal prediction of a NAR model should be the equilibrium state of its dynamics if given infinitely many iterations. However, this is infeasible in practice due to limited computational and memory budgets. To this end, we propose DEQNAR to directly solve for the equilibrium state of NAR models based on deep equilibrium networks (Bai et al., 2019) with black-box root-finding solvers and back-propagate through the equilibrium point via implicit differentiation with constant memory. We conduct extensive experiments on four WMT machine translation benchmarks. Our main findings show that DEQNAR can indeed converge to a more accurate prediction and is a general-purpose framework that consistently helps yield substantial improvement for several strong NAR backbones.

## introduction transformer @xcite has recently become the most prevailing neural architecture for sequence-to-sequence learning @xcite . transformer is originally an autoregressive (ar) sequence generative models, which adopts a sequential factorization to estimate the conditional probability of a target sequence @xmath0, • • • , y [n ] } conditioned on a source sequence x: p(y|x) = n @xmath1:n-1] , x). albeit simple and effective, such a fixed left-to-right restriction is not necessarily the unique and the best formulation for sequence modeling, limiting the design space of neural networks and applicable tasks for ar models. hence researchers are motivated to study non-autoregressive (nar) sequence generative models @xcite as an alternative to ar models, which instead use a per-token factorization p(y|x) = n @xmath2. despite their favorable decoding speed and flexible formulation to introduce constraints, nar models still lag behind their ar counterparts and require data distillation. nar models can be viewed as generating sequences by iteratively denoising from an initial guess (fig. 1(a) ). several studies based on this idea of iterative refinement show promising and competitive results compared ar models. for instance, @xcite and @xcite propose to regard nar models as denoising autoencoders, while @xcite task nar models with conditional masked language modeling. more recently, discrete denoising diffusion models have started to attract the community's attention. besides iteratively manipulating sequences of discrete tokens, researchers also find that for fully nar models @xcite , layer recurrence also calibrates intermediate continuous representations towards the target discrete sequence @xcite @xcite . in other words, fully nar and iterative-based nar models are tasked with approaching their equilibrium states, in terms of either discrete or continuous representation, which is also found in our empirical observations (fig. 1(c) ). to go a step further, we argue that nar models, including fully nar and iterative nar models, can unifiedly be regarded as a dynamical system in the form of z @xmath3, x), implying a dynamics of parallel denoising or iterative refinement process over the whole sequence (figure 1(b) ). more concretely, nar models apply a markov chain factorization to a series of intermediate predictions from the bottom up, where a neural parametric transition kernel f θ learns denoising sequences in a coarseto-fine manner, while z t is the t-th running discrete or continuous state. from such a unified dynamical system perspective, intuitively, the state of an nar system is to pursue the answer to this question, we propose to use deq networks @xcite to directly solve for such equilibrium state of nar systems. formally, given the input x, a transition kernel f ✓ parameterized by deep neural networks ✓ (e.g., transformer), we define an nar sequence generative model by the following dynamical system and solve its equilibrium state z ? as a root-finding problem: z @xmath4, x), =) z ? = rootfind(g z ; x, ✓), where g @xmath5, x) z, as aforementioned, nar models can be categorized by performing either explicit or implicit iterative refinement. in the following subsections, we will describe how to model implicit, explicit nar models and the combination of the both under the proposed deqnar framework, in accordance with different choices of the definition of the state z and the transition function f ✓ , which we summarize in to pursue the answer to this question, we propose to use d solve for such equilibrium state of nar systems. formall parameterized by deep neural networks ✓ (e.g., transforme model by the following dynamical system and solve its equ z @xmath6, x), =) z ? = rootfind(g z where z 0 is the initial condition. as aforementioned, nar models can be categorized by per refinement. in the following subsections, we will descr models and the combination of the both under the proposed different choices of the definition of the state z and the tra in transformer layer f, which sequentially computes a self-attention (san), cro and feed-forward (ffn) blocks, each of which module is followed by layer norm 2016). the initial condition z 0 = emb(hmaski) is set to be an embedding sequ tokens. in the extreme case, we assume an infinite-depth transformer which is powerfu layer is capable of refining the representation. intuitively, the quality of a ser states {z 0 , • • • , z t 1 , z t , z t+1 , • • • , z 1 } would be approximately sorted in an asc the goodness is bounded, z t must converge to some fixed point, denoted by z ? assume that z t would reach an equilibrium state which satisfies z ? = f (z ? ). ther problem of our interest becomes how to compute the equilibrium state z ? . given that z is a continuous variable, we can use any advanced black-box root e.g., newton or quasi-newton methods like broyden's methods (broyden, 1 acceleration @xcite . these methods guarantee a much faster and better-q than the case where we perform infinitely many naïve unrolling steps, which is no to computational and memory budgets. ## method: deep equilibrium nar sequence models to pursue the answer to this question, we propose to use deq networks @xcite to dire solve for such equilibrium state of nar systems. formally, given the input x, a transition kerne parameterized by deep neural networks ✓ (e.g., transformer), we define an nar sequence genera model by the following dynamical system and solve its equilibrium state z ? as a root-finding prob as aforementioned, nar models can be categorized by performing either explicit or implicit itera refinement. in the following subsections, we will describe how to model implicit, explicit n models and the combination of the both under the proposed deqnar framework, in accordance w different choices of the definition of the state z and the transition function f ✓ , which we summa in to pursue the answer to this question, we propose to use deq network solve for such equilibrium state of nar systems. formally, given the i parameterized by deep neural networks ✓ (e.g., transformer), we define model by the following dynamical system and solve its equilibrium state as aforementioned, nar models can be categorized by performing eith refinement. in the following subsections, we will describe how to m models and the combination of the both under the proposed deqnar fr different choices of the definition of the state z and the transition funct in to this end, in this paper, we reformulate sequence generation problems as solving the equilibrium state of nar models. we propose a generalpurpose framework, the deqnar, based on deep equilibrium networks @xcite , and apply it to the cases where the iterative refinement can be conducted either in continuous feature state space, discrete data state space, or a combination of both. this enables multiple preferable properties for our model over previous studies. (1) instead of naive iterative layer stacking, deqnar models define the output as fixed point of f θ given the input x, i.e., z ⋆ = f (z ⋆ , x), modeling an equilibrium representation. (2) compared with typical nar systems, the proposed deqnar permits better convergence to the equilibrium point. we can leverage any advanced black-box solvers, e.g., quasi-newton methods, to directly solve for the equilibrium solution, which often leads to better results. (3) the deqnar is also orthogonal to existing advanced techniques for nar models, for which we studied its effectiveness when combined with the current best practices, including better modeling approach (vae, gu and kong, 2021), training objective (ctc, @xcite and training strategy (glat, @xcite . we conduct extensive experiments on wmt14 english-german and wmt16 english-romanian machine translation benchmarks. based on the empirical results, our main findings are as follows: formally, let @xmath8, . . . , y [n ] ] ∈ {0, 1} n ×|v| within the vocabulary space v be a target sequence of interest, and @xmath9, . . . , x [|x|] ] be the conditional source sequence. non-autoregressive sequence-to-sequence learning aims to learn a probabilistic model p(y|x) measuring the likelihood of target sequence given its source sequence: such parameterization shares a similar form with a first-order markov chain, where the probability of y is a marginalization over all possible intermediate paths z 0•••t . the state z t evolves through layers in a bottom-up fashion, and the input u @xmath10. we provide an illustration in fig. 1(a ). for clarity of notations, we can classify nar models by the explicitness of tokens in the intermediate states: ## deqnar: a deep equilibrium nar sequence learning framework to answer this question, we propose to directly solve for such an equilibrium state of nar systems based on the use of deq networks @xcite as a critical tool. formally, given the input x, a transition kernel f θ parameterized by deep neural networks θ (e.g., transformer), we define a nar sequence generative model by the following dynamical system and solve its equilibrium state z ⋆ as a root-finding problem: as aforementioned, nar models can be categorized by performing either explicit or implicit iterative refinement. we will introduce how to accomondate implicit, explicit nar models and the combination of both under the proposed deqnar framework, in accordance with different choices of the definition of the state z and the transition function f θ , which we summarize in ## experiments we conduct extensive experiments on standard machine translation benchmarks to inspect deqnar's performance on sequence-to-sequence tasks. we demonstrate that deqnar produces better results over its nar backbones. we also show figure 2 : applying deqnar to different nar backbone models, i.e., vanilla nat @xcite , glat @xcite , ctc @xcite , glat+ctc (using greedy decoding, and vae+ctc @xcite ). note that here all the backbone models are our own implementations. ## conclusions in this work, we revisit non-autoregressive (nar) sequence generative models from the perspective of dynamical systems. we then propose deqnar that can directly solve for its equilibrium state to better estimate the desired target sequence. we conduct extensive empirical experiments demonstrating that the proposed deqnar framework can indeed converge to the equilibrium state, which consistently improves several nar backbones. limitations. while these findings are promising, there remain several limitations, e.g. the accuracyefficiency trade-off we have discussed above. another one is the existence of knowledge distillation performance bound. typical nar models need sequence-level knowledge distillation (kd) by an ar teacher model, which imposes an upper bound of performance for the nar student models. this is an obvious limitation for nar models in general, since the current strong nar baselines have been closely approaching this kd upper bound (e.g., ar transformer's 27.60 bleu on wmt14 de-en). as such, the nar community has reached a point to call for progress in eliminating the need of kd, where we notice the recent advances of kd-free nar sequence models such as da-transformer @xcite which uses directed acyclic graph (dag) for probabilistic modeling on raw data. in principle, deqnar is architecture-agnostic, and taking advantage of these new kd-free approaches is a promising future work, and we leave for further study.