Deep Equilibrium Non-Autoregressive Sequence Learning
Year: 2023

Abstract:
In this work, we argue that non-autoregressive (NAR) sequence generative models can equivalently be regarded as iterative refinement process towards the target sequence, implying an underlying dynamical system of NAR models: 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 rootfinding 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.

heading: Future Work
text: We hope that the DEQNAR can also serves as an effective and universal “solver” that can be integrated with and thus facilitates future dvances of NAR sequence approaches. 2 NAR Models as Dynamical Systems
NAR Models as Markov Process of Iterative Refinement in General. We formulate NAR models
11764
based on Transformer (Vaswani et al. , 2017) as a Markov chain. There are mainly two categories of NAR models: fully NAR and iterative-based NAR models. Both of them can be unified under a general perspective of dynamical systems conducting iterative refinement process over some intermediate state, where the parametric transition function is in a form of zt = fθ(zt−1,ut) with zt as the running state of the systems. Formally, let y = [y[1],. , y[N ]] ∈ {0, 1}N×|V| within the vocabulary space V be a target sequence of interest, and x = [x[1],. , 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:
pθ(y|x) = ∑
z0,··· ,zT pθ(y, z0,. ,zT |x)
= ∑
z0,··· ,zT pθ(y|zT ,x)
T∏
t=1
fθ(zt|zt−1,x),
where zt is the t-th intermediate state which is varied across different NAR models, f(zt|zt−1,x) is the transition function from the (t−1)-th step to the (t-th step parameterized by fθ, and p(y|zt,x) =∏
p(y[n]|zt,x) is the predicted probability made in parallel under the conditional independence assumption among the elements of y. 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 z0···T. The state zt evolves through layers in a bottom-up fashion, and the input ut = x is time-invariant or constant in sequence-to-sequence learning scenarios. 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:
• Explicit iterative refinement: Iterative NAR models (Lee et al. , 2018; Ghazvininejad et al. , 2019) perform iterative refinement within discrete space, producing discrete tokens explicitly for each iteration. The t-th system state is the discrete representation zt ∈ {0, 1}N×|V| (i. e. , the indices of tokens). The transition function f (i. e. the whole decoder) learns to refine the tokens in the previously generated sentence until meeting a certain condition (e. g. , no further improvement or reaching a maximum number of iterations). • Implicit iterative refinement: Fully NAR models (Gu and Kong, 2021; Qian et al. , 2021) can also be viewed as implicitly conducting iterative refinement within continuous feature space, given the nature of multi-layer neural networks. The t-th system state for such implicit iterative refinement is contextualized continuous representation zt ∈ RN×d (i. e. , dense vectors). The transition function f (i. e. a decoder layer) is supposed to learn to refine representations layer by layer such that the discrete data can be best described. Motivation. Based on such a dynamical system view of the NAR sequence-to-sequence learning, one can use dynamical system-inspired methods for better understanding and improved modeling. For an NAR system that conducts iterative refinement over the whole sequence towards the target sequence limt→∞ f(zt,x) = z⋆ → y, we may want to find the solution z⋆ of such a system that best estimate the target data, which is a local optima, or an equilibrium state of this system. However, as seen in Fig 1(a), the current NAR systems can be considered as resorting to a naive solver that recurrently applies the transition function f up to a manually-defined maximum iteration N , which cannot guarantee reaching the equilibrium solution, leading to a sub-optimal representation in terms of the target sequence. This motivates us to seek the answer to an arisen question: Can we find such an equilibrium state of the NAR dynamical system, which can give rise to a better solution?
3 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 (Bai et al. , 2019) 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:
zt = fθ(zt−1,x)
=⇒ z⋆ = RootFind(gz = 0;z0, θ), (1) where gz := fθ(z,x)− z,
and z0 is the initial condition. 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 Table 1. 3. 1 Case I: Implicit Iterative Refinement in Continuous (Feature) State Space
Here we explain the intuition behind how solving for the equilibrium state is connected with implicit iterative refinement through an extreme case. Assuming an infinite-depth Transformer that is powerful enough, and each layer is capable of refining the representation. Intuitively, the quality of a series of intermediate states {z0, · · · , zt−1, zt, zt+1, · · · , z∞} would be approximately sorted in an ascending order. Since the goodness is bounded, it is reasonable to assume that zt may converge to some fixed point, denoted by z⋆, which is an equilibrium state that satisfies z⋆ = f(z⋆). Therefore, the inference problem of our interest becomes how to compute the equilibrium state z⋆. Formally, for NAR models that conduct implicit iterative refinement (Gu et al. , 2018; Gu and Kong, 2021), the state zt is defined as the continuous hidden representation of intermediate layer, while the transition function fθ is parameterized a single Transformer decoder layer F , which sequentially computes a self-attention, cross-attention and feedforward blocks, each of which module is followed by layer normalization (Ba et al. , 2016). The initial condition z0 = emb(⟨mask⟩) is set to be an embedding sequence full of ⟨mask⟩ tokens. Our solution to find the continuous variable z∗ is to use advanced black-box root solving algorithms, e. g. , Newton or quasi-Newton methods like Broyden’s methods (Broyden, 1965), or Anderson acceleration (Anderson, 1965). These methods guarantee a much faster and better-quality convergence than the case where we perform infinitely many naïve unrolling steps, which is not even realistic due to computational and memory budgets. 3. 2 Case II: Explicit Iterative Refinement in Discrete (Data) State Space
As for NAR models that conduct explicit iterative refinement on discrete tokens (Lee et al. , 2018; Ghazvininejad et al. , 2019), the state zt is defined as a sequence of one-hot vectors corresponding to each of the intermediate predicted tokens. The
transition function fθ : {0, 1}N×|V| → RN×|V| is parameterized by a multi-layer Transformer decoder followed by a softmax normalization, and a final discretization operator such as argmax or sampling. The initial condition z0 = ⟨mask⟩ is set to be a sequence full of ⟨mask⟩ tokens. Two challenges exist while aiming to solve for the equilibrium point of fθ(z,x) = z regarding discrete tokens. (1) Intractable: z lies in a very high-dimensional space ({0, 1}N×|V|), whereas the cardinality of the feasible set is very small (the vocabulary size). Finding the solution is almost intractable, especially for highly non-linear neural networks. (2) Non-differentiable: Root-finding algorithms such as Newton or quasi-Newton methods require computing or estimating Jacobian inverse, which is numerically unstable or even infeasible to obtain, for transition functions f that contain non-differentiable sampling operators. Our solution is to leverage the expected embedding weighted by the softmax probabilities as a continuous surrogates of z:
zt ≈ E [emb(z̃t)] , where z̃t ∼ fθ(·|zt−1) (2)
Such approximation helps ease the two problems: (1) a point is projected to the simplex formed by feasible points, greatly restricting the search space. (2) the “soft” embedding makes the neural network differentiable. Another possible solution is to use score function gradient estimators (e. g. , REINFORCE (Williams, 1992)) for these non-differentiable operators, which, however, are known to be computationally expensive and of high variance nature. The major difference with the implicit case is that the continuous relaxation z̃t represents contextless token identity instead of contextualized deep representation. Issues may arise if the token identity is non-informative and not context-aware, preventing the model from evolving the state efficiently. This motivates us to develop the contextinformed version in the next section. 3. 3 Case III: Mixed Explicit and Implicit Iterative Refinement
In practice, explicit iterative refinement methods are aware of the strong condition signal by the immediate prediction of the last iteration, usually leading to better results than the implicit (or the fully NAR) methods. As aforementioned, however, it is non-trivial if we want to use DEQ to solve for
explicit iterative refinement, in which continuous surrogates are required. To take the best of both implicit and explicit iterative refinement, we propose an indirect way to extend DEQ by introducing layer-wise predictionawareness (Huang et al. , 2021), and refer such hybrid variant to DEQNAR-MIXED. Concretely, we make an intermediate prediction p̃(yt|·) in every layer evaluation, and feed emb[yt] the embedding of the most probable predicted token into zt:
zt ← σ(zt, emb[yt]), where ỹt = argmax p̃(·|zt)
where a fusion operator σ : Rd × Rd → Rd parameterized by a position-wise MLP. Our goal is that in such as way DEQNAR-MIXED can endow f with the awareness of running prediction made so far, which helps the model for better calibration. 3. 4 Learning via Implicit Differentiation Typically for explicit neural networks, we can directly back-propagate through the stacking layers using automatic differentiation tools (Baydin et al. , 2018). However, for implicit models like DEQ, it is computationally expensive if we unroll the iteration path of the internal optimization problem. In this section, we introduce how to train the proposed model with only knowing its equilibrium state. Moreover, we also introduce regularizations to stabilize its convergence dynamics (see Appendix B). Based on the implicit function theorem (IFT, Krantz and Parks, 2002), the DEQ model can differentiate through its fixed point without unfolding and storing intermediate states in the forward trajectory. Specifically, given fixed-point state z⋆, the task-specific loss function L (e. g. , cross-entropy), the gradients of DEQ with regard to the parameter θ and input x are given by:
∂L ∂θ = ∂L ∂z⋆
( I − ∂fθ
∂z⋆ )−1∂fθ(z⋆,x) ∂θ
∂L ∂x = ∂L ∂z⋆
( I − ∂fθ
∂z⋆ )−1∂fθ(z⋆,x) ∂x. (3)
This theorem enables us to decouple the forward and backward passes of DEQ-based models, i. e. for parameter update, we only need the final output z⋆ and do not need to perform back-propagation through the unrolled chain of forwarding passes. This allows one to train implicit networks in a modern end-to-end manner while consuming only O(1) memory cost. For the explicit case we discussed in §3. 2, we introduce a two-stage training scheme to prevent being hindered by the inaccurate and over-smooth predicted probability distribution in Equation 2. We first pretrain the NAR model as denoising autoencoders similar to (Savinov et al. , 2021) except that we use full-masking as the corruption function instead of uniform sampling from the vocabulary. We then integrate the pretrained model into our DEQNAR framework and use implicit differentiation for further finetuning, leading to the final model that learns to make predictions at its equilibrium states. On the Differences with DEQ-Transformer (Bai et al. , 2019). Note that the original work of Bai et al. (2019) has proposed DEQ-Transformer demonstrating its success for autoregressive language modeling. Compared with evolving a single word state in the DEQ-Transformer, evolving all words’ states simultaneously is a highly structured and complicated problem. Using DEQ to solve NAR problems is non-trivial since the interdependencies among all words are prone to cause inconsistency and instability, which remains challenging and yet unexplored. Related work