SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shentao Yang $^{1}$ , Shujian Zhang $^{1}$ , Congying Xia $^{2}$ , Yihao Feng $^{2}$ , Caiming Xiong $^{2}$ , Mingyuan Zhou $^{1}$
<sup>1</sup>The University of Texas at Austin  
<sup>2</sup>Salesforce Research  
shentao.yang@mccombs.utexas.edu, yihao.ac@gmail.com  
mingyuan.zhou@mccombs.utexas.edu
Aligning language models (LMs) with preferences is an important problem in natural language generation. A key challenge is that preferences are typically provided at the sequence level while LM training and generation both occur at the token level. There is, therefore, a granularity mismatch between the preference and the LM training losses, which may complicate the learning problem. In this paper, we address this issue by developing an alternate training process, where we iterate between grounding the sequence-level preference into token-level training guidance, and improving the LM with the learned guidance. For guidance learning, we design a framework that extends the pairwise-preference learning in imitation learning to both variable-length LM generation and the utilization of the preference among multiple generations. For LM training, based on the amount of supervised data, we present two minimalist learning objectives that utilize the learned guidance. In experiments, our method performs competitively on two distinct representative LM tasks — discrete-prompt generation and text summarization. Source codes are released at https://github.com/Shentao-YANG/Preference_Grounded_Guidance.
Language models (LMs) have been successfully trained with token-level cross-entropy losses, where each token position has a corresponding term in the overall training losses [1-11]. Recent studies have shown that LMs can be further improved by aligning them with preferences from human feedback [12-15] or automatic evaluation metrics [16-18]. Typically, the preferences are only provided at the sequence level, e.g., "Which of the two generated text sequences is better?" To align LMs with sequence-level preferences, there exist a variety of approaches, such as applying external filters to the training texts [19], performing supervised learning on some curated/improved datasets [20-22], and optimizing the LMs based on a learned sequence-level (pairwise-) preference predictor [14, 23-25].
While these approaches have contributed to the development of several revolutionary products [e.g., 18, 15], a mismatch issue has emerged from the perspective of guiding LM fine-tuning. Concretely, the sequence-level preference is not grounded into the token level, where LM training losses occur. This means that there is a mismatch in granularity between the feedback and training losses — the preference is coarse-grained while the training losses are fine-grained. This issue is similar to the delayed-feedback problem in reinforcement learning (RL) [26-28], where informative feedback is available only at the end of the trajectory (sequence) and not at any of the intermediate timesteps. Previous studies have noted that this problem could have a negative impact on the empirical performance of the resulting LMs [29, 30], as it introduces a more challenging learning problem characterized by higher gradient variance and lower sample efficiency to achieve the learning goal [31, 32].
To address this granularity mismatch, we focus on the following question: How can we effectively ground sequence-level preference into token-level guidance for LM fine-tuning? We propose an alternate training process that alternates between two stages: learning preference-grounded token-level guidance and improving the LM using the learned guidance. This alternate process reduces the requirement on supervised data and targets the low-data regime, e.g., few/zero-shot learning, where task-specific supervised (pre-)training is infeasible and initial LMs have weak zero-shot abilities.
To ground the sequence-level preference into token-level guidance, we propose a framework for learning a token-level "reward" function<sup>1</sup>, inspired by reward-learning-from/preferences in the imitation learning (IL) literature [33-35]. Specifically, we train the token-level rewards such that the corresponding evaluation for a generated text sequence reflects the preference among multiple alternative generations, where the preference comes from task-specific evaluation. While summation is classically used in IL to aggregate the learned token-level rewards into the text-sequence evaluation, LM tasks can be different from classical IL tasks. To cater to LM generations, our guidance-learning framework can accommodate more careful choices of the aggregation function beyond the classical summation. For instance, in generating text prompts to steer an LM for text classification, a "key token" in the prompt may be more effective than several mediocre tokens. Hence, using maximum to aggregate the token-level rewards may better reflect the text-sequence quality than summation.
To utilize the learned preference-grounded guidance in LM training, we present two minimalist learning objectives that contain only a minimal number of hyperparameters. These two objectives respectively target different amounts of supervised data in the specific LM task. We evaluate our framework on two distinct representative LM tasks: generating discrete text prompts for few-shot text classification and text summarization. On both tasks, our method exhibits competitive performance.
2 Main Method
Before diving into technical details, we will first establish the notation, provide some background on classical pairwise-preference learning, and offer an overview of our preference-grounding process.
Notation. In most LM tasks, we are given a dataset  $\mathcal{D} = \{(x^i,y^i)\}_{i = 1}^N$  of  $N$  supervised examples, where  $x$  is the input to the LM, which can be a dummy, and  $y$  is the target text-sequence. We denote the LM parameterized by  $\theta$  as  $\pi_{\theta}$ . The  $t^{\mathrm{th}}$  generated token is denoted as  $a_{t}$ , given by  $a_{t}\sim \pi_{\theta}(\cdot |s_{t})$ , where the context for token generation at step  $t\geq 0$  is denoted as  $s_t$ , consisting of the LM input  $x$  and the previously generated tokens  $a_{< t} = (a_0,\ldots ,a_{t - 1})$ . Specifically,  $s_0 = x$  and  $\forall t > 0,s_t = (x,a_{< t})$ . The full generated text-sequence of length  $T$  is denoted as  $\pmb {a} = (a_0,\dots ,a_{T - 1})$ . In most LM tasks, we have a task-specific evaluation metric  $\mathcal{R}(s_T,y)\in \mathbb{R}$  that depends on the final context  $s_T$  of the generated sequence and the target sequence  $y$ , with  $s_T = (x,\pmb {a})$ . The objective of LM training is often to maximize the expected task-specific evaluation, which can be expressed as
$$
\max _ {\boldsymbol {\theta}} \mathbb {E} _ {(x, y) \sim \mathcal {D}} \mathbb {E} _ {\boldsymbol {a} \sim \prod_ {t = 0} ^ {T - 1} \pi_ {\boldsymbol {\theta}} (a _ {t} \mid s _ {t})} \left[ \mathcal {R} (s _ {T} = (x, \boldsymbol {a}), y) \right].
$$
We model the learned token-level guidance as a bounded (reward) function  $r_{\phi}(s_t, a_t) \in [0, 1]$ , parametrized by  $\phi$ . Unlike the original sequence-level preference or evaluation that is only available at the final step  $T$ , the trained  $r_{\phi}$  can densely guide the token selection at each step  $t \leq T$ .
Pairwise Preference Learning. In reward-learning-from/preferences [e.g., 33-35], a dense reward function is learned such that the sum-aggregated reward for the entire generation trajectory aligns with the pairwise preference between two trajectories. In the context of LM generation, suppose we have two text-generation trajectories  $\tau^i$  and  $\tau^j$  associated with the same LM input and target  $(x,y)$ , taking the form  $\tau^i = \{(s_0^i,a_0^i),\dots,(s_{T^i -1}^i,a_{T^i -1}^i)\}$  with sequence lengths  $T^{i}$  and  $T^{j}$ , respectively. Assume that  $\tau^i$  is preferred over  $\tau^j$ , denoted as  $\tau^i\succ \tau^j$ . A token-level reward function  $r_{\phi}(s_t,a_t)$  is learned by requiring  $\sum_{t = 0}^{T^i -1}r_{\phi}(s_t^i,a_t^i) > \sum_{t = 0}^{T^j -1}r_{\phi}(s_t^j,a_t^j)$ . Following the Bradley-Terry model of preferences [36], the pairwise-preference loss for reward-function learning is
$$
\ell (\phi) = - \log \left[ \exp \left(\sum_ {t = 0} ^ {T ^ {i} - 1} r _ {\phi} \left(s _ {t} ^ {i}, a _ {t} ^ {i}\right)\right) / \sum_ {k \in \{i, j \}} \exp \left(\sum_ {t = 0} ^ {T ^ {k} - 1} r _ {\phi} \left(s _ {t} ^ {k}, a _ {t} ^ {k}\right)\right) \right], \tag {1}
$$
which is often interpreted as binary classification in the literature [37-39]. In Eq. (1), summation  $\sum (\cdot)$  is used to aggregate the learned token-level rewards into a parametrized sequence-level evaluation.
<sup>1</sup>We use the words "guidance" and "reward", "fine-tuning" and "training" interchangeably, depending on the specific context.
Figure 1: Overview of the proposed framework. "AVG" denotes average, which is an example of the aggregation function  $f(\cdot)$  discussed in Section 2.1. "Seq Eval" refers to the parametrized sequence-level evaluations. The model choice of the reward function and LM depends on the specific task and is discussed in Section 4.
Overview. To ground the sequence-level preference into token-level guidance for LM training and thereby address the granularity mismatch discussed in Section 1, we present an alternate learning process that alternately learns the token-level guidance and trains the LM using the learned guidance.
For learning the preference-grounded guidance, in Section 2.1 we propose a framework that learns a token-level reward function that reflects the preference among multiple generated sequences. To utilize the learned preference-grounded guidance, based on the amount of supervised data in the specific task, in Section 2.2 we present two minimalist LM training approaches that require only minimal tuning. In our framework, we iterate between the above two steps to mitigate the distribution shift between the text sequences used to train the reward function and the text sequences evaluated by the reward function during LM training, taking into account that LMs can evolve during the training process. Our alternate-learning procedure is illustrated in Fig. 1.
2.1 Token-level Guidance Learning for Preference Grounding
Instead of applying the pairwise approach discussed above, we utilize the preference among multiple generated text sequences to learn the reward function  $r_{\phi}(s_t,a_t)$ . Intuitively, we use more information to train the reward function at each optimization step. Therefore, our approach can be more efficient and effective, especially when the optimization budget is limited.
Concretely, suppose we have  $K \geq 2$  generated text-sequences  $(a^1, \ldots, a^K)$  for the same LM input and target  $(x, y)$ , with the associated generation trajectories  $(\tau^1, \ldots, \tau^K)$  and with possibly unequal sequence-lengths  $(T^1, \ldots, T^K)$ . Assume that there is a preference ordering among these  $K$  sequences, where by "preference" we mean a ranking of text sequences based on some evaluations of full text-sequences. For description simplicity, let the preference ordering be  $a^1 \succ \dots \succ a^K \Longleftrightarrow \tau^1 \succ \dots \succ \tau^K$ . We make no assumption about the source of preference. It may come from human ranking or some task-specific evaluation metric  $\mathcal{R}$  on the full text-sequences, e.g., the descending ordering of the text-sequence evaluations  $\mathcal{R}(s_{T^1}^1, y) > \dots > \mathcal{R}(s_{T^K}^K, y)$ .
For a trajectory  $\tau^k = \{(s_0^k, a_0^k), \ldots, (s_{T^k - 1}^k, a_{T^k - 1}^k)\}$ , our desired token-level reward function  $r_\phi$  generates a reward for each step as  $\{r_\phi(s_t^k, a_t^k)\}_{t=0}^{T^k-1}$ . A sequence-level evaluation  $e_\phi(\tau^k)$  for trajectory  $\tau^k$  can be obtained by  $e_\phi(\tau^k) = f(\{r_\phi(s_t^k, a_t^k)\}_{t=0}^{T^k-1})$ , where  $f(\cdot)$  is the aggregation function over all per-step rewards, e.g., the classical summation  $\sum(\cdot)$ . Our goal is to train  $r_\phi$  such that these parametrized sequence-level evaluations  $\{e_\phi(\tau^k)\}_{k=1}^K$  align with the given preference ordering  $\tau^1 \succ \dots \succ \tau^K$ . Through this, the sequence-level preference is grounded into token-level rewards  $r_\phi$ .
Under the Plackett-Luce choice model [40, 41], the parametrized sequence evaluations  $\{e_{\phi}(\tau^k)\}_{k=1}^K$  induce a probability distribution over all possible permutations of the integers  $\{1, \ldots, K\}$ . We want to maximize the likelihood of the given preference ordering  $\operatorname{ord} = (1, \ldots, K)$ , i.e.,
$$
\min  _ {\phi} \mathcal {L} (\phi) =: - \log P \left(\operatorname {o r d} \mid \left\{e _ {\phi} \left(\tau^ {k}\right) \right\} _ {k = 1} ^ {K}\right), P \left(\operatorname {o r d} \mid \left\{e _ {\phi} \left(\tau^ {k}\right) \right\} _ {k = 1} ^ {K}\right) = \prod_ {k = 1} ^ {K} \left\{\exp \left(e _ {\phi} \left(\tau^ {k}\right)\right) / \sum_ {i = k} ^ {K} \exp \left(e _ {\phi} \left(\tau^ {i}\right)\right) \right\}. \tag {2}
$$
Algorithm 1 A learning routine for the preference-grounded token-level reward function  $r_{\phi}$ .
Input: The LM  $\pi_{\theta}$ , initialized reward  $r_{\phi}$ , aggregation function  $f(\cdot)$ , reward-training steps  $M_{\mathrm{rew}}$ . for iter  $\in \{1, \ldots, M_{\mathrm{rew}}\}$  do Use  $\pi_{\theta}$  to generate  $K$  sequences  $\{\pmb{a}^k\}_{k=1}^K$ ; and get the preference ordering among  $\{\pmb{a}^k\}_{k=1}^K$ . With  $f(\cdot)$ , get sequence evaluations  $\{e_{\phi}(\tau^k)\}_{k=1}^K$  from  $r_{\phi}$ ; and optimize  $r_{\phi}$  by Eq. (2). end for
When  $K = 2$  and  $f(\cdot)$  denotes summation, Eq. (2) reduces to the classical pairwise-preference loss in Eq. (1). Therefore, our reward-learning loss can be viewed as an extension of the classical pairwise loss. Further, Eq. (2) extends the ListMLE loss [42] in recommender systems into preference learning under multiple variable-length trajectories.
Algo. 1 summarizes our reward-learning framework by describing an online-learning routine for training  $r_{\phi}$ . An offline or hybrid version can be obtained with minor changes.
The Choice of Aggregation Function  $f(\cdot)$ . In classical IL tasks such as robotics [43], the robots are trained to stand or walk as long as possible. In this scenario, summation is a natural choice for the aggregation function  $f(\cdot)$ . However, in many text generation tasks, such as summarization, the generation quality may not be directly associated with the length of the generated text sequence. Nevertheless, suppose the token-level rewards are positive (i.e.,  $r_{\phi} > 0$ ), a longer sequence naturally has a higher sum of per-step rewards than a shorter one, which can bias  $r_{\phi}$  towards automatically ranking longer sequences higher. This bias can hinder our reward-learning goal of aligning  $\{e_{\phi}(\tau^{k})\}_{k = 1}^{K}$  with the given preference ordering. A naive numeric example is additionally provided in Appendix C.
To mitigate the potential length bias in the classical summation, we discuss three alternative choices of the aggregation function  $f(\cdot)$ : average, soft maximum, and soft minimum.
Average. We define the average-aggregated sequence-level evaluation  $e_{\phi}^{\mathrm{avg}}(\tau^{k})$  for trajectory  $\tau^{k}$  as
$$
e _ {\phi} ^ {\text {a v g}} \left(\tau^ {k}\right) = \frac {C}{T ^ {k}} \sum_ {t = 0} ^ {T ^ {k} - 1} r _ {\phi} \left(s _ {t} ^ {k}, a _ {t} ^ {k}\right), \quad C = \frac {1}{K} \sum_ {k = 1} ^ {K} T ^ {k}, \tag {3}
$$
where  $C$  is the average length of the  $K$  sequences. Multiplied by the average length  $C$  has the benefit of scaling  $e_{\phi}^{\mathrm{avg}}$  to the scale of  $e_{\phi}^{\mathrm{sum}}$ , which ensures numerical-scale consistency with  $e_{\phi}^{\mathrm{sum}}$  and thus reduces hyperparameter tuning when switching from summation to average aggregation.
Soft Maximum. We define the soft-maximum-aggregated sequence-level evaluation  $e_{\phi}^{\max}(\tau^{k})$  as
$$
e _ {\phi} ^ {\max } \left(\tau^ {k}\right) = C \times \beta \cdot \log \left[ \sum_ {t = 0} ^ {T ^ {k} - 1} \exp \left(r _ {\phi} \left(s _ {t} ^ {k}, a _ {t} ^ {k}\right) / \beta\right) \right], \tag {4}
$$
where  $C$  is the average trajectory-length in Eq. (3) and  $\beta$  is the temperature parameter.
Soft Minimum. The soft-minimum-aggregated sequence-level evaluation  $e_{\phi}^{\min}(\tau^{k})$  follows Eq. (4) except for changing  $\beta$  to  $-\beta$ .
2.2 LM Training with Preference-grounded Token-level Guidance
Considering the supervised-data availability, we present two minimalist LM training objectives that utilize the learned preference-grounded guidance: 1) a REINFORCE-style update when there is no supervised data; 2) reward-weighted MLE when there are sufficient data. Our LM training directly starts from raw pre-trained checkpoints, without task-specific supervised pre-training. This choice is to keep the algorithm general and consistent in both situations we consider, since task-specific pre-training may not be feasible in the setting of few/zero-shot learning.
As shown in Algo. 1, we train the reward function  $r_{\phi}$  by the sequences sampled from LM  $\pi_{\theta}$ . Since task-specific pre-training to  $\pi_{\theta}$  is not assumed, over the course of training,  $\pi_{\theta}$  itself can evolve from a less-preferred distribution to a highly-preferred one. To mitigate the impact of this distribution shift and keep  $r_{\phi}$  as accurate guidance for LM training, we periodically re-estimate  $r_{\phi}$  during the first half of the LM-training process<sup>2</sup>, motivated by recent works in model-based RL [44-47].
2In our preliminary study, we observed that this choice ("retraining the reward model only during the first half of the LM-training process") can save about  $25 - 30\%$  compute without hurting the performance much, compared to the vanilla reward retraining ("retraining the reward model throughout the entire LM-training process").
Algorithm 2 An alternate-learning process for the reward function  $r_{\phi}$  and the LM  $\pi_{\theta}$ .
Input: The dataset  $\mathcal{D}$ , initialized LM  $\pi_{\theta}$ , initialized reward function  $r_{\phi}$ , LM-training steps  $M_{\mathrm{LM}}$ , reward-retrain period  $M_{\mathrm{re}}$ , all inputs for training the reward function specified in Algo. 1.
Initialize  $r_{\phi}$  by Algo. 1.
for iter  $\in \{1,\dots ,M_{\mathrm{LM}}\}$  do
ifiter  $\leq M_{\mathrm{LM}} / 2$  anditer  $\% M_{\mathrm{re}} = = 0$  then
Re-train  $r_{\phi}$  by Algo. 1 without re-initialization.
Optimize  $\pi_{\theta}$  by Eq. (5) or Eq. (6) with  $\mathcal{D}$  and  $r_{\phi}$ .
Without Supervised Data. When the LM  $\pi_{\theta}$  needs to discover good text generations by itself, the learned token-level reward  $r_{\phi}$  can be used to provide dense guidance on generating each token, i.e., given the generation context  $s_t$ , select the next token  $a_t$  such that  $r_{\phi}(s_t, a_t)$  is high. Intuitively, for a generation trajectory  $\tau$ , if  $\forall (s_t, a_t) \in \tau, r_{\phi}(s_t, a_t)$  is high, then the corresponding sequence-level evaluation  $e_{\phi}(\tau) = f(\{r_{\phi}(s_t, a_t)\}_{t=0}^{T-1})$  can be also high, e.g., the average or summation of token-level rewards. The associated text sequence  $\pmb{a}$  will thus be preferable since  $r_{\phi}$  is trained to reflect the sequence-level preference (Section 2.1). Through  $r_{\phi}$ , the sequence-level preference is grounded into dense token-level guidance for LM training, without granularity mismatch or feedback delay.
With the learned  $r_{\phi}$ , a minimalist implementation of this LM-training idea is the discrete-variable optimization problem
$$
\max  _ {\theta} \mathbb {E} _ {t \sim \text {U n i f o r m} \{0, \dots , T - 1 \}} \mathbb {E} _ {a _ {t} \sim \pi_ {\theta} (\cdot | s _ {t})} [ r _ {\phi} (s _ {t}, a _ {t}) ],
$$
for each timestep  $t$  of which we calculate its gradient by the classical REINFORCE method [48-50] since it can cope with a large vocabulary size. Here,  $T$  denotes a generic sequence length. Additionally, since we want multiple text generations in typical LM tasks, instead of only one, we relax the convergence of the REINFORCE method by adding a standard max-entropy gradient, which can help capture multiple good behavior-modes [51-53]. Thus, the LM  $\pi_{\theta}$  is trained by the gradient
$$
\mathbb {E} _ {t \sim \text {U n i f o r m} \{0, \dots , T - 1 \}} \left\{\mathbb {E} _ {a _ {t} \sim \pi_ {\theta} (\cdot \mid s _ {t})} \left[ r _ {\phi} \left(s _ {t}, a _ {t}\right) \cdot \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \right] + \alpha \cdot \nabla_ {\theta} \mathcal {H} \left(\pi_ {\theta} (\cdot \mid s _ {t})\right) \right\}, \tag {5}
$$
where  $\mathcal{H}(\pi_{\theta}(\cdot |s_t))$  is the Shannon entropy of  $\pi_{\theta}(\cdot |s_t)$  and  $\alpha$  is a balancing coefficient.
With Supervised Data. With a labelled dataset  $\mathcal{D} = \{(x^i, y^i)\}_{i=1}^N$  and with the learned preference-grounded guidance  $r_\phi$ , a minimalist enhancement of the classical MLE LM-training is the token-level weighted-MLE, where the per-token weight is given by the learned reward-function  $r_\phi$ . Our intention is to emphasize the important tokens in the given sequence  $y$  and downweight the unimportant ones, where the token importance given by  $r_\phi$  grounds the sequence-level preference. Intuitively, this weighting scheme can better utilize the LM capacity and the optimization budget, and may thus improve upon the vanilla supervised loss [54, 16]. Specifically, the LM  $\pi_\theta$  is trained by
$$
\min  _ {\theta} - \mathbb {E} _ {(x, y) \sim \mathcal {D}} \left[ \sum_ {t = 0} ^ {| y | - 1} w _ {t} \cdot \log \pi_ {\theta} \left(y _ {t} \mid s _ {t}\right) \right], \text {w i t h} s _ {t} = (x, y _ {<   t}) \text {a n d} w _ {t} = \frac {r _ {\phi} \left(s _ {t} , y _ {t}\right)}{\sum_ {t ^ {\prime} = 0} ^ {| y | - 1} r _ {\phi} \left(s _ {t ^ {\prime}} , y _ {t ^ {\prime}}\right)}, \tag {6}
$$
where  $|y|$  is the length of the target sequence  $y$  and  $w_{t}$  is the self-normalized token-level reward. The standard self-normalization is used to reduce the gradient variance among the samples in  $\mathcal{D}$ .
Algo. 2 sums up the entire alternate-learning process, with the reward-learning routine in Algo. 1.
Guiding Signals for LM Training. One string of works in LM training directly optimizes the LMs against the native sequence-level feedback such as the test-time metric [e.g., 3, 55-59, 32]. This choice, however, may directly suffer from the delayed-feedback issue discussed in Section 1 and the subsequent high gradient variance and low sample efficiency [31, 32]. In the recent trend of RL-based LM training, it has been common to incorporate a token-level KL penalty towards the uniform distribution [31, 60], the initial LM [23, 61], the supervised-fine-tuned model [12, 62-64, 14], or the base momentum model [65], to add to the delayed/ungrounded feedback. Although that KL penalty does impact the RL-based LM training at the token level, it is not tailored to the concrete task or the desired sequence-level feedback. When combined with the delayed-feedback issue, it
could distract the LM training from improving the received feedback/evaluation, especially at the beginning of the text-sequence generation, which can however affect all subsequent token selections. By contrast, as seen in Eq. (5), even when added a max-entropy gradient, our preference-grounded token-level guidance can still provide dense, task-specific, and feedback-oriented guidance on the selection of each token. For a more detailed discussion on the RL formulation of LM generation, the delayed-feedback issue in RL-based LM training, and delayed feedback with KL penalty, please refer to Appendix F.