SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Prompt Generation. Prior works [e.g., 6, 10, 85, 97] have shown that manual prompts can steer LMs to perform NLP tasks in the few/zero-shot setting. In general, prompts can be discrete, consisting of real token-strings; or can be continuous, where the prompts are entirely free word-embeddings that do not map to real tokens. Several works [e.g., 98-102, 83] tune continuous soft prompts using gradient descent, which typically requires some expensive gradient information [78, 103]. In this work, we apply our framework to the task of input-agnostic discrete-prompt optimization due to its challenging setting, better human understandability of the learned prompts [104, 105], potential transferability across LMs [106, 76, 60], and more robustness in the low-data regime [99]. Recent works propose some new settings such as input-dependent prompt-tuning [88], which are potential further applications of our framework and are left for future work.
Text Summarization. Apart from using RL techniques discussed in Sections 3, prior works on text summarization [e.g., 7, 89, 107-109] mainly focus on structural designs of the LMs and improvements on the source of the (pre-)training data, where the LMs are typically trained by vanilla MLE on the supervised data. In this paper, we apply our preferred-grounded token-level guidance to this task by considering a weighted-MLE objective for LM training. The weights given by the learned reward function reflect some sequence-level preference among multiple candidate summaries. Our framework thus has the potential to learn and improve from lower-quality data, and generate summaries fulfilling more general evaluation metrics, such as human preference.
Weighted MLE in NLP. Though not a very common technique, the approach of weighted MLE has been adopted in prior NLP research. For example, RAML [110] samples outputs proportionally to its exponentiated scaled "reward" (negative edit/Hamming distance) using stratified sampling. GOLD [111] frames text generation as an offline RL problem with expert demos and learns from the demos by importance weighting, where training examples with higher probability under the model are weighted higher. Besides, Ghosh et al. [112] apply the weighted MLE technique to table-to-text generation and Junczys-Dowmunt et al. [113] apply this technique to grammatical error correction for machine translation. Our token-level reward-weighted MLE in Section 2.2 adds to this research thread by emphasizing the important tokens in the supervised sequences and downweighting the unimportant tokens. This design may better utilize the LM capacity and the optimization budget. The efficacy of our reward-weighted MLE is experimentally verified in Section 4.2.
Align LMs with Preference. Similar to our paper, prior works on aligning LMs with preference typically focus on adjusting the pretrained LMs, where preference comes from human feedback or from some automatic metrics. A classical strategy is to add external filters on top of the pretrained LMs for the generated text sequences or for the training sequences [e.g., 19], where the LMs are trained using MLE on abundant supervised data. Another classical approach finetunes LMs using supervised learning (vanilla MLE) on some curated/improved datasets [20-22], or on massive highly-curated collections of tasks phrased as instructions for supervised finetuning the LMs [114-117]. Apart from supervised learning, reinforcement learning techniques have also been applied to learn from human feedback (RLHF). Similar to the discussion in Section 3, these works typically learn a sequence-level classifier that predicts human (pairwise) preferences and during LM training add a general-purpose KL penalty that is less-targeted to the specific LM task and feedback (preference, metric scores, etc.) [e.g., 14, 23-25], such as a token-level KL penalty towards the initial LM prior to training.
Alternatively, the divergence of the LMs from a target distribution can also be used as the finetuning objectives. This line of research [e.g., 118-120] formalizes controlled text generation as a constraint satisfaction problem over LM's probability distribution, with an additional divergence-minimization objective that the LMs should have a minimal KL- or  $f$ -divergence from the original pretrained model. These approaches, however, require explicit functional specification on the constraints or on the human preference, rather a more vague form of (binary) comparison between LM samples. For example, Go et al. [120] consider human preference as a probability distribution measuring how well the generated text-sequence satisfies the preference. Apart from this more demanding requirement, these approaches further require special methods to sample from the resulting LM.
To sum up, prior works on aligning LMs with preference mostly focus on an ungrounded sequence-level guidance, which can suffer from the delay-feedback issue in LM training, as discussed in Sections 1 and 3. By contrast, our preference-grounding perspective can provide a stable, data-driven, task-specific token-level guidance on LM training, and can potentially improve on the vanilla MLE, especially when the quality of the supervised data cannot be guaranteed. We experimentally validate this intuition in Section 4 and Appendix A.2.
Apart from fine-tuning the pretrained LMs, Korbak et al. [121] recently apply preference alignment to the pre-training stage of the LMs. As with prior works, the sparse sequence-level evaluation (without KL penalty/stabilizer) is directly used, to learn a token-level value function, to condition the LM generation on, or for a reward-weighted regression objective. The pre-training stage in Korbak et al. [121] is a potential further application of our framework since we make no assumption on the zero-shot ability of the initialized LMs, as discussed in Sections 2.2 and 4.3.
We also notice that a recent robotics paper [122] proposes to learn a weighted-sum aggregation together with the per-step reward, to form the sequence-level evaluation in learning the reward function, based on pairwise preference over two trajectories of equal length. Compared with this recent work, our aggregation functions in Section 2.1 do not require additional modeling and training, and therefore can be more efficient and more stable for the reward-function learning. Additionally, we do not assume that trajectory lengths are equal, as this may be infeasible for LM tasks such as text summarization. Furthermore, our framework allows utilizing the preference among more than two trajectories, rather than the classical pairwise preference. In this particular aspect, our framework can be more general than this recent work of Kim et al. [122].
F A Discussion on Applying RL Methods to LM Tasks
F.1 LM Generation as a Token-level MDP
In most LM generation tasks, there is a dataset  $\mathcal{D} = \{(x^i,y^i)\}_{i = 1}^N$  of  $N$  supervised examples, where  $x$  is the input to the LM that can be a dummy, and  $y\in \mathcal{V}$  is the target text sequence. Viewing the LM as a token-level RL policy, LM generation can be formulated as a sequential decision-making problem, specified by the Markov Decision Process (MDP)  $\mathcal{M} = (\mathbb{S},\mathbb{A},P,\mathcal{R},\gamma ,\mu_0)$  [123]. Specifically,  $\mathbb{S}$  is the state space, where the state at timestep  $t$ ,  $s_t$ , consists of the LM input  $x$  and the previously generated tokens  $a_{< t} = (a_0,\dots,a_{t - 1})$ ,  $t > 0$ , i.e.,  $s_0 = x$  and  $\forall t > 0,s_t = (x,a_{< t})$ .  $\mathbb{A}$  is the action space, which is the vocabulary  $\nu$ , and an action  $a_{t}$  at timestep  $t\geq 0$  is a token from  $\nu$ .  $P(s_{t},a_{t}): \mathbb{S}\times \mathbb{A}\to \mathbb{S}$  is the transition function that deterministically appends the newly sampled token to the end of the current state, i.e.,  $\forall t\geq 0,s_{t + 1} = (s_t,a_t) = (x,a_{\leq t})$ .  $\mathcal{R}(s_T,y):\mathbb{S}\times \mathcal{V}\rightarrow \mathbb{R}$  is the environmental reward (task-specific evaluation metric) that depends on the final state  $s_T$  of the LM-generation trajectory and the target sequence  $y$ . Here  $T$  is the ending time of the trajectory, i.e., the length of the full generated text sequence; and  $s_T = (x,a_0,\ldots ,a_{T - 1})$  is the final state of the generation trajectory consisting of the LM input  $x$  and the full generated text sequence  $\pmb {a} = (a_0,\dots,a_{T - 1})$ .  $\gamma \in [0,1]$  is the discount factor. And  $\mu_0(x):\mathbb{S}\to [0,1]$  is the distribution of the initial input  $x$ .
We denote the LM as  $\pi_{\theta}(a_t|s_t)$ , parametrized by  $\theta$ . At each timestep  $t$ ,  $\pi_{\theta}(a_t|s_t)$  generates the next token  $a_t$  given the current state  $s_t = (x, a_{<t})$ . The ultimate goal of policy learning (LM training) is to maximize the expected environmental reward  $\mathcal{R}$ , which can be expressed as
$$
\max _ {\boldsymbol {\theta}} \mathbb {E} _ {(x, y)} \mathbb {E} _ {\boldsymbol {a} \sim \prod_ {t = 0} ^ {T - 1} \pi_ {\boldsymbol {\theta}} (a _ {t} | s _ {t})} \left[ \mathcal {R} (s _ {T} = (x, \boldsymbol {a}), y) \right],
$$
where  $(x,y)$  is drawn from the corresponding sampling distribution.
F.2 Delayed Feedback in RL-based LM Training
As discussed in Appendix F.1, the environmental reward  $\mathcal{R}(s_T,y)$  is only defined on the full generated text sequence  $a$ . The token-level MDP formulation of LM generation thus meets the problem of sparse reward-signal or the delayed feedback issue discussed in Section 1. Hereafter, we will use "sparse reward (signal)" and "delayed feedback" interchangeably depending on the context, as they are used synonymously in the RL literature.
Specifically, prior works [e.g., 31, 60, 32] often manually interpolate the intermediate rewards by some non-informative values such as 0 or  $-1$ , i.e.,  $\forall t \geq 0$
$$
\mathcal {R} \left(s _ {t}, y\right) = \left\{ \begin{array}{l l} 0 \text {o r} - 1, & t <   T \\ \mathcal {R} \left(s _ {T}, y\right), & t = T \end{array} \right. \tag {7}
$$
It is clear that the reward signal is sparse. In other words, the feedback to intermediate actions/tokens is delayed until the full text-sequence has been generated.
We note that this sparse-reward/delayed-feedback problem will not be addressed by the standard actor-critic or Q-learning methods in RL. With only sparse reward-signals, it can be difficult to estimate the token-level value functions in these RL methods.
Specifically, the standard Monte Carlo estimate of the value functions is known to have high variance due to the large sampling space [123]. This problem is even severe in the LM tasks where there are exponentially many text sequences that can follow a partial sequence.
Further, as discussed in Guo et al. [31], the sparse-reward/delayed-feedback problem can also hurt the bootstrapping-style method for learning the value functions, since the standard value-function learning can suffer from "the unstable per-step bootstrapping-style training with sparse reward signals." This can subsequently harm the LM training since many actor-critic or Q-learning methods rely heavily on how accurately the learned value function(s) can assess the quality of intermediate text sequences [31, 123, 124].
F.3 Sparse Reward with KL Penalty
With the sparse-reward/delayed-feedback issue in Appendix F.2, prior works typically add a token-level KL-penalty to the sparse sequence-level environmental rewards Eq. (7). For simplicity, assume that in Eq. (7) the intermediate rewards are interpolated by 0. The KL-stabilized reward signal  $R(s_{t},a_{t},y)$  is
$$
R \left(s _ {t}, a _ {t}, y\right) = \left\{ \begin{array}{l l} - c \cdot \operatorname {K L} \left(\pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \mid \mid \pi_ {0} \left(a _ {t} \mid s _ {t}\right)\right), & t <   T - 1 \\ \mathcal {R} \left(s _ {T}, y\right) - c \cdot \operatorname {K L} \left(\pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \mid \mid \pi_ {0} \left(a _ {t} \mid s _ {t}\right)\right), & t = T - 1 \end{array} , \right. \tag {8}
$$
where  $c$  is a hyper-parameter and  $\pi_0$  is some prior distribution, such as the uniform distribution [31, 60], the initial LM prior to training [23, 61], the supervised-fine-tuned model [62, 63, 12, 14], or the base momentum model [65]. For a concrete example, see Line 224-235 of the popular trlx package's implementation.
With this KL-stabilized reward signal  $R(s_{t},a_{t},y)$ , the action-value function for the policy/LM  $\pi_{\theta}$  is
$$
\begin{array}{l} Q \left(s _ {t}, a _ {t}, y\right) = \mathbb {E} _ {\left\{a _ {t ^ {\prime}} \right\} _ {t ^ {\prime} = t + 1} ^ {T - 1} \sim \pi_ {\theta}} \left[ \sum_ {t ^ {\prime} = t} ^ {T - 1} \gamma^ {t ^ {\prime} - t} R \left(s _ {t ^ {\prime}}, a _ {t ^ {\prime}}, y\right) \mid s _ {t}, a _ {t} \right] \tag {9} \\ = \mathbb {E} _ {\left\{a _ {t ^ {\prime}} \right\} _ {t ^ {\prime} = t + 1} ^ {T - 1} \sim \pi_ {\theta}} \left[ \gamma^ {T - 1 - t} \mathcal {R} \left(s _ {T}, y\right) - c \cdot \sum_ {t ^ {\prime} = t} ^ {T - 1} \gamma^ {t ^ {\prime} - t} \mathrm {K L} \left(\pi_ {\theta} \left(a _ {t ^ {\prime}} \mid s _ {t ^ {\prime}}\right) \mid \mid \pi_ {0} \left(a _ {t ^ {\prime}} \mid s _ {t ^ {\prime}}\right)\right) \mid s _ {t}, a _ {t} \right] \\ \end{array}
$$
It is clear from Eq. (9) that the environmental reward  $\mathcal{R}(s_T, y)$  is multiplied by a factor exponentially decayed with respect to the length of the remaining horizon  $T - 1 - t$ . Without the KL penalty, the action-value  $Q(s_t, a_t, y)$  could be tiny when  $t$  is small, i.e., at the beginning of the text-sequence generation. This could make it hard to accurately model and learn the action values, echoing the previously-stated harm of the sparse-reward/delayed-feedback problem mentioned by Guo et al. [31]
Recall that the standard actor-critic and Q-learning methods in RL use the action-value function  $Q(s_{t},a_{t},y)$  as the token-level guidance (per-step critic) for policy/LM training. Due to the exponentially decaying factor  $\gamma^{T - 1 - t}$ , when the discount factor  $\gamma$  in Eq. (9) is not sufficiently large, this
token-level guidance  $Q(s_{t},a_{t},y)$  in RL-based LM training mainly reflects the (discounted) sum of future KL-penalty, rather than the actual goal of LM training — the environmental reward  $\mathcal{R}(s_T,y)$ . This phenomenon can be more evident at the beginning of the text-sequence generation, i.e., when the length of the remaining horizon  $T - 1 - t$  is long. On the other hand, learning the action-value function  $Q(s_{t},a_{t},y)$  under a large discount factor  $\gamma$  is known to be challenging [123], since the highly varying (late) future can significantly affect the current action value  $Q(s_{t},a_{t},y)$ . The selection of the discount factor  $\gamma$ , therefore, becomes a tradeoff and a challenge. Note that  $\mathcal{R}(s_T,y)$  here is generic and can represent automatic evaluation metrics or (human) preference, and that the beginning of text generation can affect all subsequent token selections. Intuitively, using Eq. (9) as the token-level guidance for policy/LM training can thus be less successful in the concrete LM task, especially when generating longer sequences, as we verified in Appendix A.2.
In the experiments (Section 4 and Appendix A.2), we compare our preference-grounding approach with RL-based baselines that estimate a standard value function similar to Eq. (9) from sparse environmental reward with KL penalty, such as the RLPrompt method [60] and the (supervised+) PPO/NLPO methods in RL4LMs [61]. We leave as future work the potential combination of our preference-grounded guidance with actor-critic and Q-learning methods in RL-based LM training.
G Further Discussion on the Guidance Re-estimation Scheme
As discussed in Section 2.2, in this paper, we deal with the most general setting where the LM training directly starts from a raw pre-trained LM, rather than an initial LM that has been fine-tuned via supervised learning on the desired dataset, such as in Stiannon et al. [12]. We also make no assumptions about the zero-shot ability of the raw pre-trained LM. We choose this setting because it is more general and naturally fits into the task of text-prompt generation, where supervised datasets of good prompts are not available and the initial LM cannot generate good prompts.
As discussed before, under this general setting, the LM  $\pi_{\theta}$  can evolve from a less-preferred distribution to a highly-preferred one, over the training process. Since our reward function  $r_{\phi}$  is trained by text sequences sampled from  $\pi_{\theta}$ , there is a distribution shift between the sequences used to train  $r_{\phi}$  during reward-function learning, and the sequences evaluated by  $r_{\phi}$  during LM training, especially after  $\pi_{\theta}$  has been sufficiently improved. To keep  $r_{\phi}$  as accurate guidance for LM training, a natural idea is to refine  $r_{\phi}$  periodically on the text generations from the latest LM, leading to our reward-function retraining scheme.
We emphasize that the reward-function retraining scheme does not give our method an unfair advantage over the baseline methods. In particular, RLPrompt [60] and RL4LMs' methods [61] retrain their value-functions in every optimization step, and thus, they query the environmental reward in every optimization step. Specifically, in Algorithm 1 of the RL4LMs paper, the penalized reward  $\hat{R}_t$  is calculated in each optimization step, whose calculation requires the true environmental reward  $R$  (Eq. (1) of the RL4LMs paper). Besides, in the codebase of RLPrompt, this environmental interaction is implemented in this line, which is queried in every optimization step, as seen in this line. In the notion of Reinforcement Learning from Human Feedback (RLHF), this every-step interaction is similar to asking humans to score the LM generations in every training step, which can be infeasible. By contrast, in our paper, we reduce the frequency of these environmental interactions by retraining the guidance model only periodically and only during the first half of the LM-training process.
Though the motivation of this reward-function retraining scheme comes from model-based RL (Section 2.2), we notice that some prior RLHF works do implement similar ideas. For example, Page 2 of Ziegler et al. [23] mentions that "... we continue to collect additional data and retrain our reward model as the policy improves (online data collection)." Page 2 of Stiennon et al. [12] mentions that "We can then gather more human data using samples from the resulting policy, and repeat the process." Page 5 of Menick et al. [25] and Page 20 of Bai et al. [24] also have similar discussions. Based on these, our reward-function retraining scheme is both well-motivated and practical, even with human rankings in RLHF.
H Potential Negative Societal Impacts
Since our framework can ground the sequence-level preference into token-level guidance for LM training and can be not tied to a specific preference source, it is possible that this framework may be
used to train ill-intended LMs by grounding some malicious or unethical preferences. This potential negative impact may be mitigated by closer monitoring the datasets on which our framework operates.
Since our token-level guidance is learned by grounding sequence-level preference, a potential failure case of our framework will be when the preference orderings are very noisy. In this situation, the learned guidance may not be meaningful and hence could even deteriorate the subsequent utilization of it in LM training.
Even though we have shown in Section 4.3 that it can be beneficial to use more than two sequences to learn the token-level guidance, it can be practically challenging to obtain a high-quality ranking among many candidate text sequences, e.g., when the number of sequences is more than seven.
Besides, the reward-function retraining scheme may incur some additional computational complexity, compared with training the reward function only once and fixing it throughout the LM-training process.
J Computational Resources
The experiments are conducted on NVIDIA GeForce RTX 3090 and NVIDIA A100 GPUs. Depending on the specific task and setting, several models could be trained concurrently on a single GPU.