SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Combining steps 1-3, we obtain:
$$
p _ {d a t a} \left(x _ {i}\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right) \geq p _ {d a t a} \left(x _ {i}\right) \cdot g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right) \geq p _ {d a t a} \left(x _ {j}\right) \cdot g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right) \geq p _ {d a t a} \left(x _ {j}\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right). \tag {20}
$$
Further, the following steps show that the gradient is less than 0:
1.  $p_{fg}(x_i)\geq p_{fg}(x_j)$  , from theorem 1   
2.  $f^{\prime}g(p_{fg}(x_i))\geq f^{\prime}g(p_{fg}(x_j))$  , from step 1, the increasing property of  $g$  , and the convexity of  $f$    
3.  $\left.\frac{\partial\mathcal{L}_{fg}(p_\theta = p_{fg}^\prime)}{\partial\alpha}\right|_{\alpha = 0} <   0$  , from step 2 and inequality 20 (21)
Proof. The Shannon entropy of distribution  $p$ , denoted by  $H_{p}$ , is defined as  $H_{p} = -\sum_{x} p(x) \log p(x)$ . Consider a function  $h(\Delta x) = -(x_1 + \Delta x) \log (x_1 + \Delta x) - (x_2 - \Delta x) \log (x_2 - \Delta x)$ . It's first-order derivative  $h'(\Delta x) = \log (x_2 - \Delta x) - \log (x_1 + \Delta x)$ . Assuming  $x_1 \geq x_2$ , we observe that  $h'(\Delta x) < 0$  when  $\Delta x > 0$ , so the entropy decreases when we reduce  $x_2$  and increase  $x_1$  accordingly. The transformation from  $p_g$  to  $p_{fg}$  can be viewed as a series of such adjustments, which implies that the Shannon entropy of  $p_{fg}$  is less than or equal to the Shannon entropy of  $p_g$ .
Proof. The theorem guarantees the existence of a positive integer  $m$  such that  $p_{fg}(x_i) \geq p_g(x_i)$  for all  $i < m$  and  $p_{fg}(x_i) \leq p_g(x_i)$  for all  $i \geq m$ . In the case where  $n < m$ , we have  $p_{fg}(x_i) \geq p_g(x_i)$  for all  $i$  with  $1 \leq i \leq n < m$ . This leads to the inequality  $\sum_{i=1}^{n} p_{fg}(x_i) \geq \sum_{i=1}^{n} p_g(x_i)$ . In the case where  $n \geq m$ , we have  $p_{fg}(x_i) \leq p_g(x_i)$  for all  $i$  with  $m \leq n < i$ . Therefore, we can write  $\sum_{i=1}^{n} (p_{fg}(x_i) - p_g(x_i)) = -\sum_{i=n+1}^{|X|} (p_{fg}(x_i) - p_g(x_i)) \geq 0$ . Consequently, the inequality  $\sum_{i=1}^{n} p_{fg}(x_i) \geq \sum_{i=1}^{n} p_g(x_i)$  also holds. Therefore, in both cases, the corollary is proved.
A.4 Proof of Theorem 4
Proof. We prove this theorem by contradiction. Assuming the theorem does not hold, there must exist an index index  $(i,j)$  with  $i < j$  and  $0 \leq p_{fg}(x_i) - p_g(x_i) < p_{fg}(x_j) - p_g(x_j)$ . In this case, we can construct a distribution  $p_{fg}'$  with a lower loss than  $p_{fg}$ , which contradicts the optimality of  $p_{fg}$ . Specifically, let  $p_{fg}'$  be identical to  $p_{fg}$  except for the changes  $p_{fg}'(x_i) = p_{fg}(x_i) + \alpha$  and  $p_{fg}'(x_j) = p_{fg}(x_j) - \alpha$ , where  $0 \leq \alpha \leq p_{fg}(x_j)$ . Then we can calculate the gradient of loss  $\mathcal{L}_{fg}(p_\theta = p_{fg}')$  with respect to  $\alpha$ :
$$
\left. \frac {\partial \mathcal {L} _ {f g} \left(p _ {\theta} = p _ {f g} ^ {\prime}\right)}{\partial \alpha} \right| _ {\alpha = 0} = p _ {d a t a} \left(x _ {j}\right) \cdot f ^ {\prime} g \left(p _ {f g} \left(x _ {j}\right)\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right) - p _ {d a t a} \left(x _ {i}\right) \cdot f ^ {\prime} g \left(p _ {f g} \left(x _ {i}\right)\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right). \tag {22}
$$
Our goal is to demonstrate that  $\left.\frac{\partial\mathcal{L}_{fg}(p_{\theta} = p_{fg}^{\prime})}{\partial\alpha}\right|_{\alpha = 0} < 0$ , which would contradict the assumption that  $p_{fg}$  minimizes  $\mathcal{L}_{fg}$ , thereby proving the theorem.
Given the optimality of  $p_g$ , we have  $p_{data}(x_i) \cdot g'(p_g(x_i)) \geq p_{data}(x_j) \cdot g'(p_g(x_j))$ , otherwise we can reduce  $p_g(x_i)$  to obtain a lower loss. From Theorem 1, we know that  $p_{fg}(x_i) \geq p_{fg}(x_j)$ , and because  $f$  is convex and  $g$  is increasing, we have  $f'g(p_{fg}(x_i)) \geq f'g(p_{fg}(x_j))$ . Using these inequalities, we obtain a upperbound of equation 22:
$$
\begin{array}{l} \left. \frac {\partial \mathcal {L} _ {f g} \left(p _ {\theta} = p _ {f g} ^ {\prime}\right)}{\partial \alpha} \right| _ {\alpha = 0} / \left(p _ {d a t a} \left(x _ {i}\right) \cdot g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)\right) \\ = \frac {p _ {d a t a} (x _ {j}) \cdot f ^ {\prime} g \left(p _ {f g} (x _ {j})\right) \cdot g ^ {\prime} \left(p _ {f g} (x _ {j})\right)}{p _ {d a t a} (x _ {i}) \cdot g ^ {\prime} \left(p _ {g} (x _ {i})\right)} - \frac {p _ {d a t a} (x _ {i}) \cdot f ^ {\prime} g \left(p _ {f g} (x _ {i})\right) \cdot g ^ {\prime} \left(p _ {f g} (x _ {i})\right)}{p _ {d a t a} (x _ {i}) \cdot g ^ {\prime} \left(p _ {g} (x _ {i})\right)} \\ \leq \frac {p _ {d a t a} (x _ {j}) \cdot f ^ {\prime} g (p _ {f g} (x _ {j})) \cdot g ^ {\prime} (p _ {f g} (x _ {j}))}{p _ {d a t a} (x _ {j}) \cdot g ^ {\prime} (p _ {g} (x _ {j}))} - \frac {p _ {d a t a} (x _ {i}) \cdot f ^ {\prime} g (p _ {f g} (x _ {i})) \cdot g ^ {\prime} (p _ {f g} (x _ {i}))}{p _ {d a t a} (x _ {i}) \cdot g ^ {\prime} (p _ {g} (x _ {i}))} \\ = \frac {f ^ {\prime} g \left(p _ {f g} \left(x _ {j}\right)\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)} - \frac {f ^ {\prime} g \left(p _ {f g} \left(x _ {i}\right)\right) \cdot g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)} \\ \leq f ^ {\prime} g \left(p _ {f g} \left(x _ {j}\right)\right) \cdot \left(\frac {g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)} - \frac {g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)}\right). \tag {23} \\ \end{array}
$$
To demonstrate that  $\left.\frac{\partial\mathcal{L}fg(p\theta = p_{fg}^{\prime})}{\partial\alpha}\right|_{\alpha = 0} < 0$ , we only need to prove the following inequality:
$$
\frac {g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)} - \frac {g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)} <   0. \tag {24}
$$
Let  $\Delta x = p_{fg}(x_i) - p_g(x_i) < p_{fg}(x_j) - p_g(x_j)$ . As a result,  $p_{fg}(x_j) > \Delta x + p_g(x_j)$ , and thus  $g'(p_{fg}(x_j)) < g'(p_g(x_j) + \Delta x)$ . This allows us to further simplify the inequality:
$$
\frac {g ^ {\prime} \left(p _ {f g} \left(x _ {j}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)} - \frac {g ^ {\prime} \left(p _ {f g} \left(x _ {i}\right)\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)} <   \frac {g ^ {\prime} \left(p _ {g} \left(x _ {j}\right) + \Delta x\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)} - \frac {g ^ {\prime} \left(p _ {g} \left(x _ {i}\right) + \Delta x\right)}{g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)}. \tag {25}
$$
To establish that the right-hand side of the above inequality is non-positive, we can apply the logarithm transformation and show the following inequality instead:
$$
\log \left( \right.g ^ {\prime} \left(p _ {g} \left(x _ {j}\right) + \Delta x\right) - \log \left(g ^ {\prime} \left(p _ {g} \left(x _ {j}\right)\right)\right) \leq \log \left(g ^ {\prime} \left(p _ {g} \left(x _ {i}\right) + \Delta x\right) - \log \left(g ^ {\prime} \left(p _ {g} \left(x _ {i}\right)\right)\right). \right. \tag {26}
$$
Let's denote  $h(x) = \log(g'(x))$ ,  $x_1 = p_g(x_i)$ , and  $x_2 = p_g(x_j)$ . The above inequality can be simplified to:
$$
h \left(x _ {2} + \Delta x\right) - h \left(x _ {2}\right) \leq h \left(x _ {1} + \Delta x\right) - h \left(x _ {1}\right), \tag {27}
$$
where  $\Delta x\geq 0$  and  $x_{2}\leq x_{1}$  according to Theorem 1. The above inequality holds when  $h(x)$  is a convex function. The second-order derivative of  $h(x) = \log (g^{\prime}(x))$  is:
$$
h ^ {\prime \prime} (x) = \frac {g ^ {\prime \prime \prime} (x) g ^ {\prime} (x) - g ^ {\prime \prime} (x) ^ {2}}{g ^ {\prime} (x) ^ {2}}. \tag {28}
$$
Therefore,  $h(x)$  is a convex function under the condition  $g'''(x) \cdot g'(x) \geq g''(x)^2$ . This verifies that  $\left. \frac{\partial \mathcal{L}_{fg}(p_{\theta} = p_{fg}')}{\partial \alpha} \right|_{\alpha = 0} < 0$ , completing the proof by contradiction.
B Experimental Settings
B.1 Machine Translation
B.1.1 Datasets and Metrics
Datasets We conduct experiments on widely used translation benchmark: WMT14 English-German (EN-DE, 4.5M), where the validation and test sets are newtest2013 and newtest2014 respectively. We apply BPE [48] with 32K merge operations to learn a joint vocabulary on the tokenized data. Considering the major topic of this research is how to learn from a real-world data distribution, we don't apply any tricks that may have an influence on the distribution, e.g., knowledge distillation.
Metrics The overall quality of translation is assessed using metrics BLEU [38] and COMET [43]. In the case of non-autoregressive models, we additionally quantify the prediction confidence and translation fluency of the generated output. Prediction confidence is measured with negative log-likelihood (NLL) of model generation. A lower NLL value indicates a more focused model distribution and higher prediction confidence. To evaluate translation fluency, we utilize an external pre-trained autoregressive language model. The generated translation is fed to the language model using teacher forcing and the resulting perplexity (PPL) is calculated as a measure of fluency. A lower external PPL score indicates a higher level of fluency.
B.1.2 Implementation Details
Architectures In order to validate the overall efficacy of convex-composition loss, we perform experiments using various model architectures. We adopt Transformer-base [63] as our autoregressive baseline and Vanilla-NAT [17], CMLM [14] and CTC [46] as our non-autoregressive baselines. We apply uniform copy to construct decoder inputs in Vanilla-NAT and CTC. The decoder length in CTC is set to  $2 \times$  the source length.
Training Although training with convex-composition loss offers the desirable property of optimality, it can encounter gradient vanishing issues during initialization as analyzed previously. To mitigate this, we employ a two-step training approach: MLE pre-training followed by fine-tuning with convex-composition loss. This approach allows us to avoid numerical gradient issues while still benefiting from the optimality achieved through convex composition. For training with convex-composition loss, we set the exponent hyperparameter  $k$  to 1 for the autoregressive model and tune it from \{1,2,3,5,8\} on the validation set for non-autoregressive models. Throughout both MLE and convex-composition training, all models are optimized using the Adam optimizer [24] with  $\beta = (0.9,0.98)$  and  $\epsilon = 10^{-8}$ . Detailed information regarding other training hyperparameters can be found in Table 7.
Table 7: Settings of training hyperparameters on WMT14 EN  $\leftrightarrow$  DE dataset.
Transformer	Vanilla-NAT	CMLM	CTC
MLE	Convex	MLE	Convex	MLE	Convex	MLE	Convex
batch size	32k	32k	64k	256k	64k	256k	64k	256k
learning rate	7e-4	2e-4	5e-4	3e-4	5e-4	3e-4	5e-4	3e-4
warmup steps	4k	1k	10k	500	10k	500	10k	500
training steps	200k	50k	300k	10k	300k	10k	300k	10k
dropout	0.1	0.1	0.3	0.3	0.3	0.3	0.3	0.1
weight decay	0	0	0.01	0.01	0.01	0.01	0.01	0.01
label smoothing	0.1	0.1	0.1	0	0.1	0	0.01	0
length loss factor	-	-	0.1	0.01	0.1	0.01	-	-
Decoding For the autoregressive model, we set the beam length to 5 by default and tune the length penalty on the validation set unless stated otherwise. For Vanilla-NAT and CTC, we utilize fully non-autoregressive argmax decoding. In the case of CMLM, we employ 5 length candidates and
<sup>5</sup>We use checkpoint Unbabel/wmt22-comet-da to compute COMET score. It is available at https://github.com/Unbabel/COMET.
<sup>6</sup>We use checkpoint transformer_lm.wmt19.de to compute the external PPL score. It is available at https://github.com/facebookresearch/fairseq/tree/main/examples/language_model.
disable iteration for inference. The decoding speedup is measured with a batch size of 1 on GeForce RTX 3090 GPUs.
B.2 Abstractive Summarization
B.2.1 Datasets and Metrics
We conduct experiments on two widely used summarization benchmarks: CNN/DailyMail [18] and Xsum [34]. CNN/DailyMail contains 220K articles from the Daily Mail newspaper and 93K articles from CNN. Each article contains a bullet point summary consisting of multiple sentences. We use the non-anonymized variant following [47, 29]. After the pre-processing, there are 311,971 <article, summary> pairs. XSum consists of 227K online articles from the British Broadcasting Corporation (BBC), containing professionally written single-sentence summaries. After the preprocessing, there are 226,677 <article, summary> data pairs. In order to maintain consistency with previous works [26, 40], we employ GPT-2 tokenizer to tokenize raw CNN/DailyMail data, and Berttokenizer to tokenize raw Xsum data. The summarization quality is measured with ROUGE-1, ROUGE-2 and ROUGE-L [28] as discussed in [6].
B.2.2 Implementation Details
In our summarization experiments, most of the implementation details of the Transformer align with those used in translation. However, there are a few modifications to ensure consistency with previous work [26]. We apply layer normalization to the embeddings. The attention dropout is set to 0.1, and the weight decay is set to 0.01. We utilize beam search with a size of 4 during decoding. The length penalty, max_len_b, and min_len are set to 2.0, 140, and 55, respectively on CNN/DailyMail dataset. We use a length penalty of 1.2 on Xsum dataset. For CNN/DailyMail dataset, we additionally employ a tri-gram repetition prevention trick.
B.3 Large Language Models
For the development of LLMs, we utilize LLaMA-7B and LLaMA-13B [61] as our foundation models. We conduct instruction tuning using the Alpaca dataset by GPT4 [65, 60], which comprises 52K instruction-following demonstrations. Instead of the standard cross-entropy loss employed during instruction tuning, we adopt the convex-composition loss of exponential form to fine-tune foundation models.
The generative capability of LLMs is also evaluated on the two representative closed-ended text generation tasks: machine translation and text summarization. For machine translation, we follow previous works [22, 70, 69, 30] to evaluate the translation capability on four WMT22 translation tasks (Chinese-to-English, English-to-Chinese, German-to-English, and English-to-German). For text summarization, we follow Liu et al. [30] to conduct the evaluation on CNN/DailyMail Dataset [18]. We employ beam search with a beam size of 4 for machine translation and 2 for summarization. The prompt for machine translation is "Translate the following sentences from [SRC] to [TGT]." The prompt for summarization is "Write a brief and focused summary of the passage that follows."
C Effects of  $k$  on AR Models
We study the effects of exponent hyper-parameter  $k$  on autoregressive models. Table 8 presents the BLEU scores of autoregressive models as the exponent  $k$  varies, showing that the optimal performance is achieved when  $k = 1$ . Other choices of  $k$ , such as  $k = 0.5$  or 0.75, also yield improvements, predominantly in the context of the greedy search setting.
Table 8: BLEU scores of autoregressive models as the exponent k varies on WMT14 EN-DE test set.
k-th Power	0.5	0.75	1	2	3
Greedy	26.89	26.89	26.92	26.78	26.13
Beam5	27.62	27.74	27.78	27.49	26.76
D Correlations on Other NAR Models
In Section 4.2, we present compelling evidence in support of the mode collapse property of convex function effectively mitigating the multimodality issue in the NAR model. This evidence is derived from the strong correlation observed between model entropy and generation fluency in the CMLM model, as demonstrated in Table 4. In this section, we provide additional evidence for other NAR models to further support our findings in Table 9 and 10.
Table 9: Prediction confidence (Output NLL) and generation fluency (External PPL) of Vanilla-NAT on WMT14 EN-DE test set.
k-th Power	1	2	3	5	8
Confidence (Output NLL) ↓	23.34	16.17	11.32	6.25	6.27
Fluency (External PPL) ↓	1000.06	730.91	463.78	344.56	353.40
Table 10: Prediction confidence (Output NLL) and generation fluency (External PPL) of CTC on WMT14 EN-DE test set.
k-th Power	1	2	3	5	8
Confidence (Output NLL) ↓	18.74	13.88	11.20	7.69	5.55
Fluency (External PPL) ↓	174.79	142.80	134.28	137.07	154.60
E Results on Alternative Choice of Convex Function
In addition to the exponential function, we have also explored another choice of convex function in our framework of convex-composition loss. In this section, we discuss the results of the choice of power function, i.e.,  $\mathcal{L}_f(\theta) = -\mathbb{E}_{x\sim p_{data}(x)}[-(-\frac{\log(p_\theta(x))}{T})^k], 0\leq k\leq 1$ . The results obtained from applying the power function in convex-composition loss are presented in Table 11.
Table 11: Results of BLEU scores by applying power function in convex-composition loss. We denote the MLE baseline by using  $k = 1.0$ . We employ greedy decoding for Transformer. In cases where training fails, we use "N/A" to denote such instances.
k-th Power	0.1	0.3	0.5	0.7	1.0
Transformer	26.64	26.68	26.60	26.52	26.48
Vanilla-NAT	N/A	N/A	10.74	10.51	10.41
We have observed that the benefits of applying the power function within the convex composition framework are significantly marginal compared to the exponential function, especially in the case of Vanilla-NAT. In addition, we have found the training process may encounter difficulties or failure when  $k$  is approaching 0. We attribute such problem to the shape of  $f'(g(p_{\theta}(x)))$  when power function is applied, i.e.,  $k \cdot (-\frac{\log(p_{\theta}(x))}{T})^{k - 1}$ .
As shown in Figure 3, the value of  $f'(g(p_{\theta}(x)))$  will approach a constant 1 as  $k$  approaches 1. This phenomenon arises due to the reduction in the convexity of function  $f$ , resulting in a decrease in gain. In case of  $k$  approaching 0, the situation is even worse where  $f'(g(p_{\theta}(x)))$  will experience a sudden increase from an extremely small value near 0. These factors result in an unstable training process and contribute to the power function being less suitable within the framework of convex-composition loss.
Figure 3: Shapes of  $f'(g(p_{\theta}(x)))$  when different convex functions are applied.
F Results on Diverse Generation
We study the effects of convex functions on VAE-based text generation models by replacing the log-probability-based reconstruction loss in ELBO with the convex-composition loss. Formally, we train the model using the following loss:
$$
\mathbb {E} _ {z \sim q (z | x)} - f \left(\frac {1}{T} \log p (x | z)\right) + \mathrm {K L} (q (z | x) \| p (z)), \tag {29}
$$
where we opt for the convex function  $f$  to be  $e^{kx}$ ,  $k \geq 0$ . We perform experiments within the context of conditional generation, utilizing a VAE-based non-autoregressive model [54, 16] for the task of machine translation. During inference, we randomly sample the latent variable 3 times to generate diverse texts. We assess the quality with BLEU score computed against reference (reference-BLEU) and measure the diversity with BLEU score computed against each other (pairwise-BLEU). The average value and standard derivation are reported in Table 12.
Table 12: Reference-BLEU and Pairwise-BLEU scores of VAE-based NAT models trained with different objectives on WMT14 EN-DE test set. The texts are generated by sampling the latent distribution 3 times.
ELBO	Convex + KL
Reference-BLEU	16.23±.14	23.35±.04
Pairwise-BLEU	29.52±.20	91.91±.03
During the training process, we have observed that KL divergence tends to vanish more readily when the convex functions are applied. We attribute this phenomenon to the smaller norms of gradients associated with the convex-composition loss. As a result, the gradient of the KL divergence dominates the model update, leading to the KL divergence vanishing. We note VAE-based text generation models trained using the convex-composition loss exhibit a higher generation quality while suffering from poor diversity, which is consistent with the mode collapse property of convex function.
G Analysis of Convex Learning and Knowledge Distillation
With the ability to capture a concentrated distribution from datasets exhibiting a multi-modal distribution, the proposed convex learning approach shows similar dynamics to knowledge distillation [19], a technique which encourages the student model to imitate the output of the teacher model. To compare the two methods, we utilize autoregressive Transformer as the teacher and apply sequence-level knowledge distillation [23] to construct a dataset of lower complexity, and train the models using different losses.
Table 13: BLEU scores of autoregressive and vanilla-NAT models trained with or without knowledge distillation (KD) on WMT14 EN-DE test set.
Transformer	Vanilla-NAT
MLE	Convex	MLE	Convex
w/ KD	27.73	27.80	19.18	23.17
w/o KD	27.57	27.78	10.41	16.74
Figure 4: BLEU scores of autoregressive models as beam size varies with or without knowledge distillation (KD) on WMT14 EN-DE test set.
The results in Table 13 and Figure 4 demonstrate that convex learning and knowledge distillation have similar effects on text generation models. Both methods lead to significant improvements on non-autoregressive models and bridge the performance gap between greedy and beam search of autoregressive models. It is worth noting that training with the convex-composition loss avoids the intricate process of training an additional teacher model and decoding the whole training set to achieve the improvements. Moreover, convex-composition loss can be combined with knowledge distillation to further enhance the performance.