N-gram Prediction and Word Difference Representations for Language Modeling
DongNyeong Heo Daniela Noemi Rim Heeyoul Choi

[
    [
    September 9, 2024
=====================

§ ABSTRACT 
Causal language modeling (CLM) serves as the foundational framework underpinning remarkable successes of recent large language models (LLMs). Despite its success, the training approach for next word prediction poses a potential risk of causing the model to overly focus on local dependencies within a sentence. While prior studies have been introduced to predict future N words simultaneously, they were primarily applied to tasks such as masked language modeling (MLM) and neural machine translation (NMT). In this study, we introduce a simple N-gram prediction framework for the CLM task. Moreover, we introduce word difference representation (WDR) as a surrogate and contextualized target representation during model training on the basis of N-gram prediction framework. To further enhance the quality of next word prediction, we propose an ensemble method that incorporates the future N words' prediction results. Empirical evaluations across multiple benchmark datasets encompassing CLM and NMT tasks demonstrate the significant advantages of our proposed methods over the conventional CLM. 






§ INTRODUCTION

With the remarkable advancements in deep learning techniques, neural language modeling has become a central component in modern natural language processing (NLP) tasks, such as natural language understanding (NLU), neural machine translation (NMT) and question answering. Among the approaches to language modeling, causal language modeling (CLM), which predicts the next word given the previous words, is a widely employed language modeling framework. For example, prominent large language models (LLMs) like GPT-2 <cit.> and GPT-3 <cit.> rely on CLM as their primary training framework. Despite their successful applications, the prevalent next word prediction manner can inadverently lead models to overfit to local dependencies rather than capturing long-term dependencies between words. This tendency arises from some phrases or paired words that have strong dependencies with each other, such as "Barack Obama" and "Harry Potter" <cit.>. 





A way of mitigating this problem involves predicting not solely the next word but also subsequent words in later time-steps such as N-gram prediction. Researchers <cit.> have adopted this N-gram prediction methodology for the masked language modeling (MLM) during the pre-training phase of LLMs <cit.>. Similar approaches have been applied to the NMT task <cit.>. However, these methods often require significant modifications to the model architecture, a different loss function than the conventional cross-entropy loss, or an expansion of the vocabulary for N-grams.


 

This paper introduces a novel N-gram prediction framework designed specifically for CLM and proposes innovative methods aimed at fortifying this framework. The contributions of this work can be summarized as follows.
(1) A simple N-gram prediction for CLM: we propose a simple N-gram prediction integrated to existing CLM models. Except for an additional multi-layer perceptron (MLP) layer, our method does not require other modifications to model architecture, loss function, and vocabulary.
(2) Word difference representation: we propose to use the embedding vectors' difference between contiguous words, termed word difference representation (WDR), as a surrogate representation for individual words. Departing from the conventional approaches that employing a fixed word embedding as target representation, we provide diverse WDR as target representations in accordance with context. We discovered this method can vary backpropagated gradient during training so that it can enhance generalizability. The algorithmic reversibility of WDR preserves the feasibility of the above simple N-gram prediction method.
(3) An ensemble method suitable for the CLM task: we propose an ensemble method designed to refine the next word prediction by leveraging other multiple N predictions from the N-gram prediction.

Our preliminary and primary experimental results, conducted several CLM benchmark datasets, highlight the gradual improvements in perplexity achieved by our proposed simple N-gram framework, the WDR diverse target representations, and ensemble method when compared to several baseline models. Our qualitative analysis focusing on gradient elucidates the advantage of the WDR method from the perspective of optimization generalizability. In addition to the main CLM task, we demonstrate the applicability and advantages of our proposed approaches to the NMT task, which is a conditional form of the CLM task.





§ BACKGROUND: CONVENTIONAL CLM


Since the work of <cit.>, neural network-based language modeling has been developed and become mainstream in language modeling. As background knowledge, we describe the conventional training framework of CLM (the next word prediction) in this section.

A sentence consists with words, X={x_1, x_2, ⋯, x_T}, x ∈𝒱, where T means the sequence length of the sentence and 𝒱 is the vocabulary set. Conventional CLM computes the likelihood of a word conditioned on its preceding words in the sentence, p(x_t|x_<t). For processing, words are mapped to embedding vectors <cit.>, and the encoded hidden state at time-step t is formulated as follows:

    𝐡_t=Enc_θ({𝐱^e_1,𝐱^e_2, ⋯, 𝐱^e_t-1}) ∈ℝ^d,

where 𝐱^e_t ∈ℝ^d means the embedded vector of x_t. Enc_θ is an encoder model with its parameter set θ. d is the dimension of the encoded hidden state and the embedding vector spaces. Recently, most language models use Transformer <cit.> as their encoder architecture. After encoding, the encoded hidden state is linearly transformed to a logit value of each word in a vocabulary set 𝒱. Finally, the likelihood of the predicted word is formulated as follows:

    p(x̂_t|x_<t;θ)    = softmax(𝐱̂^l_t), 
    𝐱̂^l_t    = 𝐖^l𝐡_t=𝐖^l𝐱̂^e,l_t,

where 𝐖^l ∈ℝ^|𝒱| × d is the weight matrix of the logit layer. 

To help the understanding of our idea, we note that a parameter vector of logit layer's weight is another word embedding set that is mapped to the target word, that is 𝐖^l=[𝐱^e,l_1,𝐱^e,l_2,⋯,𝐱^e,l_|𝒱|]^⊤. In this point of view, the encoded hidden state, 𝐡_t, is the predicted word embedding vector of the logit layer, 𝐱̂^e,l_t. Then, the inner product between 𝐖^l and 𝐱̂^e,l_t outputs the predicted score of each embedding that indicates how the predicted word embedding is similar to the logit layer's original word embedding.

Finally, the model learns to minimize the negative log-likelihood (NLL) loss as follows:

    ℒ(X,θ)=-∑_t=1^Tlog p(x̂_t=x_t|x_<t;θ).

This loss becomes the minimum when the model exactly predicts the logit layer's embedding of the target word, that is 𝐱̂^e,l_t=𝐱^e,l_t. This process is illustrated in Fig.<ref>(a).



§ PROPOSED METHODS

In this section, we propose three ideas: (1) a simple N-gram CLM, (2) word difference representation N-gram CLM, and (3) an ensemble method over N-gram predictions. 



 §.§ Simple N-gram CLM


First, we propose a simple N-gram prediction on the conventional framework of CLM. The core idea is adding an MLP layer to predict a future word given the same hidden state of the conventional CLM. This process is formulated as follows:

    𝐱̂^e,l_t+n=MLP^n(𝐡_t).

For instance, assuming N is 3, two MLP layers, MLP^1 and MLP^2, are employed and predict 𝐱̂^e,l_t+1 and 𝐱̂^e,l_t+2, respectively, as shown in Fig.<ref>(b).

The limited capability of the MLP layer to learn an effective function from a large and complicated dataset may regularize the main encoder, Enc_θ, to encode a simultaneously informative hidden state for all N-gram predictions. This regularization might be beneficial to prevent the model to overly focus on local dependencies.

We compute the likelihoods of the future target words, p(x̂_t+1|x_<t;θ) and p(x̂_t+2|x_<t;θ) in the above example, following each logit layer and the softmax function. Instead of using individual logit layers for each future word prediction, we share the parameters of all logit layers, including the conventional CLM model's logit layer. Therefore, this approach increases just a small amount of parameters for each additional MLP layer. Furthermore, it re-uses the original (unigram) vocabulary set for the future word prediction, not an additional large vocabulary set of N-grams.

The loss for n-th future word prediction is as follows:

    ℒ_n(X,θ)=-∑^T-n_t=1log p(x̂_t+n= x_t+n|x_<t;θ).

As like Eq.(<ref>), this loss becomes minimum when the model exactly predicts the future target word's embedding, i.e., 𝐱̂^e,l_t+n=𝐱^e,l_t+n. The total loss for the training of this simple N-gram CLM model is the mixture of Eq.(<ref>) and Eq.(<ref>) as follows:

    ℒ^tot_N(X,θ)=1/2ℒ(X,θ)+1/2(N-1)∑^N-1_i=1ℒ_i(X,θ).

Notably, we do not equally take the average of the original loss, Eq.(<ref>), with other losses, since the next word typically has stronger dependencies with the preceding words than other future words. In other words, averaging the entire set of loss terms together might introduce excessive regularization.






 §.§ Word Difference Representation (WDR) N-gram CLM

To use a more informative target than simple N-gram CLM, we introduce the idea of WDR which is a contextualized surrogate representation of words within a sentence. Basically, it is based on a form of word embedding compositions: the difference vector, 𝐱^e_t+1-𝐱^e_t. Since <cit.> demonstrated that arithmetic compositions of learned word embedding can convey semantic meanings, many researches have explored the word embedding compositionality <cit.>. Their studies utilized composed word embeddings as inputs to models, instead of original word embeddings, showcasing their advantages across various NLP tasks.

Unlike the prior research, we provide WDR to the model as the target to predict, rather than utilizing it as input. The difference vector of contiguous words offers a different representation for the word depending on its adjacent words. Therefore, by leveraging WDR as the target, we expect the model can learn more diverse targets than previous works. Generating WDR is simple repetition of vector subtractions which is computationally cheap and easy to parallelize, so it does not impose a high computational cost. Moreover, generating WDR is reversible, so that original embedding vectors can be reconstructed from WDR. This property facilitates the development of WDR-based N-gram CLM integrating the same framework of the simple N-gram CLM without a significant modification. Detailed explanations elucidating these advantages are provided in the subsequent sections.




  §.§.§ Definition of n-level WDR


As we briefly mentioned above, we use the difference of contiguous embedding vectors as the base of WDR. Given an embedding vector sequence {𝐱^e_1,𝐱^e_2,⋯,𝐱^e_T}, the 1-level WDR at the time-step t is defined as follows:

    Δ_1𝐱^e_t   =
            𝐱^e_t+1-𝐱^e_t    if  1≤ t < T, 
    𝐱^e_T    if  t = T.

In an inductive manner, the n-level WDR at the time-step t when n>1 is defined as follows:

    Δ_n𝐱^e_t   =
            Δ_n-1𝐱^e_t+1-Δ_n-1𝐱^e_t    if  1≤ t < T, 
    Δ_n-1𝐱^e_T=𝐱^e_T    if  t = T.


As an alternative of the above n-level WDR definition, we explored the opposite direction to subtract the contiguous vectors, that is Δ_n-1𝐱^e_t-Δ_n-1𝐱^e_t+1. In our internal empirical studies, we discovered the alternative design achieved similar performances. Therefore, we follow the design of Eq. <ref> throughout this paper.

Based on the definitions of Eqs. <ref> and <ref>, the n-level WDR can be represented by the composition of original word embeddings. For example, the 2 and 3-level WDRs at time-step t can be represented as follows: Δ_2𝐱^e_t=𝐱^e_t+2-2𝐱^e_t+1+𝐱^e_t and Δ_3𝐱^e_t=𝐱^e_t+3-3𝐱^e_t+2+3𝐱^e_t+1-𝐱^e_t, respectively. With this manner, we can derive the formulation of n-level WDR as follows:

    Δ_n𝐱^e_t=∑^n_i=0ni(-1)^i𝐱^e_t+(n-i),

where ni=n!/(n-i)!i! is the binomial coefficient. This equation holds for every positive integer of n 
and for every time-step t when t ≤ T-n.

See Appendix <ref> for a proof of this equation.

As we mentioned earlier, n-level WDR is reversible to the original word embedding. For the 1-level WDR, 𝐱^e_t+1 can be reconstructed by adding 𝐱^e_t to Δ_1𝐱^e_t. Likewise, 𝐱^e_t+n can be reconstructed by adding -∑^n_i=1ni(-1)^i𝐱^e_t+(n-i) to Δ_n𝐱^e_t (note that the first term of the right-hand side of Eq.(<ref>) is 𝐱^e_t+n). For simplicity, we use a new notation for the conjugate term that reconstructs the original embedding by addition to the n-level WDR as follows:

    Δ^r_n𝐱^e_t=-∑^n_i=1ni(-1)^i𝐱^e_t+(n-i),



    Δ^r_n𝐱^e_t   =
            𝐱^e_t    if  n = 1, 
    
                -∑^n_i=1ni(-1)^i𝐱^e_t+(n-i)   if  n > 1.


This leads to Δ_n𝐱^e_t+Δ^r_n𝐱^e_t=𝐱^e_t+n. The conjugate term for reconstruction, Δ^r_n𝐱^e_t, can be obtained by Eq.(<ref>) or  iterative operations of Eq.(<ref>).





  §.§.§ Training of WDR N-gram CLM


We develop the WDR-based N-gram CLM from the framework of simple N-gram CLM.
To achieve the mentioned goal that providing the WDR as the target of the model, we apply the definitions and derivations in Sec.<ref> to the logit layer's embeddings.

Following the idea of the simple N-gram CLM described in Sec.<ref>, we employ MLP layers for predictions of N-gram. However, in WDR N-gram CLM, the MLP^n layer outputs Δ_n𝐱̂^e,l_t instead of 𝐱̂^e,l_t+n. Then we produce its corresponding conjugate term, Δ^r_n𝐱^e,l_t, based on the logit layer's embedding matrix. Adding those two, Δ_n𝐱̂^e,l_t+Δ^r_n𝐱^e,l_t, yields 𝐱̂^e,l_t+n as in the simple N-gram CLM. Then, we take the same processes of the logit, likelihood, and loss computations as in the simple N-gram CLM.

An essential design of this framework is detachment of the produced conjugate term, Δ^r_n𝐱^e,l_t, from the backpropagation process. Absence of this detachment might lead the model to adjust the logit layer's weight matrix in a distorted manner, because the input of the logit layer is recursively produced from itself. 

In WDR N-gram CLM, the minimum value of NLL loss of x_t+n prediction, Eq.(<ref>), is achieved when 𝐱̂^e,l_t+n=𝐱^e,l_t+n, which is Δ_n𝐱̂^e,l_t+Δ^r_n𝐱^e,l_t=Δ_n𝐱^e,l_t+Δ^r_n𝐱^e,l_t based on the equation led by Eq.(<ref>). Because the conjugate term, Δ^r_n𝐱^e,l_t, is detached, the model would learn to predict Δ_n𝐱^e,l_t, which is true n-level WDR. 
In other words, WDR N-gram CLM learns to predict composed word embeddings, offering diverse and contextualized target representations, even for the same target word. The entire process of WDR trigram CLM example is illustrated in Fig.<ref>(c). 

 



  §.§.§ How Diverse Are WDR-based Target Representations?


In order to gain a more profound understanding of WDR as target representations, we explored how WDR would diversify target representations compared to the conventional CLM or the simple N-gram CLM. As we mentioned in Sec.<ref> and Sec.<ref>, the conventional CLM and the simple N-gram CLM utilize the logit layer's embeddings as target representations to predict. To see the practical examples of these target representations, we collected 1,270 representations from the logit layer's embedding matrix of the pre-trained conventional CLM model (`TF' in the preliminary experiment, Sec.<ref>). The 1,270 representations correspond to all the tokens of randomly selected 10 sentences from the Penn TreeBank (PTB) <cit.> testset. Also, we computed 1 and 2-level WDRs with the collected embeddings, and added them to the collection, resulting in 3,810 representations in total. Finally, we reduced the dimension of the total collection to 2-dimension with t-SNE algorithm <cit.>. 

Fig.<ref> shows the collected representations in a 2-dimensional space. The first plot illustrates the original embeddings, 𝐱^e,l. Note that the representations of frequent words, such as `to' may be included more times than other words in the collection. We interpret that this is the reason why t-SNE places frequent words (e.g., `in', `to', and `the') distant from other less frequent words to resemble the non-uniform distribution of the collection. 
On the other hand, the 1-level WDR representations, Δ_1𝐱^e,l, look more diverse compared to the original embeddings as in the second plot. For example, by composing adjacent words such as `want', `unable', `returned', into the frequent word `to', it diversifies the embedding representations according to its previous word as in the third plot which is zoomed in. The 2-level WDR looks more diverse even compared to 1-level WDR as in the last plot. Based on this analysis, we expect WDR N-gram CLM to give more diverse target representations than other methods, such as conventional CLM and the simple N-gram CLM.



 §.§ Ensemble Method to Refine the Next Word Prediction Leveraging N-gram Predictions


We propose a new ensemble method to incorporate the N-gram predictions into the process of the next word prediction. The encoder model, such as Transformer, outputs {𝐡_2,𝐡_3,⋯,𝐡_t} given the embedded input sentence {𝐱^e_1,𝐱^e_2,⋯,𝐱^e_t-1}. The encoded hidden state 𝐡_i represents the computed hidden state given the inputs up to time-steps (i-1). At testing, in addition to the predicted embedding 𝐱̂^e,l_t from the conventional CLM, MLP^n layer of N-gram CLM can estimate the target word for time t given 𝐡_t-n. Therefore, we can get N predicted embeddings for the current time-step. We ensemble these predicted embeddings just before the logit layer using the following formulation:

    𝐱̂^e,l_t,ens=(1-λ)𝐱̂^e,l_t+λ/N-1∑^N-1_i=1MLP^i(𝐡_t-i),

where


    𝐱̂^e,l_t,ens=(1-λ)𝐱̂^e,l_t+λ[1/N-1∑^N-1_i=1𝐱̂^e,l_t,i],

where 𝐱̂^e,l_t,i is the output of MLP^i layer given 𝐡_t-i.

λ is a scalar value between 0 and 1. It controls the influences of future word predictions (but derived from past time-steps) on the current word prediction. Similar to the rationale behind the dominance of the original NLL loss in its total loss formulation, Eq.(<ref>), we do not equally average the original predicted embedding with others. In the case of WDR-based N-gram CLM, we ensemble MLP^i(𝐡_t-i)+Δ^r_i𝐱^e,l_t-i=𝐱̂^e,l_t in the summation part in Eq.(<ref>).

After this ensemble computation, we input it to the logit layer and compute the next word's likelihood. At testing, this ensemble likelihood result is used to compute perplexity (PPL) in CLM tasks or serving as candidate scores for beam search in NMT tasks.



§ EXPERIMENTS AND RESULTS

To assess the performances of our proposed methods, we conducted CLM and NMT experiments on multiple benchmark datasets. 



 §.§ Causal Language Modeling (CLM)

For the CLM task, we executed two experiments: preliminary and primary. The preliminary experiment was dedicated to monitor the dynamics of two hyperparameters: N and λ toward the performance. In contrast, we only report the results of the best hyperparameters in the primary experiment's demonstration.






  §.§.§ Data Description

PTB (-, 0.9M tokens, 10K vocabulary), WikiText-2 (W2, 2M tokens, 33K vocabulary), Text8 (T8, 15M tokens, 254K vocabulary), and WikiText-103 (W103, 103M tokens, 268K vocabulary) <cit.>. To ensure standardization and transparency in our data-related processes (e.g., download, tokenization, vocabulary, and train/valid/testsets splitting), we relied on open sources. Specifically, the W2 and T8 datasets were sourced from the GitHub repository[https://github.com/chakki-works/chazutsu], while the PTB and W103 datasets were sourced from the Tensorized Transformer <cit.>'s GitHub repository[https://github.com/szhangtju/The-compression-of-Transformer]. In the primary experiment, we used the whole datasets, whereas the preliminary experiment was conducted solely on the PTB dataset.





  §.§.§ Models and Training


For the baseline model of the preliminary experiment, we implemented Transformer (TF) encoder-based CLM. The total number of parameters of the TF baseline is 12M, and our proposed simple and WDR methods increase only 0.1M parameters per an additional MLP layer (note that the logit layer's parameters are all shared). The details of model architecture and training method for the preliminary experiment are described in Table <ref> (in Appendix <ref>) in the column of `Small Enc. TF CLM'.

For the baseline models of the primary experiment, we trained the two baseline models that are advanced ones based on TF: tensorized transformer (TT) <cit.> and Reformer (RF)[https://github.com/lucidrains/reformer-pytorch] <cit.>. We mostly followed their reported configurations, except some minor changes such as the number of tokens in a mini-batch and learning rates. The details of these changes for each dataset are described in Table <ref> (in Appendix <ref>). As a result, the total numbers of parameters of (TT, RF) models according to datasets are (6.7M, 15.3M) for PTB and W2, (82.4M, 236.6M) for T8 and W103, respectively. Our proposed simple and WDR methods increase the number of parameters by 0.1M and 0.5M, respectively, per an additional MLP layer regardless of the type of dataset.

On top of the baseline models, we applied our proposed method, and we call them `TF+Sim', `TF+WDR', `TT+Sim', `TT+WDR', `RF+Sim', and `RF+WDR'. We varied N from 2 to 4 and λ from 0.0 to 0.6 for every experiment of our proposed methods. In the demonstration of the primary experiment results, we report the result of the best hyperparameter setting of each model. These settings are reported in the `CLM Task' column of Table <ref> (in Appendix <ref>).






  §.§.§ Preliminary Experimental Results


Table <ref> presents the outcomes of the preliminary experiments. We trained the model of each configuration five times with different seeds, and we report the average PPL scores. Both `TF+Sim' and `TF+WDR' surpass the performances of the conventional CLM baseline. This observation aligns with findings from previous studies on other tasks <cit.>. The ensemble method consistently improves performance compared to the non-ensemble ones (where λ=0.0). It usually achieves the best scores at λ=0.4 for both the `TF+Sim' and `TF+WDR' models. Also, we observed that the `TF+WDR' model maintains strong performance even at λ=0.6, while the `TF+Sim' model does not. This implies that `TF+WDR' generally generates more accurate predictions for future words. Moreover, `TF+WDR' tends to outperform their `TF+Sim' counterparts in each setting. These findings collectively suggest that the WDR training approach offers benefits over N-gram prediction methodologies.






  §.§.§ Gradient Diversity Analysis

As an additional exploration of the advantages of WDR, we checked the connection between the diverse target representations and its benefit during training. Given the evidence in Sec.<ref> that WDR gives more diverse target representations compared to other CLMs, it is plausible to guess the backpropagated gradients are also diverse. To quantify this property, we measured `gradient diversity (GD)' <cit.> which is formulated as follows:

    GD(𝒟,θ)   =∑^|𝒟|_i=1||g_i||^2_2/||∑^|𝒟|_i=1g_i||^2_2,
       = ∑^|𝒟|_i=1||g_i||^2_2/∑^|𝒟|_i=1||g_i||^2_2+∑_i j⟨ g_i,g_j ⟩, 
    
        g_i   =∇_θℒ^tot_N(X_i,θ),

where 𝒟={X_1,X_2, ⋯,X_|𝒟|} is a mini-batch, ||·||^2_2 is the squared L^2 norm operation, ⟨·,·⟩ is the inner product operation, and ∇_θ is gradient operator with respect to θ. This metric is large when the inner product terms in denominator are small, which means the gradients are different from each other. 

We measured GD of the `TF+Sim N=4' and `TF+WDR N=4' models in Table <ref> during training. The GDs over epochs are presented in Fig.<ref>. `TF+WDR N=4' usually has higher GD than `TF+Sim N=4'. As the stochastic property of stochastic gradient descent is known for noisy gradient which enhances generalizability compared to full-batch gradient descent <cit.>, higher GD may offer similar advantages due to the stochastic property. Given this understanding, we believe WDR-based training could be beneficial for improving generalization.












  §.§.§ Primary Experimental Results

Table <ref> presents the entire results of the primary experiments (6 models on 4 datasets). Results show that, with the exception of TT-based models on W2, our proposed N-gram CLMs consistently either match or surpass the baseline CLMs, even without the ensemble method. Remarkably, WDR N-gram CLMs generally improve performance on top of the simple N-gram CLMs. Upon applying our proposed ensemble method, they generally exhibit improvements over their non-ensemble counterparts, except the models trained on W103. Notably, the effect of ensemble method is relatively significant in smaller datasets (PTB and W2) in contrast to larger datasets (T8 and W103). Based on these results, we argue that our proposed methods have actual advantages on various models and datasets for the CLM task.







 §.§ Neural Machine Translation




  §.§.§ Data Description

Since NMT includes language modeling as a part of the decoder, we view the NMT could be an appropriate additional experimental task to demonstrate the effectiveness of our proposed approach in addition to the main CLM tasks.  We conducted NMT experiments on several datasets: `IWSLT14 English-German'(En-De, 160K training pairs) <cit.>, `WMT14 English-German'(En-De, 3.9M training pairs) <cit.>, and `WMT18 English-Turkish' (En-Tr, 207K training pairs) <cit.>. We used the same preprocessing, tokenization and subword byte-pair encoding methods with <cit.>. We used 10K, 10K, 32K most frequents subwords to organize vocabularies for datasets, respectively.







  §.§.§ Models and Training


As a baseline, we used our implementation of Transformer (TF) <cit.> in the encoder-decoder architecture. We used the small Transformer for the `IWSLT14 En-De' and `WMT18 En-Tr' datasets,

and the base Transformer for the `WMT14 En-De' dataset. The total number of parameters of small and base TF baselines are 32M and 77M, respectively. We applied our simple and WDR N-gram CLM methods onto the decoder parts of the baselines, `TF+Sim' and `TF+WDR'. Each additional MLP layer in our simple and WDR methods increases the number of parameters by around 0.5M. Information about the models and how TF models are optimized can be found in the columns labeled 'Small Enc-Dec TF NMT' and 'Base Enc-Dec TF NMT' in Table  <ref>. Also, the hyperparameters (N and λ) for `TF+Sim' and `TF+WDR' are described in the `NMT Task' column of Table <ref> (in Appendix <ref>).



As a more closely related baseline, bag-of-words (BOW) NMT was proposed to predict the whole words in the context of the original NMT task <cit.>. However, their approach was not applied to the TF architecture, and they evaluated the model only on the English-Chinese translation dataset of NIST. To ensure a fair comparison, we re-implemented BOW NMT based on our TF architecture and compared with our proposed method. Following their prescribed approach, we integrated the computed loss of whole words prediction into the original loss. 





  §.§.§ BLEU Results

Table <ref> presents the experiment results of the models on each testset with SacreBLEU <cit.> as the evaluation metric. Our proposed `TF+Sim' and `TF+WDR' models exhibit usually enhanced performances compared to the `TF' and `BOW NMT' baselines. `TF+WDR' always outperforms its counterpart of `TF+Sim'. Notably, the integration of the ensemble method from both of `TF+Sim' and `TF+WDR' further increases performances. Specifically, we note that `TF+WDR' with ensemble method improved performances by 0.7 1.5 BLEU scores compared to `TF' baseline on the both translation directions of `IWSLT14 En-De', and German-to-English translations of `WMT14 En-De' testsets. 

To explain why N-gram prediction approaches are more effective for German-to-English translation compared to English-to-German translation in `IWSLT14 En-De' and `WMT14 En-De' experiments, we hypothesize that the difference in word diversity between the two languages plays a role. We analyzed the `WMT14 En-De' training dataset (subword-level tokenized) and found that English has around 33.6K unique unigrams and 6.7M unique bigrams, while German has around 34.9K unique unigrams and 9.3M unique bigrams. This suggests that German-to-English translation might have simpler local dependencies to learn compared to English-to-German translation due to the lower number of unique bigrams. Considering simple local dependencies might lead to the over-fitting problem, we believe that this is a potential reason why N-gram prediction approaches, which can help mitigate over-fitting to local dependencies, are more effective for German-to-English translation. 















§ CONCLUSION

In this work, we have constructed an advanced N-gram prediction framework tailored specifically to causal language modeling. 
In addition to the construction of this framework, our work includes the introduction of new strategies for providing diverse target representations and an ensemble method over the predicted N words. Extensive experiments on language modeling and neural machine translation have confirmed the practical benefits of the proposed method.





§ LIMITATIONS

Given the demonstrated performance improvements of the WDR-based N-gram CLM, we tried to apply the WDR method to other tasks beyond CLM, such as the MLM task. In addition to the standard loss function of MLM, which involves predicting the masked word <cit.>, we added new loss terms to predict n-level WDR target representations of the masked position. For this experiment, we utilized the CrammedBERT model <cit.>, a streamlined variant of BERT that facilitates faster pre-training while maintaining competitive performance on the GLUE benchmark. We integrated the WDR approach into this model and conducted a comparative analysis with the original CrammedBERT configuration. Further experimental details are provided in Appendix <ref>.

Table <ref> (in Appendix <ref>) presents the results of our experiments comparing CrammedBERT and the applications of WDR models on the GLUE test set. While the application of 2-level WDR resulted in a 1.0 point increase in the average GLUE score, the performance benefits of the WDR method is less consistent across individual sub-tasks compared to the benefits observed in the CLM tasks. We attribute this result to the fundamental difference between the CLM and MLM tasks. Specifically, in MLM, when the WDR method combines the masked word embedding with the embeddings of the next words, such information is already provided as input. This partial visibility of the target representation might lead to an unexpected optimization behavior, such as the model disproportionately focusing on the right-side (future) context which is incorporated in the target, rather than considering the entire context.

Since there are prior works for N-gram prediction within the MLM framework <cit.>, we believe we can apply the WDR method to the prior works by combining the only masked words when WDR is calculated to solve the aforementioned issue. We expect that the high gradient diversity characteristic of the WDR method may offer additional benefits to the prior MLM framework.




§ ACKNOWLEDGEMENTS











§ APPENDIX





 §.§ Proof of Eq.(<ref>)


We provide a proof of Eq.(<ref>) with the induction method. To avoid confusion, we temporarily change the notation of Δ_n𝐱^e_t in conjecture Eq.(<ref>) to Δ̂_n𝐱^e_t until it is proved. Based on the definitions of the 1 and n-level WDR, Eq.(<ref>) and Eq.(<ref>), we can verify the initial condition, that is n=1, holds as follows:

    Δ_1𝐱^e_t   =𝐱^e_t+1-𝐱^e_t
       =10(-1)^0𝐱^e_t+1+11(-1)^1𝐱^e_t
       = ∑_i=0^11i(-1)^i𝐱^e_t+(1-i)
       =Δ̂_1𝐱^e_t.

Therefore, the conjecture holds for the initial condition. Then, by following the induction method, we assume the conjecture at n-level is true, that is Δ̂_n𝐱^e_t=Δ_n𝐱^e_t. Then, the (n+1)-level WDR from the definition Eq.(<ref>) is derived to Δ_n+1𝐱^e_t=Δ_n𝐱^e_t+1-Δ_n𝐱^e_t=Δ̂_n𝐱^e_t+1-Δ̂_n𝐱^e_t. Each term is derived as follows:

    Δ̂_n𝐱^e_t+1   =n0(-1)^0𝐱^e_t+n+1+n1(-1)^1𝐱^e_t+n+ 
       ⋯ +nn-1(-1)^n-1𝐱^e_t+2+nn(-1)^n𝐱^e_t+1, 
    
        -Δ̂_n𝐱^e_t   =n0(-1)^1𝐱^e_t+n+n1(-1)^2𝐱^e_t+n-1+ 
       ⋯ +nn-1(-1)^n𝐱^e_t+1+nn(-1)^n+1𝐱^e_t, 
    Δ̂_n𝐱^e_t+1-Δ̂_n𝐱^e_t   =n0(-1)^0𝐱^e_t+n+1+(n0+n1)(-1)^1𝐱^e_t+n+ 
       ⋯+(nn-1+nn)(-1)^n𝐱^e_t+1+nn(-1)^n+1𝐱^e_t
       =n+10(-1)^0𝐱^e_t+n+1+n+11(-1)^1𝐱^e_t+n+ 
       ⋯+n+1n(-1)^n𝐱^e_t+1+n+1n+1(-1)^n+1𝐱^e_t
       =∑^n+1_i=0n+1i(-1)^i𝐱^e_t+(n+1-i)
       =Δ̂_n+1𝐱^e_t.

Note that the binomial coefficient, ni, is the n-th row and i-th value of Pascal's triangle, and it satisfies ni-1+ni=n+1i. Based on this outcome, the conjecture holds for (n+1)-level if the n-level is true. Therefore, the conjecture is proved.




 §.§ Experiment Details



















We trained the models described in Sec. <ref> and Sec. <ref> following the configurations described in Table <ref> for Transformer-based models, `TF', and the configurations reported in the previous works' papers <cit.> with several changes as described in Table  <ref> for the primary CLM baselines, `TT' and `RF'. For Transformer-based models' experiments, we saved the best checkpoint based on the validation results. We early stopped the training whenever the model does not beat its previous best performance for the `Patience' times on the validation <cit.>. For the primary CLM baselines, we followed the pre-defined total training iterations. Table <ref> describes the specific configurations, such as N and λ, we used for our proposed N-gram CLMs, simple-based and WDR-based. 

About the information of our computational environment, we used a single NVIDIA RTX3090 GPU for the large CLM datasets, such as T8 and W103, and a GTX1080Ti GPU for the small CLM datasets, such as PTB and W2. On average, they took 1 day and 3 hours, respectively, for training. We used 4x NVIDIA RTX3090 GPUs for the large NMT datasets, such as WMT14 English-German, and 2x GTX1080Ti GPUs for the small NMT datasets, such as IWSLT14 English-German and WMT18 English-Turkish. On average, they took 3 days for training.





 §.§ Masked Language Modeling Experiment


We adhered to the environmental settings established by CrammedBERT <cit.> for all aspects of our study, including dataset preprocessing, model configurations, pre-training, fine-tuning procedures, and evaluations. Comprehensive details of these settings can be found in the associated GitHub repository[https://github.com/JonasGeiping/cramming]. Building on the CrammedBERT architecture, we apply the WDR method that is analogous to the method conducted in our WDR-based N-gram CLM experiment. Specifically, we utilized N additional MLP layers designed to predict n-level WDRs alongside the original word embedding at the masked position. These n-level WDRs are calculated by composing the next words of the masked word. The final loss is computed as the average of the original loss and the additional losses derived from the WDR method, with the original and additional losses being averaged unequally, as described in Section <ref>.

Table <ref> presents the experimental results for CrammedBERT and our proposed models, evaluated on the GLUE test set following fine-tuning. We varied the number of grams, N, from 1 to 3. The results indicate that the application of 2-level WDR yields an increase of 1.0 point in the average GLUE score. However, the performance improvements across individual sub-tasks are not consistently superior; in some cases, they were similar to or worse than the baseline.







