SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chenze Shao\*1,2, Zhengrui Ma\*1,2, Min Zhang3 & Yang Feng†1,2
<sup>1</sup> Key Laboratory of Intelligent Information Processing
Institute of Computing Technology, Chinese Academy of Sciences
$^{2}$  University of Chinese Academy of Sciences
<sup>3</sup> School of Future Science and Engineering, Soochow University
chenzeshao@tencent.com, mazhengrui21b@ict.ac.cn
zhangminmt@hotmail.com, fengyang@ict.ac.cn
Maximum likelihood estimation (MLE) is a statistical method used to estimate the parameters of a probability distribution that best explain the observed data. In the context of text generation, MLE is often used to train generative language models, which can then be used to generate new text. However, we argue that MLE is not always necessary and optimal, especially for closed-ended text generation tasks like machine translation. In these tasks, the goal of model is to generate the most appropriate response, which does not necessarily require it to estimate the entire data distribution with MLE. To this end, we propose a novel class of training objectives based on convex functions, which enables text generation models to focus on highly probable outputs without having to estimate the entire data distribution. We investigate the theoretical properties of the optimal predicted distribution when applying convex functions to the loss, demonstrating that convex functions can sharpen the optimal distribution, thereby enabling the model to better capture outputs with high probabilities. Experiments on various text generation tasks and models show the effectiveness of our approach. It enables autoregressive models to bridge the gap between greedy and beam search, and facilitates the learning of non-autoregressive models with a maximum improvement of  $9+$  BLEU points. Moreover, our approach also exhibits significant impact on large language models (LLMs), substantially enhancing their generative capability on various tasks. Source code is available at https://github.com/ictnlp/Convex-Learning.
Text generation is an important field within natural language processing that aims to generate human-like texts for specific tasks. It can be broadly divided into two categories: open-ended and closed-ended text generation. Open-ended tasks encourage the model to produce novel and diverse outputs without a specific expected outcome or structure. Representative tasks in this category include language modeling [41, 7], chatbot [64], storytelling [13], etc. In contrast, closed-ended tasks are more constrained and adhere to specific rules or formats. Representative tasks in this category include machine translation [8, 2], text summarization [45], etc.
In recent years, learning neural probabilistic models with maximum likelihood estimation has become the dominant approach for both open-ended and closed-ended text generation [5, 2, 7]. Maximum likelihood estimation (MLE) is a statistical method used to estimate the parameters of a probability distribution that maximize the likelihood of the observed data [33]. Since directly maximizing the
*Equal contribution. Order determined by coin flip.
† Corresponding author: Yang Feng
likelihood can be numerically unstable, it is common to minimize the negative log-likelihood loss function, which is also referred to as cross-entropy loss. It is equivalent to minimizing Kullback-Leibler (KL) divergence [25, 1] between the true distribution and the predicted distribution, which ensures that the optimal predicted distribution is the true data distribution.
While MLE has gained widespread adoption, it does not always align with the objective of text generation, especially for closed-ended text generation tasks such as translation and summarization. In these tasks, the goal of the model is to generate the most appropriate response, rather than producing diverse outputs. For example, in the task of machine translation, though there may exist multiple translations for the same input sentence, we usually want the most accurate and commonly used translation result. Generally speaking, the desired output can be mathematically defined as the output with the maximum probability in the true data distribution, which does not necessarily require the model to estimate the entire data distribution with MLE.
In terms of generating the most probable output, MLE is also suboptimal for current neural text generation models. For autoregressive models, even if the model can perfectly fit the data distribution, it still requires decoding algorithms like greedy or beam search to generate the output, which do not guarantee the exact result with the maximum probability. To our knowledge, only Stahlberg and Byrne [55] proposed an exact decoding algorithm for autoregressive models, but it is too slow for practical applications. The limitation in exact decoding can be overcome by non-autoregressive models [17, 14], which independently predict the output at each position. However, fitting the data distribution by MLE is theoretically beyond the ability of non-autoregressive models [21]. In light of these issues, alternative training objectives should be considered to better address the specific requirements of text generation without incurring the shortcomings associated with MLE.
Based on the analysis above, MLE is suboptimal that it trains the model to estimate the data distribution, which complicates the training and decoding of text generation models. It would be advantageous if the model could converge to a sharper optimal distribution under an alternative loss function, as this would enable autoregressive models to easily find high probability outputs and also allow non-autoregressive models to converge to a better distribution. Exploring loss functions with this characteristic could lead to improved performance and efficiency of neural text generation models, particularly for closed-ended tasks.
In this paper, we propose a novel class of training objectives based on convex functions, which help text generation models capture highly likely outputs without estimating the entire data distribution. Intuitively, the concave shape of log-probability discourages the model from assigning a large prediction probability to a single sample, as the marginal benefit diminishes with increasing probability. If the learning criterion is convex or less concave, then intuitively the model would converge to a sharper distribution, which is the motivation of this work. We further investigate the theoretical properties of the optimal predicted distribution when applying convex functions to the loss. Our findings demonstrate that convex functions can sharpen the optimal distribution, allowing the model to better capture outputs with high probabilities.
Experiments on various closed-ended text generation tasks and models show the effectiveness of our approach. Specifically, it enables autoregressive models to bridge the gap between greedy and beam search, and facilitates the learning of non-autoregressive models with a maximum improvement of  $9+$  BLEU points. Moreover, our approach also exhibits significant impact on large language models, substantially enhancing their generative capability on various tasks.
2.1 Maximum Likelihood Estimation
Maximum likelihood estimation (MLE) is a statistical method used to estimate the parameters of a probability distribution that best explain the observed data. This is achieved by maximizing a likelihood function so that the observed data is most probable. Since directly maximizing the likelihood can be numerically unstable, it is common to minimize the negative log-likelihood loss function, also referred to as cross-entropy loss. Given the data distribution  $p_{data}$  and a parametric model with parameters  $\theta$ , MLE training minimizes:
$$
\mathcal {L} _ {M L E} (\theta) = - \mathbb {E} _ {x \sim p _ {\text {d a t a}} (x)} [ \log p _ {\theta} (x) ]. \tag {1}
$$
MLE can be viewed as an attempt to minimize KL divergence between the true underlying distribution of the data  $p_{data}$  and the estimated distribution  $p_{\theta}$  provided by the model [1]. The following equation reveals the relationship between MLE loss and KL divergence:
$$
\mathcal {D} _ {K L} \left(p _ {d a t a} \mid \mid p _ {\theta}\right) = \sum_ {x} p _ {d a t a} (x) \log \frac {p _ {d a t a} (x)}{p _ {\theta} (x)} = \mathcal {L} _ {M L E} (\theta) - H _ {d a t a}, \tag {2}
$$
where  $H_{data}$  is the Shannon entropy of the data distribution, which remains constant with respect to the model parameter  $\theta$ . Therefore, the MLE loss and KL divergence share the same minimizer that the estimated distribution  $p_{\theta}$  equals to the true distribution  $p_{data}$ . By minimizing the MLE loss, the predicted distribution is encouraged to be as close as possible to the true data distribution. In the context of text generation, this ensures that the model learns to generate text that closely resembles the text in the training data.
The above discussion can be extended to conditional scenarios. In such cases, the log-likelihood loss can be expressed as:
$$
\mathcal {L} _ {M L E} (\theta) = - \mathbb {E} _ {c \sim p _ {\text {d a t a}} (c)} [ \mathbb {E} _ {x \sim p _ {\text {d a t a}} (x | c)} [ \log p _ {\theta} (x | c) ] ], \tag {3}
$$
where  $c$  represents the input context. This extension allows the MLE framework to accommodate a wide range of text generation tasks such as machine translation, summarization, dialogue system, etc.
2.2 Text Generation Models
Based on how the sequence probability is factorized, neural text generation models can be broadly categorized into two types: autoregressive (AR) models and non-autoregressive (NAR) models. Autoregressive models generate text sequentially, predicting one token at a time based on the previously generated tokens. In AR models, the probability of generating a sequence  $x = (x_{1}, x_{2}, \dots, x_{T})$  is factorized as:
$$
p _ {\theta} (x | c) = \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {t} \mid x _ {<   t}, c\right), \tag {4}
$$
where  $c$  represents the input context. With the autoregressive decomposition, AR models can perfectly fit the data distribution if it satisfies  $p_{\theta}(x_t|x_{<t},c) = p_{data}(x_t|x_{<t},c)$  for every  $x,c,t$ . In inference, AR models can perform deterministic decoding like greedy/beam search to generate a high probability output, or sample from the model distribution to generate diverse outputs.
In contrast to autoregressive models, non-autoregressive models [17, 14] generate text in parallel, predicting all tokens simultaneously without conditioning on previously generated tokens. This approach can significantly speed up the generation process, as it removes the sequential dependency between tokens. In NAR models, the generation probability is factorized as:
$$
p _ {\theta} (x | c) = \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {t} \mid c\right). \tag {5}
$$
Unlike AR models, NAR models can efficiently find the most likely output by using argmax decoding at each step. However, MLE is beyond the ability of NAR models since they are theoretically unable to fit the data distribution. Huang et al. [21] showed that KL divergence from  $p_{\theta}$  to  $p_{data}$  is bounded by a non-negative constant:
$$
\mathcal {D} _ {K L} \left(p _ {d a t a} \mid \mid p _ {\theta}\right) \geq \mathcal {C} = - H _ {d a t a} (x | c) + \sum_ {t = 1} ^ {T} H _ {d a t a} \left(x _ {t} \mid c\right), \tag {6}
$$
The MLE loss is minimized when NAR models achieve the equality by ignoring sequential dependency and predicting  $p_{\theta}(x_t|c) = p_{data}(x_t|c)$ . Therefore, NAR models trained with MLE often suffer from reduced performance, as they lack the ability to model dependencies between tokens.
In this section, we will explore alternative loss functions for the learning of text generation models, which overcomes the limitations of MLE. We begin by introducing a general learning framework that allows arbitrary loss functions. Next, we discuss the benefits of applying convex functions to the loss within this framework. Finally, we use convex functions to construct composite loss functions, which can be used in practical text generation scenarios.
3.1 General Learning Framework
For simplicity of notation, we omit condition  $c$  in the probabilities, with the data distribution represented as  $p_{data}(x)$  and the model predicting the distribution  $p_{\theta}(x)$ . The derived theoretical results hold in both unconditional and conditional settings.
First, we introduce the general learning framework for text generation, characterized by the following loss function:
$$
\mathcal {L} _ {f} (\theta) = - \mathbb {E} _ {x \sim p _ {\text {d a t a}} (x)} [ f \left(p _ {\theta} (x)\right) ], \tag {7}
$$
where  $f$  is an arbitrary function of the prediction probability  $p_{\theta}(x)$ . We impose some basic requirements on  $f$ : (1) The domain of function  $f$  should contain the interval  $(0, 1]$ ; (2)  $f$  must be differentiable on the interval  $(0, 1]$  since we need to compute its gradient; and (3)  $f$  should be an increasing function on  $(0, 1]$  to encourage the model to generate the current sample. Under this framework, we can explain maximum likelihood estimation as a special case of  $f = \log$ , which is a differentiable and increasing function within the interval  $(0, 1]$ . We also establish some reasonable assumptions:
Assumption 1 (Countability of Sample Space). The sample space  $\mathcal{X}$  is countable, which allows us to enumerate all samples in a systematic way. Note that  $|\mathcal{X}|$  can be either finite or infinite.
Assumption 2 (Distinctness of Sample Probabilities). In the data distribution  $p_{data}$ , the probabilities of all samples are distinct, which allows us to arrange samples in a strictly descending order of sample probabilities.
Assumption 1 naturally holds in text generation tasks due to the inherent discreteness of textual data. With a countable sample space and probabilities lying in a dense subspace of real number, it is reasonable to assume the distinctness of sample probabilities. While Assumption 2 is not strictly necessary, removing it would introduce many corner cases that would complicate the subsequent analysis. In the following, we will assume that Assumptions 1-2 always hold, and we arrange the samples such that  $p_{data}(x_1) > p_{data}(x_2) > \dots > p_{data}(x_i) > \dots$ . Since the sample space  $\mathcal{X}$  is countable, the loss function in Equation 7 can be reformulated as follows:
$$
\mathcal {L} _ {f} (\theta) = - \sum_ {i = 1} ^ {| \mathcal {X} |} p _ {\text {d a t a}} \left(x _ {i}\right) \cdot f \left(p _ {\theta} \left(x _ {i}\right)\right). \tag {8}
$$
In this framework, our primary focus is to analyze the probability distribution  $p_{\theta}$  that the model is inclined to predict when the loss function is  $\mathcal{L}_f$ . We use  $p_f$  to denote the optimal distribution that minimizes the loss  $\mathcal{L}_f$ , which represents the expected outcome of the model. If  $\mathcal{L}_f$  has multiple optimal distributions, we use  $p_f$  to denote an arbitrary optimal distribution. This choice does not harm the generality of our analysis, as the subsequent discussion is applicable to all optimal distributions. Currently, it is only established that the optimal distribution for the MLE loss  $\mathcal{L}_{\log}$  is the data distribution  $p_{\mathrm{log}} = p_{\mathrm{data}}$ . For other loss functions, the following theorem reveals a general property of the optimal distribution. With samples organized in descending order of their probabilities in the data distribution, i.e.,  $p_{\mathrm{data}}(x_1) > p_{\mathrm{data}}(x_2) > \dots > p_{\mathrm{data}}(x_i) > \dots$ , the optimal distribution of an arbitrary function  $f$  maintains this order as  $p_f(x_1) \geq p_f(x_2) \geq \dots \geq p_f(x_i) \geq \dots$ . The proofs for the theorems presented in this paper can be found in Appendix A.
Theorem 1. Given an arbitrary differentiable and increasing function  $f$ , the optimal distribution  $p_f$  satisfies  $p_f(x_1) \geq p_f(x_2) \geq \dots \geq p_f(x_i) \geq \dots$ .
In the following, we will further explore the properties of optimal distributions associated with specific loss functions.
3.2 Loss with Convex Function
In certain text generation scenarios that require precise and deterministic outputs, it is beneficial for the model to converge to an optimal distribution that is sharper than the data distribution. In
3When  $\mathcal{X}$  is countably infinite, an arbitrary sequence of sample probabilities forms a convergent series since their sum is 1. This guarantees the existence of a maximum point in the series, ensuring that the sample probabilities can be arranged in a strictly descending order.
this section, we demonstrate that this objective can be achieved by employing convex functions as learning criterion.
The MLE loss function is based on log-probability, which is a concave function whose gradient decreases as the probability increases. The concave shape of the learning criterion prevents the model from assigning a large prediction probability to a single sample, since the marginal benefit diminishes as the probability increases. If function  $f$  is convex, then intuitively the model would converge to a sharper distribution. The following theorem validates this intuition, which shows that the optimal distribution  $p_f$  is a one-hot distribution when  $f$  is convex.
Theorem 2. If  $f$  is an increasing convex function on  $[0,1]$ , then the optimal distribution  $p_f$  is a one-hot distribution that  $p_f(x_1) = 1$  and  $p_f(x_i) = 0, i > 1$ .
The one-hot characteristic of the optimal distribution is advantageous for text generation models seeking precise and deterministic outputs. For autoregressive models, they do not need the computationally expensive beam search decoding any more if the model distribution is nearly one-hot. For non-autoregressive models, they suffer from reduced performance under MLE due to their inability to fit the data distribution. However, fitting a one-hot optimal distribution is well within their capabilities, allowing these models to generate high-quality outputs.
However, the direct application of loss with convex functions in training text generation models comes with an inherent limitation, impeding its practical utility. Specifically, the gradient of the parameter  $\theta$  tends to be very small when the prediction probability approaches 0, thereby rendering the training process inefficient. The gradient of  $\theta$  can be formulated as follows:
$$
\begin{array}{l} \frac {\partial \mathcal {L} _ {f} (\theta)}{\partial \theta} = - \mathbb {E} _ {x \sim p _ {d a t a} (x)} [ f ^ {\prime} (p _ {\theta} (x)) \cdot \frac {\partial p _ {\theta} (x)}{\partial \theta} ] \\ = - \mathbb {E} _ {x \sim p _ {\text {d a t a}} (x)} \left[ f ^ {\prime} \left(p _ {\theta} (x)\right) \cdot p _ {\theta} (x) \cdot \sum_ {t = 1} ^ {T} \frac {\partial \log \left(p _ {\theta} \left(x _ {t}\right)\right)}{\partial \theta} \right], \tag {9} \\ \end{array}
$$
where we have omitted the autoregressive history condition of  $p_{\theta}(x_t)$  for simplicity. The equation above indicates that the gradient is proportional to the sentence probability  $p_{\theta}(x)$ . In text generation models, the sentence probability  $p_{\theta}(x)$  is the product of token probabilities  $p_{\theta}(x_t)$ , which causes  $p_{\theta}(x)$  to be typically close to 0, especially when the model is newly initialized.
To counter this effect, the gradient  $f'(p_{\theta}(x))$  would need to approach infinity as  $p_{\theta}(x)$  approaches 0. For instance, the log-probability function has the gradient  $\frac{1}{p_{\theta}(x)}$ , which offsets the impact of  $p_{\theta}(x)$  such that  $f'(p_{\theta}(x)) \cdot p_{\theta}(x) = 1$ . However, for an increasing convex function  $f(p_{\theta}(x))$  whose gradient increases with  $p_{\theta}(x)$ , its gradient must be bounded when  $p_{\theta}(x)$  approaches 0, leading to an extremely small gradient update for the parameter  $\theta$  during training. This inherent limitation of loss with convex functions poses a significant hurdle to their practical applications.
3.3 Loss with Convex-composition Function
3.3.1 Theoretical Analysis
In the preceding discussion, we illustrate that while convex functions can induce a desirable one-hot optimal distribution, their inherent limitations during training pose significant impediments to practical applications. Consequently, we consider a relaxation of the convexity requirement, with the objective of rendering the function  $f$  less concave. This approach aims to obtain an optimal distribution that is sharper than  $p_{data}$ , thereby providing a practical solution that augments model performance without sacrificing training feasibility.
The standard loss function in maximum likelihood estimation is the negative log-probability, where log-probability is a concave function that yields a smooth optimal distribution. To render the learning criterion less concave, we propose a convex-composition approach that combines a convex function  $f$  with the original concave function  $g$ . This composition yields the following loss function:
$$
\mathcal {L} _ {f g} (\theta) = - \sum_ {i = 1} ^ {| \mathcal {X} |} p _ {\text {d a t a}} \left(x _ {i}\right) \cdot f g \left(p _ {\theta} \left(x _ {i}\right)\right), \tag {10}
$$
where  $f$  is an increasing convex function and  $g$  is an increasing concave function. The objective of this composition is to moderate the concavity of the overall loss function, thereby allowing for a sharper optimal distribution. The subsequent theorem and corollaries outline the theoretical properties associated with the optimal distribution under this function composition framework.
Theorem 3. Let  $f$  be an increasing convex function and  $g$  be an increasing concave function. Then, there exists a positive integer  $m$  such that the following inequalities hold:
1.  $p_{fg}(x_i) \geq p_g(x_i)$  for all  $i < m$ ,
2.  $p_{fg}(x_i) \leq p_g(x_i)$  for all  $i \geq m$ .
Corollary 1. The Shannon entropy of  $p_{fg}$  is less than or equal to the Shannon entropy of  $p_g$ .
Corollary 2. For any  $n \in \{1, 2, \ldots\}$ , the sum of the probabilities of the  $n$  most probable samples increases:  $\sum_{i=1}^{n} p_{fg}(x_i) \geq \sum_{i=1}^{n} p_g(x_i)$ .
Theorem 3 indicates that the convex-composition loss function tends to allocate higher probabilities to the more probable samples, while simultaneously diminishing the probabilities assigned to less probable ones, resulting in a sharper optimal distribution. Corollary 1 quantitatively establishes this observation, demonstrating that the incorporation of a convex function into the loss function effectively sharpens the optimal distribution, as evidenced by a reduction in the Shannon entropy of  $p_{fg}$  compared with  $p_g$ . Furthermore, Corollary 2 reveals an increase in the cumulative probability of the  $n$  most probable samples. Consequently, text generation models are better equipped to capture the highly probable outputs without explicitly modeling the data distribution.
In the above analysis, we only assume the original loss  $g$  to be an increasing concave function. By imposing specific conditions on  $g$ , we can derive more desirable properties from the optimal distribution  $p_{fg}$ , as demonstrated in the subsequent theorem:
Theorem 4. Let  $f$  be an increasing convex function and  $g$  be an increasing concave function. If  $g$  satisfies  $g''(x) \cdot g'(x) \geq g''(x)^2 > 0$  for all  $x \in (0,1)$ , then the difference between  $p_{fg}$  and  $p_g$  exhibits a monotonic order:  $p_{fg}(x_1) - p_g(x_1) \geq p_{fg}(x_2) - p_g(x_2) \geq \ldots \geq p_{fg}(x_{m-1}) - p_g(x_{m-1}) \geq 0$ , where  $m$  is the positive integer described in Theorem 3.