EMO: EARTH MOVER DISTANCE OPTIMIZATION FOR AUTO-REGRESSIVE LANGUAGE MODELING

Dear reviewer, we sincerely appreciate your time reading the paper. Here we do our best to address your concerns.

• MLE/EMO-induced distribution: For per-token loss, both MLE and EMO are minimized theoretically when the model distribution $Q_{\theta}$ is one-hot. However, in practical training which involves iterative stochastic gradient descent on human text data, both MLE and EMO optimize model distribution to become more human-like. Moreover, as stated in Section 4.1.1 Training Details, the final loss of EMO is the additive combination of MLE and DEMD. The difference between the final distributions induced by MLE and EMO emerge as the aggregated effect of per-token gradient update throughout training, where EMO explicit penalizes low-quality tokens and takes more account into plausible alternatives. We demonstrate this quantitatively and qualitatively

• Quantitative difference: We show the differences between MLE and EMO induced distributions by measuring the forward cross-entropy $\text{CE}^{\text{forward}}=-\sum{P\log{Q_{\theta}}}$, reverse cross-entropy $\text{CE}^{\text{reverse}}=-\sum{Q_{\theta}\log{P}}$, and their harmonic mean against a reference distribution $P$ provided by a larger language model GPT-Neo. The GPT-2 results of $Q_{\theta}^{\text{MLE}}$, $Q_{\theta}^{\text{EMO}}$, and a smoothed one-hot distribution are shown in the table below. As can be seen, both $Q_{\theta}^{\text{EMO}}$ and $Q_{\theta}^{\text{MLE}}$ deviate far from one-hot and $Q_{\theta}^{\text{EMO}}$ is more similar to reference distribution by better harmonizing recall and precision. | | | WikiText-2 | | | WritingPrompt | | |:---:|:---:|:---:|:---:|:---:|:---:|:---:| | | $\text{CE}^{\text{reverse}}$↓ | $\text{CE}^{\text{forward}}$↓ | Harmonic Mean↓ | $\text{CE}^{\text{reverse}}$↓ | $\text{CE}^{\text{forward}}$↓ | Harmonic Mean↓ | | $Q^{\text{one-hot}}$ | 3.22 | 16.83 | 5.15 | 3.32 | 17.74 | 5.33 | | $Q_{\theta}^{\text{MLE}}$ | 4.50 | 4.33 | 4.22 | 4.25 | 3.92 | 3.99 | | $Q_{\theta}^{\text{EMO}}$ | 3.94 | 4.45 | 3.97 | 3.54 | 4.01 | 3.66 |

  We also report the Distinct N-gram of the sampled outputs from $Q_{\theta}^{\text{MLE}}$, $Q_{\theta}^{\text{EMO}}$, and reference text. As can be seen, $Q_{\theta}^{\text{EMO}}$ is also more similar to human text distribution in various aspects of the sampled outputs. | | | | WikiText2 | | | | | | WritingPrompt | | | |---|:---:|:---:|:---:|:---:|:---:|---|:---:|:---:|:---:|---|:---:| | | | | Dist$_{n}$ | | MAUVE | | | | | Dist$_{n}$ | | MAUVE | | | n=1 | n=2 | n=3 | n=4 | - | | | n=1 | n=2 | n=3 | n=4 | - | | $Q_{\theta}^{\text{MLE}}$ | 0.748 | 0.971 | 0.994 | 0.998 | 77.5 | | | 0.806 | 0.983 | 0.995 | 0.997 | 83.6 | | $Q_{\theta}^{\text{EMO}}$ | 0.698 | 0.955 | 0.990 | 0.997 | 87.5 | | | 0.762 | 0.972 | 0.993 | 0.996 | 87.4 | | Human | 0.687 | 0.958 | 0.988 | 0.996 | 100 | | | 0.764 | 0.974 | 0.993 | 0.997 | 100 |

• Qualitative difference: Here we show some predicted next token probability from $Q_{\theta}^{\text{MLE}}$ and $Q_{\theta}^{\text{EMO}}$. We show the top 5 predicted tokens and their probabilities.

  • Prefix: Donald Trump has multiple occupations, such as
  • $Q_{\theta}^{\text{MLE}}$: the(0.393), a(0.325), construction(0.100), executive(0.093), president(0.088),...
  • $Q_{\theta}^{\text{EMO}}$: a(0.366), the(0.353), president(0.100), executive(0.091), construction(0.091),...
  • Prefix: Elephant is the largest living
  • $Q_{\theta}^{\text{MLE}}$: animal(0.275),creature(0.231),elephant(0.215),mammal(0.197),ant(0.079),...
  • $Q_{\theta}^{\text{EMO}}$: animal(0.268),creature(0.257),mammal(0.216),elephant(0.196),rept(0.062),...

• Greedy/Beam search: The results of greedy/beam search of $Q_{\theta}^{\text{MLE}}$ and $Q_{\theta}^{\text{EMO}}$ on WikiText-2 test set are shown in the table below. More greedy decoding results of LLaMa2 are presented in our General Response. | | Method | MAUVE | BLEU | ROUGE-1 | |---|:---:|:---:|:---:|:---:| | Greedy Search | $Q_{\theta}^{\text{MLE}}$ | 7.10 | 21.64 | 38.8 | | | $Q_{\theta}^{\text{EMO}}$ | 11.0 | 22.27 | 38.0 | | Beam Search | $Q_{\theta}^{\text{MLE}}$ | 9.93 | 21.67 | 37.0 | | | $Q_{\theta}^{\text{EMO}}$ | 13.14 | 22.37 | 37.2 |

• Difference between EMO and word embedding-based loss: We would like to modestly point out that the word embedding-based loss described in the paper mentioned by the reviewer is indeed a weighted variant of MLE. In their method, the gradient of the ground-truth token's probability is weighted by the sum of its Euclidean distance to all tokens in the target vocabulary, .i.e, $\nabla Q_{\theta}(y_i)\sum_k d(E(v_k),E(v_i))$. In contrast, EMO involves the sum of gradients of all tokens' probability, weighted by their respective expected transport cost. We will cite this paper in the final version.

Reference:

[1] Meister et. al: On the Efficacy of Sampling Adapters. ACL 2023

Thanks for your reply.

1. I appreciate the clarification. I didn't notice that the final loss also includes MLE, so it turns out that the distribution of EMO is not as sharp as I expected. However, I disagree with the author's claim that "both MLE and EMO optimize model distribution to become more human-like." Based on my analysis, the expected EMO loss is minimized when the model distribution Q is a one-hot distribution that only outputs the token i that maximizes $\sum_{j=1}^{|V|}P(v_j)C(v_i, v_j)$. Consequently, the perplexity of EMO is higher than that of MLE, indicating a less human-like distribution. If DEMD were used solely as the loss function, I believe the resulting model distribution would be sharper and less human-like, leading to a significantly higher perplexity. A similar phenomenon can be observed in language GANs [1].

2. I maintain my point that the paper's motivation is flawed, and I suggest the author consider presenting the story of EMO from a different perspective.

3. I acknowledge my previous misunderstanding of the word embedding-based loss and appreciate the clarification provided.

I have a few questions:

1. Could you explain the concept of Q^{one-hot}?

2. In the WikiText-2 test set results for greedy/beam search, why are the MAUVE scores substantially lower than those obtained through unbiased sampling, as shown in Table 1?

[1] https://arxiv.org/pdf/1811.02549.pdf

Thanks for reading our responses. Regarding your questions, our answers are as follows:

• The concept of $Q^{\text{one-hot}}$ refers to the one-hot distribution constructed based on the ground-truth next token. The various divergence measures of $Q^{\text{one-hot}}$ we've reported in our response serve as a reference to indicate how far are $Q_{\theta}^{\text{EMO}}$) from such one-hot distribution. Compared to $Q_{\theta}^{\text{MLE}}$, $Q_{\theta}^{\text{EMO}}$ is closer to the reference distribution $P$ as measured by the harmonic mean of forward/reverse cross-entropy.
• This is because MAUVE score measures how similar the model distribution $Q_{\theta}$ is to the data distribution $P$. Concretely, MAUVE computes the area under an interpolated KL divergence curve. It rewards model distribution that balances recall($Q_{\theta}$ has sufficient cover of $P$'s mode) and precision($Q_{\theta}$ does not overestimates regions where $P$ is low). Two identical distributions will have a MAUVE score of 100. Both greedy and beam search are maximization-based decoding strategies that discard all possible outputs except the most possible one. This is equivalent to a high-precision but significantly low-recall output distribution. Therefore, the MAUVE scores of greedy and beam search are usually much smaller than sampling-based methods, e.g., unbiased/top-p/top-k/typical sampling. The same phenomenon was also observed in the original MAUVE paper as well as a prior work[1]'s Table 1.

In the following we would like to explicate more about your comments:

• We would like to kindly clarify that, in the context of auto-regressive language modeling(where we have a collection of text corpus for training), the mathematically lowest loss at each time step is achieved when the ground-truth next token's probability under $Q_{\theta}$ is 1, for both MLE and DEMD: (1) for MLE, -log1=0, and (2) for DEMD, 1-11cos(0)=0. Both MLE and EMO optimize model towards such solution, but through different trajectories, i.e., gradient signal. The gradient signal of EMO informs more about "bad" tokens and alternative "good" tokens.
• "higher perplexity": MLE has lower perplexity is expected because MLE's objective is to minimize perplexity($e^{-\log{Q_{\theta}(x)}}$), which is the so-called recall-prioritization[2]. EMO has higher perplexity because it cares more about diversity in the range of other "good" tokens.
• "less human-like distribution": According to an exhaustive array of quantitative metrics we've reported, i.e., MAUVE, Harmonic mean of forward/reverse cross-entropy, Distinct N-gram, and the alignment experiments upon LLaMa2, the EMO-induced distribution is more human-like. We would much appreciate if the reviewer could elaborate further why EMO is less human-like.

Please let us know you have further questions and we will try our best to solve your issues.

References:

[1] Su et.al: A Contrastive Framework for Neural Text Generation. NeurIPS 2022.

[2] Meister et. al: On the Efficacy of Sampling Adapters. ACL 2023

Dear reviewer, we sincerely appreciate your time reading the paper. Here we do our best endeavor to address your concerns:

• Choice of Baselines: Our primary rationale for selecting baselines is to compare methods that optimize different probability divergence measures between model distribution and data distribution. Among our baselines, MLE optimizes the forward cross-entropy, TaiLr focuses on the total variation distance and MixCE mixes forward and reverse cross-entropy. D2GPo did not propose to optimize new divergence measures but instead focused on a new target distribution. Moreover, we didn't find the publicly available official implementation of D2GPo. Therefore, we didn't include D2GPo as our baseline in our submitted draft. Our reimplementation of D2GPo shows similar performance to MLE.
• Choice of Cosine Distance: The reasons for using cosine distance between llm head embedding are three-fold:
  • The general quadratic form of DEMD using a $|V|\times|V|$ transport cost matrix $C$ is $O(|V|^2)$. By utilizing the cosine distance, we can decrease the computation complexity from $O(|V|^2)$ to $O(|V|H)$(Eq. 13->Eq. 14).
  • The output logits of LM are computed based on the dot-product between hidden states and LLM head embedding. The computation of cosine distance is also based on dot-product between normalized llm head embedding, exploiting the information stored in llm head. In fact, the transport cost matrix can also be computed based on the lm head of the model undergoing fine-tuning, which gets updated to incorporate the distribution shift from the pre-training corpus to task-specific training data. We term this variant of EMO as $\text{EMO}_{\text{adaptive}}$and conduct experiments to verify this and the results of language modeling are shown in the table below: | | WikiText-2 | WritingPrompts | |:---:|:---:|:---:| | MLE | 77.5 | 83.6 | | EMO | 87.5 | 87.4 | | EMO$_{\text{adaptive}}$ | 88.1 | 88.3 |
  • The main purpose of policy gradient methods such as REINFORCE and PPO is to enable differentiable learning towards potentially non-differentiable metrics. Moreover, the reward is usually a sequence-level delayed one. In contrast, the research objective of EMO focuses on exploring more effective probability divergence measures at the token level through the lens of Earth Mover Distance. In the LLM area, the former usually happens in the alignment stage while EMO is more suitable to be applied to (1) domain-specific fine-tuning; (2) continual pre-training and (3) instruction-tuning. We have included additional results on instruction tuning upon LLaMa2 in the general response. Please refer to it for more details.
• Perplexity: You are right, the evaluation is conducted over the sampled sentences instead of the test set. The $\text{PPL}^{\text{oracle}}$ reported in Table 2 is computed using the Oracle data generator on the model-generated samples. A lower $\text{PPL}^{\text{oracle}}$ indicates higher precision. Here we additionally include the $\text{PPL}^{\text{test}}$ on the test set in the table below: | | $\text{PPL}^{\text{test}}$↓ | $\text{PPL}^{\text{oracle}}$↓ | MAUVE↑ | |:---:|:---:|:---:|:---:| | MLE | 70.1 | 114.5 | 77.5 | | EMO | 74.9 | 55.9 | 83.4 |
• Label Smoothing: The MAUVE scores of label smoothing(smoothing factor set to 0.1) on language modeling tasks are shown below. Label smoothing performs even worse than MLE. The reason is that the number of plausible alternatives to ground-truth tokens is typically much smaller than the number of improbable tokens. Thus, label smoothing blindly puts more probability mass on low-quality tokens, leading to worse precision. | | WikiText-2 | WritingPrompt | PTB | AG News | |---|:---:|:---:|:---:|:---:| | MLE | 77.5 | 83.6 | 76.1 | 75.0 | | Label Smoothing | 75.8 | 81.6 | 71.9 | 73.8 |

Thanks for reading our responses. Regarding your questions, our further responses are as follows:

• On one hand, the diversity of MLE often come at the cost of generting adverse outputs due to its low precision. For language generation systems, high precision is arguably a higher priority since a single bad output can leave a lasting poor impression on the user. On the other hand, EMO's diversity manifests in the range of good tokens, which is more aligned/similar to human text distribution, as indicated by the Distinct N-gram, MAUVE score, and harmonic mean of forward/reverse cross-entropy.
• You are right, the reward can also be a word-level one, depending on what the environment can offer. To summarize, the main differences between EMO and RL-based training include: (1) motivation: exploring alternative probability distance measures for language modeling v.s. steer language model towards certain aspects (2) formula: weighted combination of gradient of each token in the vocabulary v.s. reward-augmented maximum likelihood of tokens in the output, and (3) application scenario: domain-specific fine-tuning/adaptation, instruction-tuning, and continual pre-training v.s. alignment.

Thank you for spending time reading our reponse. Regarding your concern, our responses are as follows:

• The number of synonyms/semantically-similar tokens is typically much smaller than that of improbable tokens. Therefore, by explicitly signifying training signal of low-quality tokens and the relative goodness of high-quality tokens, $Q_{\theta}^{\text{EMO}}$ is less diverse in terms of full output distribution because a large portion of probability mass are removed from those low-quality tokens and reassigned to a smaller set of plausible tokens. In other words, EMO is more diverse in the range of good tokens. Another evidence to show this that, even after applying advanced inference-time sampling strategies(e.g., temperature(T=0.7), top-p(p=0.9), and top-k(k=20)), MLE still underperform EMO because the distribution rescaling behavior of these sampling strategies is token-agnostic. On WikiText-2, the MAUVE score of MLE with top-p and top-k are 81.5, 80.3 respectively.
• Pursuing diversity in the full-distribution does not entail a human-like distribution. As indicated by the results we reported above, EMO leads to more human-like distribution in terms of (1) output N-gram diversity (2) MAUVE score and (3) Harmonic mean of forward/reverse cross-entropy.

Dear reviewer, we genuinely thank you for dedicating your time to reviewing the paper. Here, we strive to effectively respond to the issues you raised.

• recall-prioritization: Conceptually, as described in Section 2.2.1, a high recall means that high likelihood tokens under $P$ shall also have high likelihood under $Q_{\theta}$. In contrast, precision focuses on measuring whether low-quality tokens (unlikely under $P$) have low probabilities under $Q_{\theta}$. Ideally, with unlimited training data and model capacity, as well as a perfect optimizer, MLE-learned distribution will conform to the real data distribution. However, in practice, due to the lack of explicit training signals to penalize low-quality tokens, MLE-learned distribution tends to put non-trivial probability mass over low-quality samples.

  The operationalization of recall and precision in language modeling has also been investigated by Meister et. al[1]. They quantitiatively measure the recall by forward cross-entropy $\text{CE}^{\text{forward}}=-\sum{P\log{Q_{\theta}}}$, and precision by reverse cross-entropy $\text{CE}^{\text{reverse}}=-\sum{Q_{\theta}\log{P}}$. The recall-prioritization(low-precision) behavior of MLE can be observed by comparing the original distribution against the distributions modified by various inference-time truncation methods, e.g., top-p, top-k.

  Here we demonstrate this by reporting $\text{CE}^{\text{forward}}$, $\text{CE}^{\text{reverse}}$, and their harmonic mean against a reference distribution $P^{\text{ref}}$ provided by a larger GPT-Neo model. Results on the test set of WikiText-2 and WritingPromptare are shown in the following table. As can be seen, EMO emerges as a training objective that strikes a better balance between recall and precision. | | | WikiText-2 | | | WritingPrompt | | |:---:|:---:|:---:|:---:|:---:|:---:|:---:| | | $\text{CE}^{\text{reverse}}$↓ | $\text{CE}^{\text{forward}}$↓ | Harmonic Mean↓ | $\text{CE}^{\text{reverse}}$↓ | $\text{CE}^{\text{forward}}$↓ | Harmonic Mean↓ | | $Q^{\text{one-hot}}$ | 3.229 | 16.83 | 5.154 | 3.316 | 17.742 | 5.326 | | $Q_{\theta}^{\text{MLE}}$ | 4.503 | 4.333 | 4.219 | 4.250 | 3.920 | 3.994 | | $Q_{\theta}^{\text{EMO}}$ | 3.940 | 4.455 | 3.973 | 3.538 | 4.008 | 3.664 |

• diversity: Though EMO is better at harmonizing recall and precision, it does not directly translate to higher diversity in the model generations. In fact, the MLE-tuned model may generate more diverse outputs because it allocated probability mass more than it should to low-quality outputs. We present the Distinct N-gram($\frac{|\text{unique n-grams}|}{|\text{total n-grams}|}$) of GPT-2-generated continuations on the WikiText-2/WritingPrompt test set in the table below. Compared to MLE, the EMO-tuned model shows n-gram repetitiveness and MAUVE score closer to that of human texts by notable margins. | | | | WikiText2 | | | | | | WritingPrompt | | | |---|:---:|:---:|:---:|:---:|:---:|---|:---:|:---:|:---:|---|:---:| | | | | Dist$_{n}$ | | MAUVE | | | | | Dist$_{n}$ | | MAUVE | | | n=1 | n=2 | n=3 | n=4 | - | | | n=1 | n=2 | n=3 | n=4 | - | | $Q_{\theta}^{\text{MLE}}$ | 0.748 | 0.971 | 0.994 | 0.998 | 77.5 | | | 0.806 | 0.983 | 0.995 | 0.997 | 83.6 | | $Q_{\theta}^{\text{EMO}}$ | 0.698 | 0.955 | 0.990 | 0.997 | 87.5 | | | 0.762 | 0.972 | 0.993 | 0.996 | 87.4 | | Human | 0.687 | 0.958 | 0.988 | 0.996 | 100 | | | 0.764 | 0.974 | 0.993 | 0.997 | 100 |

• Qualitative difference: In addition to the quantitative difference between MLE and EMO, here we show some predicted next token probabilities from $Q_{\theta}^{\text{MLE}}$ and $Q_{\theta}^{\text{EMO}}$. We show the top-5 predicted tokens along with their probability due to space constraints. EMO tends to a

  • Prefix: Donald Trump has multiple occupations, such as
  • $Q_{\theta}^{\text{MLE}}$: the(0.393), a(0.325), construction(0.100), executive(0.093), president(0.088),...
  • $Q_{\theta}^{\text{EMO}}$: a(0.366), the(0.353), president(0.100), executive(0.091), construction(0.091),...
  • Prefix: Elephant is the largest living
  • $Q_{\theta}^{\text{MLE}}$: animal(0.275),creature(0.231),elephant(0.215),mammal(0.197),ant(0.079),...
  • $Q_{\theta}^{\text{EMO}}$: animal(0.268),creature(0.257),mammal(0.216),elephant(0.196),rept(0.062),...

• Beam Search: We present the MAUVE, BLEU, and ROUGE-1/2/L scores of model outputs generated by beam search(beam width set to 4) on the test set of WikiText-2. With beam search, EMO still displays higher similarity with human texts. For more results, please refer to our general responses where we compare the quality of greedy decoding generations from fine-tuned LLaMa2 models. | | MAUVE | BLEU | ROUGE-1 | ROUGE-2 | ROUGE-L | |:---:|:---:|:---:|:---:|:---:|:---:| | $Q_{\theta}^{\text{MLE}}$ | 9.93 | 21.31 | 37.0 | 24.7 | 33.6 | | $Q_{\theta}^{\text{EMO}}$ | 13.14 | 22.34 | 37.2 | 22.9 | 33.9 |

Thanks very much for your valuable suggestions and comments. We sincerely appreciate your time in reading the paper, and our point-to-point responses to your comments are given below.

• related work section: Thanks for your advice, we will include a related work section discussing how optimal transport was previously explored and exploited in NLP.

• perplexity: We present additional results of fine-tuned GPT-2's perplexity on the test set of WikiText-2 and WritingPrompts in the following table. In addition, we also report the token-level averaged forward cross-entropy $\text{CE}^{\text{forward}}=-\sum{P\log{Q_{\theta}}}$, reverse cross-entropy $\text{CE}^{\text{reverse}}=-\sum{Q_{\theta}\log{P}}$, and their harmonic mean, computed using the learned model distribution $Q_{\theta}$and the reference distribution $P_{\text{ref}}$ provided by a language model with stronger capacity, i.e., GPT-Neo. These metrics comprehensively delineate the behavior differences between MLE and EMO. Based on the results, $Q_{\theta}^{\text{EMO}}$ exhibits slightly higher ppl compared to $Q_{\theta}^{\text{MLE}}$. However, EMO excels at striking the balance between recall and precision, resulting in a more human-like distribution. | | Perplexity | MUAVE | $\text{CE}^{\text{forward}}$ | $\text{CE}^{\text{reverse}}$ | Harmonic Mean | |---|---|---|---|---|---| | $Q_{\theta}^{\text{MLE}}$ | 29.8 | 77.5 | 4.33 | 4.503 | 4.219 | | $Q_{\theta}^{\text{EMO}}$ | 30.9 | 87.5 | 4.45 | 3.940 | 3.973 |

• Decoding: Would you please kindly explain "discussion on decoding, ie generation"? If you are referring to the decoding method, we employ unbiased sampling as our primary decoding method as it allows us to explore the learned distribution in an unbiased way[1,2]. With EMO, we expect the output from the learned model distribution to be more similar to human distribution than MLE. Besides MAUVE, we also include the Distinct N-gram($\frac{|\text{unique n-grams}|}{|\text{total n-grams}|}$) of model-generated continuations on the test set of WikiText-2 and WritingPrompt. The closer to that of human text, the better. As shown in the table, EMO also displays higher similarity in terms of n-gram characteristics.
  | | | WikiText-2 | | | | WritingPrompt | | | |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| | | n=1 | n=2 | n=3 | n=4 | | |n=1 | n=2 | n=3 | n=4 | | MLE | 0.738 | 0.971 | 0.994 | 0.998 | | | 0.806 | 0.983 | 0.995 | 0.997 | | EMO | 0.698 | 0.955 | 0.990 | 0.997 | | | 0.762 | 0.972 | 0.993 | 0.996 | | Human | 0.687 | 0.948 | 0.988 | 0.996 | | | 0.764 | 0.974 | 0.993 | 0.997 |

Reference:

[1] Zhang et al.: MIXCE: Training Autoregressive Language Models by Mixing Forward and Reverse Cross-Entropies. ACL 2023

[2] Eikema et al.: Is MAP decoding all you need? the inadequacy of the mode in neural machine translation. ICCL 2020