SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shangtong Gui $^{1,2,3}$ ,Chenze Shao $^{2,3}$ ,Zhengrui Ma $^{2,3}$ ,Xishan Zhang $^{1,4}$ ,Yunji Chen $^{1,3}$ ,Yang Feng $^{2,3*}$
$^{1}$ State Key Lab of Processors,
$^{2}$ Key Laboratory of Intelligent Information Processing,
<sup>3</sup>University of Chinese Academy of Sciences
4Cambricon Technologies
Non-autoregressive Transformer(NAT) significantly accelerates the inference of neural machine translation. However, conventional NAT models suffer from limited expression power and performance degradation compared to autoregressive (AT) models due to the assumption of conditional independence among target tokens. To address these limitations, we propose a novel approach called PCFG-NAT, which leverages a specially designed Probabilistic Context-Free Grammar (PCFG) to enhance the ability of NAT models to capture complex dependencies among output tokens. Experimental results on major machine translation benchmarks demonstrate that PCFG-NAT further narrows the gap in translation quality between NAT and AT models. Moreover, PCFG-NAT facilitates a deeper understanding of the generated sentences, addressing the lack of satisfactory explainability in neural machine translation.[2]
Autoregressive machine translation models rely on sequential generation, leading to slow generation speeds and potential errors due to the cascading nature of the process. In contrast, non-autoregressive Transformer (NAT) models [13] have emerged as promising alternatives for machine translation. The Vanilla NAT model assumes that the target tokens are independent of each other given the source sentence. This assumption enables the parallel generation of all tokens. However, it neglects the contextual dependencies among target language words, resulting in a severe multi-modality problem [13] caused by the rich expressiveness and diversity inherent in natural languages. Consequently, Vanilla NAT models exhibit significant performance degradation compared to their autoregressive counterparts.
Several approaches have been proposed to alleviate the multi-modality problem in NAT models while maintaining their computational efficiency. CTC-based NAT [11, 22] depicted in Figure 1(b) addresses this issue by implicitly modeling linear adjacent dependencies through token repetition and blank probability. This approach helps alleviate the multi-modality of token distributions at each position compared to Vanilla NAT. Similarly, DA-Transformer [14, 37, 26] shown in Figure 1(c) introduces a hidden Markov model and employs a directed acyclic graph to model sentence probability, enabling the model to capture dependencies between adjacent tokens and reduce the multi-modal nature of the distribution.
*Corresponding author: Yang Feng
$^{2}$ Code is publicly available at https://github.com/ictnlp/PCFG-NAT.
(a) Vanilla NAT
(b) CTC-NAT
(c) DA-Transformer
(d) PCFG-NAT
Figure 1: Illustration of Vanilla NAT, CTC-NAT, DA-Transformer, and PCFG-NAT
However, despite these enhancements, the modeling of only adjacent and unidirectional dependencies limits the model's ability to capture the rich semantic structures with non-adjacent and bidirectional dependencies present in natural language. For example, considering the token need as the previous token, the distribution of the next token is still affected by the issue of multi-modality, as the word boat can be modified by various rich prefixes in the corpus.
To overcome this constraints, we propose PCFG-NAT, a method that leverages context-free grammar to endow the model with the capability to capture the intricate syntactic and semantic structures present in natural language. For instance, as illustrated in Figure 1(d), PCFG-NAT can directly model the robust and deterministic correlation between non-adjacent tokens, such as You need boat. In the parse tree generated by PCFG-NAT for the target sentence, the nodes connected to the tokens You need boat constitute the main chain of the tree, capturing the left-to-right semantic dependencies through linear connections. Additionally, PCFG-NAT can learn relationships between modifiers, such as are gonna and a bigger, and the words they modify by constructing a local prefix tree for the nodes on the main chain. This allows for the modeling of right-to-left dependencies between the modified words and their modifiers.
We conduct experiments on major WMT benchmarks for NAT (WMT14 En  $\leftrightarrow$  De, WMT17 Zh  $\leftrightarrow$  En, WMT16 En  $\leftrightarrow$  Ro), which shows that our method substantially improves the translation performance and achieves comparable performance to autoregressive Transformer [40] with only one-iteration parallel decoding. Moreover, PCFG-NAT allows for the generation of sentences in a more interpretable manner, thus bridging the gap between performance and explainability in neural machine translation.
In this section, we will present the necessary background information to facilitate readers' comprehension of the motivation and theoretical underpinnings of PCFG-NAT. In Section 2.1, we will provide a formalization of the machine translation task and introduce the concepts of autoregressive Transformer [40] and Vanilla NAT [13], which will serve as the basis for our proposed approach. Subsequently, in Section 2.2, we will provide a precise and rigorous definition of PCFG, accompanied by illustrative examples aimed at enhancing clarity and comprehension.
2.1 Non-autoregressive Translation
Machine translation can be formally defined as a sequence-to-sequence generation problem. Given source sentence  $X = \{x_{1}, x_{2}, \dots, x_{S}\}$ , the translation model is required to generate target sentence  $Y = \{y_{1}, y_{2}, \dots, y_{T}\}$ . The autoregressive translation model predicts each token on the condition of all previous tokens:
$$
P (Y | X) = \prod_ {t = 1} ^ {T} p \left(y _ {t} \mid y _ {<   t}, X\right). \tag {1}
$$
Figure 2: An example parse tree for You are gonna need a bigger boat.
Applied Rules	Sequence
Start	V1
V1 → V0You V5,V0 → ε	You V5
V5 → V3 need V9	You V3 need V9
V3 → V0are V4,V0 → ε	You are V4 need V9
V4 → gonna	You are gonna need V9
V9 → V7 boat V10	You are gonna need V7 boat V10
V7 → V0a V8,V0 → ε	You are gonna need a V8 boat V10
V8 → bigger	You are gonna need a bigger boat V10
V10 → .	You are gonna need a bigger boat .
Table 1: Derivation of the parse tree. Used production rules for every step are listed, where  $\epsilon$  stands for an empty string.
The autoregressive inference is time-consuming as it requires generating target tokens from left to right. To reduce the decoding latency, Gu et al. [13] proposes non-autoregressive Translation, which discards dependency among target tokens and predicts all the tokens independently:
$$
P (Y | X) = \prod_ {t = 1} ^ {T} p \left(y _ {t} | X\right). \tag {2}
$$
Since the prediction does not rely on previous translation history, the decoder inputs of NAT are purely positional embeddings or the copy of source embeddings. Vanilla NAT trains a length predictor to determine the sentence length when decoding. During the training, the golden target length is used. During the inference, the predictor predicts the target length and the translation of the given length is obtained by argmax decoding.
2.2 Probabilistic Context-free Grammar
Probabilistic Context-free Grammar is built upon Context-Free Grammar (CFG). A CFG is defined as a 4-tuple:
$$
G = (\mathcal {N}, \mathcal {T}, \mathcal {R}, S),
$$
where  $\mathcal{N}$  is the set of non-terminal symbols,  $\mathcal{T}$  is the set of terminal symbols,  $\mathcal{R}$  is the set of production rules and  $S$  is the start symbol. PCFG extends CFG by associating each production rule  $r\in R$  with a probability  $P(r)$ . Production rules are in the form:
$$
A \rightarrow \alpha , \tag {3}
$$
where  $A$  must be a single non-terminal symbol, and  $\alpha$  is a string of terminals and/or nonterminals,  $\alpha$  can also be empty. Starting with the sequence that only has start symbol  $S$ , we can replace any non-terminal symbol  $A$  with  $\alpha$  based on production rule  $A \rightarrow \alpha$ . The derivation process repeats until the sequence only contains terminal symbols and the derived hierarchical structure can be seen as a Parse Tree. Those production rules and their probabilities can be applied regardless of the contexts, which provides Markov property for the joint probability of the parse tree. For every rule,  $A \rightarrow \alpha$  applied in a derivation, all of the symbols in  $\alpha$  are the children nodes of  $A$  in order. The probability of a parse tree  $t$  is formulated as:
$$
P (t) = \prod_ {r \in R (t)} (P (r)) ^ {n _ {r}}, \tag {4}
$$
where  $R(t)$  is the set of production rules applied in  $t$  and  $n_r$  is the number of times that  $r$  is applied. For sentence You are gonna need a bigger boat., Figure 2 shows an example parse tree and how the sentence is derived from the start symbol  $V_1$ .
In the following sections, we present the formal definition of a novel variant of PCFG designed for PCFG-NAT in Section 3.1 and the architecture of the PCFG-NAT model in Section 3.2. Finally, we describe the training algorithm and decoding algorithm for PCFG-NAT in Section 3.3 and Section 3.4.
3.1 Right Heavy PCFG
Despite the strong expressive capabilities of PCFGs, the parameter induction process becomes challenging due to the vast latent variable space encompassing the possible parse tree structures of a single sentence. Moreover, training PCFGs incurs significant time complexity [35]. To address these challenges, we propose a novel variant of PCFG called Right-Heavy PCFG (RH-PCFG). RH-PCFG strikes a balance between expressive capabilities and computational complexity, mitigating the difficulties associated with parameter induction and reducing training time.
3.1.1 Derived Parse Tree of RH-PCFG
RH-PCFG enforces a distinct right-heavy binary tree structure in the generated parse tree. Specifically, for each non-terminal symbol in the tree, the height of the right subtree increases proportionally with the length of the derived sub-string of the non-terminal symbol, while the height of the left subtree remains constant. This characteristic significantly reduces the number of possible parse trees for a given sentence, thereby reducing training complexity.
In the context of a right-heavy binary tree, the traversal process begins with the root node and proceeds recursively by selecting the right child of the current node as the subsequent node to visit. This sequential path, known as the Main Chain, is guided by a design principle inspired by linguistic observations in natural language. Specifically, it reflects the tendency for sentence structures to display left-to-right connectivity, with local complements predominantly positioned to the left of the central core structure. Within this framework, the left subtree of each node along the main chain is referred to as its Local Prefix Tree. The parse tree in Figure 2 represents a right heavy tree. In this tree,  $V_{1}, V_{5}, V_{9}, V_{13}$  form the main chain, and their left subtrees correspond to their respective local prefix trees.
3.1.2 Support Tree of RH-PCFG
To make sure that all the parse trees are the right heavy tree, we first build up a Support Tree as the backbone for the parse trees derived from RH-PCFG. The construction of the support tree follows two steps: At first, given a source sentence with a length of  $L_{x}$  and an upsample ratio of  $\lambda$ , we sequentially connect  $\lambda L_{x} + 1$  nodes with right links, which become the candidates for the main chain nodes in parse trees. Next, excluding the root node, we append a complete binary tree of depth  $l$  as the left subtree to each node. For the root node, only one left child is added. The added nodes represent local prefix tree candidates to the attached main chain candidates. Consequently, for a source sentence of length  $L_{x}$ , the number of non-terminal symbols  $m$  of the RH-PCFG is
$$
m = \lambda L _ {x} \cdot 2 ^ {l} + 2. \tag {5}
$$
The nodes within the support tree are assigned sequential indexes through an in-order traversal, and these indexes correspond to the corresponding non-terminal symbols sharing the same index. Detailed construction algorithm can be found in Appendix A.1. Figure 3 shows an example of PCFG parse tree with  $L_{x} = 3$ ,  $\lambda = 1$ ,  $l = 2$ .
Figure 3: An example support tree of Right Heavy PCFG
In RH-PCFG, the set of non-terminals is  $\mathcal{N} = \{V_0,\dots,V_{m - 1}\}$  and the set of terminals  $\mathcal{T}$  is equivalent to the vocabulary. The collection of production rules is represented by  $R$ , and the starting symbol is  $V_{1}$ . The production rules in RH-PCFG have the following form:
$$
V _ {0} \rightarrow \epsilon , \tag {6}
$$
$$
V _ {i} \rightarrow a, a \in \mathcal {T}, i \neq 0, \text {I s L e a f} (i) \tag {7}
$$
$$
V _ {i} \rightarrow V _ {j} a V _ {k}, V _ {i}, V _ {j}, V _ {k} \in \mathcal {N}, a \in \mathcal {T}, (\text {I n L e f t R e a c h} (j, i) \vee j = 0) \wedge \text {I n R i g h t R e a c h} (k, i) \tag {8}
$$
$\epsilon$  stands for an empty string.  $\operatorname{IsLeaf}(i)$  indicates whether the non-terminal symbol  $V_{i}$  corresponds to a leaf node in the support tree.  $\operatorname{InLeftReach}(j,i)$  denotes that, in the support tree, the node corresponding to  $V_{j}$  is in the left subtree of the node corresponding to  $V_{i}$ .  $\operatorname{InRightReach}(k,i)$  signifies that the node corresponding to  $V_{k}$  is in the right subtree of the node corresponding to  $V_{i}$  in the support tree, and if  $V_{i}$  corresponds to a main chain node,  $V_{k}$  has to correspond to a main chain node at the same time.
3.2 Architecture
The main architecture of PCFG-NAT is identical to that of the Transformer [40], with the additional incorporation of structures to model the probabilities of the PCFG. The input sentence to be translated is fed into the encoder of the model. Additionally, based on the length  $L_{x}$  of the source sentence, we compute the number of corresponding non-terminal symbols  $m$  in the RH-PCFG. Then  $m$  positional embeddings are used as inputs to the decoder. Through multiple layers of self-attention modules with a global attention view, cross-attention modules that focus on distant information, and subsequent feed-forward transformations, we obtain  $m$  hidden states  $h_0, \dots, h_{m - 1}$ . These hidden states are associated with non-terminal symbols of the same index and are utilized to calculate the rule probabilities related to each corresponding non-terminal symbol.
Figure 4: The architecture of PCFG-NAT model. A PCFG is inserted between the decoder layers and the output layer, which captures the dependency between different tokens. A parse tree is derived to generate the sentence You are gonna need a bigger boat. and is highlighted with bold black lines.
The probability of the rule  $P(V_{i} \to a)$  can be regarded as the word prediction probability at each position, and we can directly calculate it with the softmax operation. We use  $h_{i}$  to denote the hidden state representation of the non-terminal symbol  $V_{i}$  and use  $W_{o}$  to denote the weight matrix of the output layer. Formally, we have:
$$
P \left(V _ {i} \rightarrow a\right) = \operatorname {S o f t m a x} \left(W _ {o} h _ {i}\right). \tag {9}
$$
For the rule  $P(V_{i} \to V_{j}aV_{k})$ , we can independently consider the prediction probability of word  $a$  and the transition probability between the three non-terminal symbols. Formally, we decompose it into two parts:
$$
P \left(V _ {i} \rightarrow V _ {j} a V _ {k}\right) = P \left(<   V _ {j}, V _ {k} > | V _ {i}\right) \cdot P (a | V _ {i}). \tag {10}
$$
The prediction probability  $P(a|V_i)$  can be calculated with the softmax operation like in Equation 9. To normalize the distribution of  $P(< V_j, V_k > |V_i)$ , we have to enumerate all valid combinations of  $j, k$ . Let  $\mathbf{Child}(i) = \{ < o, u > |V_i \to V_0aV_u \in \mathcal{R} \}$  denotes the set of children combinations of the non-terminal symbol  $V_i$ , we can formally define the transition probability as:
$$
q _ {i} = W _ {q} h _ {i}, \quad q _ {j} = W _ {l} h _ {j}, \quad q _ {k} = W _ {r} h _ {k}, \tag {11}
$$
$$
P \left(<   V _ {j}, V _ {k} > \mid V _ {i}\right) = \frac {\exp \left(q _ {i} ^ {T} q _ {j} + q _ {i} ^ {T} q _ {k} + q _ {j} ^ {T} q _ {k}\right)}{\sum_ {<   o , u > \in \mathbf {C h i l d} (i)} \exp \left(q _ {i} ^ {T} q _ {o} + q _ {i} ^ {T} q _ {u} + q _ {o} ^ {T} q _ {u}\right)}, \tag {12}
$$
where  $W_{q}, W_{l}, W_{r}$  are learnable model parameters.
3.3.1 CYK Algorithm for RH-PCFG
We train PCFG-NAT with the maximum likelihood estimation, which sums up all the possible parse trees for the given sentence  $Y$  conditioned on source sentence  $X$  and model parameters  $\theta$ :
$$
\mathcal {L} (\theta) = - \log P (Y | X, \theta) = - \log P \left(V _ {1} \Rightarrow y _ {0} y _ {1} \dots y _ {n - 1}\right). \tag {13}
$$
In PCFG, the CYK algorithm [35] is employed to compute the probability of a given sentence derived from the start symbol. The CYK algorithm follows a bottom-up dynamic programming approach to fill in the table  $S$ . Each element  $S_{i,j}^{a}$  in the table represents the sum of probabilities of that  $V_{a}$  derives the consecutive sub-string  $y_{i}\ldots y_{j}$  of the target sentence, denoted as  $P(V_{a}\Rightarrow y_{i}\ldots y_{j})$ . Therefore the size of table  $S$  is  $\mathcal{O}(mn^2)$ . Based on the definition of  $S_{i,j}^{a}$ , we have recursion formula
$$
S _ {i, j} ^ {a} = \sum_ {k = i + 1} ^ {j - 1} \sum_ {\langle b, c \rangle \in \mathbf {C h i l d} (a)} P \left(V _ {a} \rightarrow V _ {b} y _ {k} V _ {c}\right) S _ {i, k - 1} ^ {b} S _ {k + 1, j} ^ {c} \tag {14}
$$
For every non-terminal symbol  $V_{a}$ , the complexity filling all the  $S_{i,j}^{a}$  is the product of the choice of span  $< i, j >$ , splitting point  $k$ , and children symbols  $< b, c >$ .
Local Prefix Tree For the non-terminal symbol  $V_{a}$  in a local prefix tree, the rules specify that all values of  $b$  and  $c$  are strictly fewer than the size of the left attached tree in the support structure, denoted as  $2^{l} - 1 < d = 2^{l}$ . Similarly, the variables  $i$ ,  $j$ , and  $k$  are also strictly limited to values fewer than  $d$ , since it is the max length of sub-strings that  $V_{a}$  can derive. Please note that  $l$  represents the max depth of the complete binary tree added as the left subtree to each node. Therefore the complexity filling all the  $S_{i,j}^{a}$  is  $\mathcal{O}(d^{5})$  for  $V_{a}$  and all  $V_{a}$  in Local Prefix Tree should be  $\mathcal{O}(d^{5}m)$  where  $m$  is the number of non-terminal symbols.
Main Chain For the non-terminal symbol  $V_{a}$  that serves as a main chain, based on the rules of the model, the variable  $b$  has strictly fewer than  $d$  choices, while  $c$  has a maximum of  $m / d$  choices. In a right-heavy parse tree, the symbols in the main chain are responsible for deriving a complete suffix of the target sentence as they are the right offspring of the root. Consequently,  $j$  is always equal to  $n - 1$ , and variable  $i$  has  $n$  choices. However,  $k$  has only  $d$  choices since the length of derived sub-string of  $V_{b}$  is less than  $d$ . Considering the filling of all  $S_{i,j}^{a}$  entries for the main chain symbols, the time complexity is  $\mathcal{O}(mdn)$ . Summing up the complexities for all  $m / d$  main chain symbols, the total complexity is  $\mathcal{O}(m^2 n)$ .
Based on the above analysis, the overall time complexity of the CYK algorithm for RH-PCFG is  $\mathcal{O}(d^5 n + m^2 n)$ . Such complexity is within the acceptable range, while the training complexity of general form PCFG, such as Chomsky normal form [5], is  $\mathcal{O}(n^3 |\mathcal{R}|)$ . Detailed pseudocode for the training algorithm can be found in the Appendix A.2.
3.3.2 Glancing Training
Glancing training, as introduced in the work of Qian et al. [30], has been successfully applied in various NAT architectures [12, 14]. In the case of PCFG-NAT, we also incorporate glancing training. This training approach involves replacing a portion of the decoder inputs with the embeddings of the target sentence  $Y$ . To determine the ratio of decoder inputs to be replaced with target embeddings, we employ the masking strategy proposed by Qian et al. [30]. There is a mismatch between the decoder length  $m$  and target length  $n$ , so we have to derive an alignment between them to enable the glancing. To incorporate Glancing Training into PCFG-NAT, we adapt the approach by finding the parse tree  $T_{Y}^{*} = \arg \max_{T_{Y} \in \Gamma(Y)} P(Y, T_{Y}|X)$  with the highest probability that generates the target sentence  $Y$ , where  $\Gamma(Y)$  stands for all parse trees of  $Y$ .  $T_{Y}^{*}$  can be traced by repalacing the sum operation with max in the CYK algorithm. We provide the pseudocode in the Appendix A.3.
3.4 Inference
We propose a Viterbi [41] decoding framework for PCFG-NAT to find the optimal parse tree  $T^{*} = \arg \max_{T} P(T|X, L)$  under the constraint of target length  $L$ . For each non-terminal symbol  $V_{a}$ , we
maintain a record  $M_L^a$  that represents the highest probability among sub-strings of length  $L$  that can be derived by  $V_{a}$ . Similar to the CYK algorithm, we can compute  $M_L^a$  in a bottom-up manner using dynamic programming, considering all non-terminal symbols  $V_{a}$  and lengths  $L$ .
$$
M _ {L} ^ {a} = \max  _ {<   b, c > \in \mathbf {C h i l d} (a), o \in \mathcal {T}, k = 1.. L - 1} P \left(V _ {a} \rightarrow V _ {b} o V _ {c}\right) M _ {L - 1 - k} ^ {b} M _ {k} ^ {c} \tag {15}
$$
To recover the parse tree with the maximum probability given  $L$ , we need to keep track of the values of  $b$ ,  $c$ , and  $k$  that lead to  $M_L^a$  and start the construction from  $M_L^1$ . The detailed pseudocode is provided in the Appendix A.4. According to Shao et al. [37], we rerank the output sentences with probability and length factors and then output the sentence with the highest score.
Dataset We conduct our experiments on WMT14 English  $\leftrightarrow$  German (En-De, 4.5M sentence pairs), WMT17 Chinese  $\leftrightarrow$  English (Zh-En, 20M sentence pairs) and WMT16 English  $\leftrightarrow$  Romanian (En-Ro, 610k sentence pairs). We use the same preprocessed data and train/dev/test splits as Kasai et al. [18]. For all the datasets, we learn a joint BPE model [36] with 32K merge operations tookenize the data and share the vocabulary for source and target languages. The translation quality is evaluated with sacreBLEU [29] for WMT17 En-Zh and tokenized BLEU [28] for other benchmarks.
Models	Iter	WMT14	WMT17	WMT16	Average Gap	Speedup
En-De	De-En	Eh-Zh	Zh-En	En-Ro	Ro-En
Transformer [40]	N	27.66	31.59	34.89	23.89	34.26	33.87	0	1.0×
CMLM [8]	10	24.61	29.40	-	-	32.86	32.87	-	2.2×
SMART [10]	10	25.10	29.58	-	-	32.71	32.86	-	2.2×
DisCo [17]	≈4	25.64	-	-	-	-	32.25	-	3.5×
Imputer [34]	8	25.0	-	-	-	-	-	-	2.7×
CMLMC [15]	10	26.40	30.92	-	-	34.14	34.13	-	1.7×
Vanilla NAT [13]	1	11.79	16.27	18.92	8.69	19.93	24.71	14.31	15.3×
CTC [21]	1	17.68	19.80	26.84	12.23	-	-	-	14.6×
FlowSeq[25]	1	18.55	23.36	-	-	29.26	30.16	-	1.1×
AXE [9]	1	20.40	24.90	-	-	30.47	31.42	-	14.2×
OAXE [6]	1	22.4	26.8	-	-	-	-	-	14.2×
CTC + GLAT [30]	1	25.02	29.14	30.65	19.92	-	-	-	14.6×
DA-Transformer + Lookahead [14]	1	26.55	30.81	33.54	22.68	32.31*	32.73*	1.27	14.0×
DA-Transformer + Joint-Viterbi [37]	1	26.89	31.10	33.65	23.24	32.46*	32.84*	1.00	13.2×
FA-DAT [26]	1	27.47	31.44	34.49	24.22	-	-	0.41	13.2×
PCFG-NAT	1	27.02	31.29	33.60	23.40	32.72	33.07	0.84	12.6×
Table 2: Results on raw WMT14 En-De, WMT17 Zh-En and WMT16 En-Ro. 'Iter' means the number of decoding iterations, and  $N$  is the length of the target sentence. The speedup is evaluated on the WMT14 En-De test set with a batch size of 1. \* indicates the results of our re-implementation.