SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In some relatively "ideal" settings, prior works have attempted to learn task-specific token-level guidance for LM training. For instance, Shi et al. [66] use inverse RL, Guo et al. [67] propose a hierarchical approach, and Yang et al. [68] learn LM discriminators; but these methods require abundant expert data for supervised (pre-)training, making them infeasible for the few/zero-shot settings we consider. Under the same requirement of sufficient expert data, Lin et al. [69] learn a sequence-level adversarial-ranking reward and Yu et al. [70] train a GAN structure. They both use Monte-Carlo rollouts to simulate intermediate rewards, which can be computationally expensive and have high variance. Le et al. [71] use some values related to the sequence evaluation without explicitly learning per-token rewards. Pang et al. [72] learn a token-level error predictor for machine translation, but they rely on expert error-span annotations for each translation, which is highly demanding.
By contrast, we propose a versatile framework for learning task-specific token-level guidance for LM training that can ground the sequence-level preference. Our approach is not limited to standard LM tasks and is also suitable for the low-data regime, with few assumptions about expert-data availability or preference source. In our experiments, we compare our method to recent RL-based approaches that train LM under delayed/ungrounded feedback with KL penalty. We discuss additional related works on prompt generation, text summarization, and aligning LMs with preferences in Appendix E.
We test our framework on two distinct representative text-sequence generation tasks: 1) input-agnostic discrete text-prompt generation for few-shot text-classification (Section 4.1), 2) the classical text summarization (Section 4.2). Our LM training directly starts from raw pre-trained checkpoints from HuggingFace [73], without task-specific supervised pre-training. Depending on the LM  $\pi_{\theta}$  used in the specific task, our reward function  $r_{\phi}$  can be implemented as either a decoder-only or an encoder-decoder model. Similar to prior works [e.g., 32, 25], given a text sequence  $\pmb{a}$  and an LM input  $x$ , the causal mask in transformers enables us to get the learned guidance  $r_{\phi}(s_t, a_t)$  at each step of the sequence in parallel. Source codes have been publicly released.
4.1 Input-agnostic Discrete-prompt Generation
Overview. In discrete text-prompt generation [e.g., 10, 74], we input a discrete text-prompt  $\pmb{a}$  and an observation sequence  $o$  to a large pre-trained downstream LM  $\pi_{\mathrm{DLM}}(\cdot | \pmb{a}, o)$  to directly classify text  $o$ , without finetuning  $\pi_{\mathrm{DLM}}$ . We follow the classical setting [e.g., 75, 60] to perform classification by selecting tokens corresponding to some predefined class labels. In our input-agnostic setting, the generated prompt is independent of the observation  $o$ . During inference time, only the learned prompts are used and the LM  $\pi_{\theta}$  is discarded. The initial input  $x$  to  $\pi_{\theta}$  is a dummy, and the target  $y$  is the class label. We also adopt the standard few-shot setting [76], where both the training and validation sets have 16 (o, y)-pairs per class. With a fixed length  $T$ , the goal is to find discrete text-prompts  $\pmb{a} = (a_0, \dots, a_{T - 1})$  that have high test accuracy. We simulate the sequence-level preference by the stepwise metric in Deng et al. [60], i.e., the higher value the better prompt. This choice ensures a fair comparison and avoids a potential overfitting — training and testing the LM on the same evaluation metric "accuracy". Appendix D discusses more details about the prompt task.
LM Training, Implementation, and Datasets. Since the prompt-generation task does not assume the availability of supervised data — the ground-truth prompts, the LM  $\pi_{\theta}$  is trained by the REINFORCE-style update in Section 2.2 to discover highly-accurate prompts by itself. We implement our framework on the codebase of RLPrompt [60], and adopt the standard datasets and most hyperparameter settings in it. Reward training is reconstituted every 1000 steps during the first 6000 steps of the LM training process and has early stopping. Reward function is learned with 5 sampled sequences and the temperature in Eq. (4) is set as  $\beta = 2$ . The coefficient  $\alpha$  in Eq. (5) is  $\alpha = 2^{-3}$ . Appendix A.2 discusses the choices of these hyperparameters. The length of the generated prompts is fixed at 5. We
Table 1: Test accuracy on the prompt task. Best overall result is bold and best discrete-prompt result is underlined if different. The reported results are mean (standard deviation). We denote "BB Tuning-50" for Black-Box Tuning with mixed discrete and soft prompts that tunes the 50 soft tokens; and "AVG", "SUM", "MIN", "MAX" for our method with aggregation function average, summation, soft minimum, and soft maximum (Section 2.1).
SST-2	Yelp P.	AG News
Finetuning	Few-shot Finetuning	80.6 (3.9)	88.7 (4.7)	84.9 (3.6)
Continuous Prompt	Soft Prompt Tuning [83]	73.8 (10.9)	88.6 (2.1)	82.6 (0.9)
BB Tuning-50 [78]	89.1 (0.9)	93.2 (0.5)	83.5 (0.9)
AutoPrompt [84]	75.0 (7.6)	79.8 (8.3)	65.7 (1.9)
Discrete Prompt	Manual Prompt [85]	82.8	83.0	76.9
In-Context Demo [10]	85.9 (0.7)	89.6 (0.4)	74.9 (0.8)
Instructions [86]	89.0	84.4	54.8
GrIPS [87]	87.1 (1.5)	88.2 (0.1)	65.4 (9.8)
RLPrompt [60]	90.5 (1.5)	94.2 (0.7)	79.7 (2.1)
Ours (AVG / SUM)	92.6 (1.7)	94.7 (0.6)	82.8 (1.5)
Ours (MIN)	91.9 (1.8)	94.4 (0.8)	82.4 (1.1)
Ours (MAX)	91.2 (2.5)	94.8 (0.5)	83.3 (1.4)
test on three popular few-shot datasets in prior work [e.g., 77, 78]: two sentiment binary-classification datasets SST-2 [79, 80] and Yelp Polarity [81], and a topic four-way-classification dataset AG News [81, 82]. Additional details on the experiment and datasets are provided in Appendix B.1.
Results. We compare three variants of our framework with finetuning and with baselines in discrete- and continuous-prompt generation. Since the generated prompts all have length 5, in this task, the average aggregation is equivalent to summation. Table 1 shows the test accuracy, where we rerun the codebase of RLPrompt [60] under the same random seeds and evaluation script as our method. Other baseline results are from the literature [60, 88].
On all three tested datasets, our method shows competitive and stable results against the strong baselines not only in discrete-prompt generation, but also in heavier continuous-prompt tuning and finetuning the large downstream LM. Based on Section 3, the performance improvement achieved by our method compared to RLPrompt suggests that utilizing the token-level guidance learned by our approach, which grounds the task-specific preference, can be more effective than learning under delayed/ungrounded feedback with KL penalty. Further, on both Yelp P. and AG News, using MAX aggregation is better than the classical summation. Table 3 in Appendix A shows examples of good generated prompts and their test accuracy. For instance, high-quality prompts on the AG News dataset often contain a topic classification keyword, such as "Tags" and "Category". This aligns with our intuition that good prompts may be identified by a (few) "key" token(s), as discussed in Sections 1 and 2.1. Thus, the (soft-)maximum aggregation may better reflect prompt quality than summation.
4.2 Text Summarization
Overview. In the summarization task, we follow the standard setting [e.g., 89, 61], where a set of supervised samples is available. The LM input  $x$  is the text to be summarized and the target  $y$  is the given summary. We simulate the sequence-level preference by the classical Meteor score [90] and report the standard ROUGE scores [91], to avoid overfitting evaluation metrics as in the prompt task.
LM Training, Implementation, and Datasets. Since a supervised dataset  $\mathcal{D}$  is available in this task, the LM  $\pi_{\theta}$  can be trained by the weighted-MLE objective in Section 2.2. This objective could be more stable and computationally efficient than REINFORCE-style methods in tasks of long-sequence generation. Due to limited computing resources, unless explicitly mentioned, we use the standard T5-small model [89] for both the LM and reward function. The reward training is simply 1 epoch of training on randomly sampled  $10\%$  of the training set and is repeated every 0.5 epochs during the first 2 epochs of LM training. Reward function is learned with 3 sampled sequences and again the temperature  $\beta = 2$  in Eq. (4). Additional experiment details are in Appendix B.2. We test on the standard setting of two news summary datasets: CNN/DailyMail (CNN/DM) [92] and XSum [93].
Results. We compare four variants in our framework with the standard supervised fine-tuning and RL-based methods PPO and NLPO in RL4LMs [61] under the environmental reward Meteor — both
3There are small discrepancies between our reported RLPrompt results and the original paper's. We have confirmed our reproduced results both with RLPrompt's authors and with the recent TEMPERA paper [88].
Table 2: Results on text summarization. We bold the best result of each metric on each dataset. The results of Lead-3 on CNN/DM are from Ramamurthy et al. [61] and on XSum are from Lewis et al. [7]. Other baseline results are from our reruning RL4LMs' codebase [61] using T5-small. Number reporting formats follow Table 1.
CNN/DailyMail	XSum
ROUGE-1	ROUGE-2	ROUGE-L	ROUGE-1	ROUGE-2	ROUGE-L
Lead-3	40.10	17.50	36.30	16.30	1.60	11.95
Supervised	38.88 (0.02)	16.22 (0.05)	32.58 (0.04)	31.79 (0.02)	9.68 (0.01)	24.70 (0.03)
PPO	39.16 (0.51)	17.37 (0.33)	33.77 (0.37)	23.18 (0.31)	4.46 (0.19)	16.07 (0.32)
Supervised + PPO	39.17 (0.65)	17.29 (0.44)	33.76 (0.53)	28.24 (0.39)	7.68 (0.13)	20.02 (0.23)
NLPO	38.90 (0.35)	17.22 (0.35)	33.51 (0.42)	22.97 (0.23)	4.53 (0.13)	15.62 (0.35)
Supervised + NLPO	39.27 (0.60)	17.41 (0.36)	33.85 (0.42)	28.08 (0.16)	7.68 (0.20)	19.88 (0.16)
Ours (AVG)	40.94 (0.02)	18.78 (0.03)	38.17 (0.03)	33.62 (0.03)	11.17 (0.02)	26.33 (0.05)
Ours (SUM)	40.70 (0.06)	18.48 (0.05)	37.93 (0.08)	33.27 (0.09)	10.83 (0.07)	25.90 (0.06)
Ours (MIN)	40.78 (0.06)	18.67 (0.03)	38.01 (0.04)	33.57 (0.02)	11.14 (0.02)	26.30 (0.03)
Ours (MAX)	39.98 (0.08)	18.06 (0.03)	37.26 (0.06)	32.50 (0.14)	10.46 (0.12)	25.58 (0.12)
with and without task-specific supervised pre-training. For a fair comparison, the baseline results are from our rerunning RL4LMs' codebase with a T5-small model as our method. Table 2 shows the mean and standard deviation of ROUGE-1/2/L score across three random seeds.
On both datasets, our method shows favorable and stable performance against the classical and recent baselines. The better results of our method over supervised fine-tuning confirm the improvement of our reward-weighted MLE over the vanilla supervised loss, as discussed in Section 2.2. As in the prompt task, the gain of our method over RL-based baselines may indicate the benefit of utilizing our preference-grounded token-level guidance over learning under delayed feedback with KL penalty. In this task, using average as the aggregation function outperforms the classical summation. This confirms our idea in Section 2.1 on avoiding the interference of unequal sequence-lengths in training  $r_{\phi}$ . Using MIN is also suitable for this task, since it is not confounded by individual lengths and reflects overall text quality. Unlike the prompt task, using MAX is unsuitable, since good summaries can hardly be identified by a few keywords. Overall, these results show the importance of customizing the aggregation choice for the specific LM task, a key feature of our guidance-learning framework.
Further, to verify the performance of our method under a larger LM, we change the average variant of our method in Table 2 from T5-small to using T5-base LM. Fig. 2 (a) - (e) compares our method on CNN/DM against the baselines, with an additional metric BertScore [94]. The baseline results are directly cited from RL4LMs [61] and are the per-metric best across their three environmental rewards. Table 4 in Appendix A.1 shows the detailed numbers. It is clear that our method performs favorably against these strong baseline methods, especially in the ROUGE-L, BERTScore, Meteor, and ROUGE-2 metrics. To further clarify our method, we conducted a human study under the T5-base LM. The results are in Fig. 2 (f), with detailed setup and numerics in Table 5 of Appendix A.1. It is clear that this human evaluation on the summarization task supports the improvements in ROUGE, Meteor, and BertScore by our method. Further scaling-up of our method is left as future work.
This section discusses the following three research questions to better understand our framework.
(a): What will be the performance if we switch to using preference-based sequence-level guidance?
To further study the gain of grounding preference into token-level guidance, we change the preference-based token-level reward in our method to the corresponding sequence-level reward. Fig. 3 shows the results when applying this change to our best variants in the prompt and summarization tasks in Sections 4.1 and 4.2, including the T5-base LM in Section 4.2, in comparison to the best corresponding baselines. For summarization, we plot the average ROUGE scores, i.e., (ROUGE-1 + ROUGE-2 + ROUGE-L) / 3. Table 6 in Appendix A.1 shows each ROUGE metric with standard deviation.
We see that learning and using preference-based sequence-level guidance does not provide a significant advantage over those baselines that mostly directly work with the task-specific native sequence-level feedback — the results are even much worse than the baselines in some datasets. Besides, the results of our sequence-level variants are generally less stable. These echo the harm of the delayed-feedback issue discussed in Section 1. Overall, this set of comparisons confirms that the gain of our framework
<sup>4</sup>We carefully tuned the RL4LMs' baselines on several hyperparameters, which is detailed in Appendix B.2. <sup>5</sup>The "ROUGE-L" here refers to "Rouge-LSum" in RL4LMs and HuggingFace, as detailed in Appendix B.2.
(a) ROUGE-1
(b) ROUGE-2
(c) ROUGE-L
(d) Meteor
(e) BertScore
(f) Human
Figure 2: CNN/DM summarization of our method and baselines under T5-base LM. "Sup" denotes "Supervised". "Ref" denotes the ground-truth reference summary. Except for the human study in (f), baseline results are directly cited from RL4LMs [61] and are the per-metric best across their three environmental rewards.
Figure 3: Performance of our method using sequence-level and token-level preference-based guidance. "Best Baseline" refers to the best result in the baseline discrete-prompt methods for the prompt task, and the best result over all baseline methods for the summarization task. Error bars show one standard deviation.
mainly comes from our preference-grounding perspective, i.e., learning and using a preference-based token-level guidance, rather than simply learning and using "a preference-based guidance."
(b): How does our method perform if we remove the reward-function retraining scheme?
To study the effect of guidance re-estimation, we remove the reward-function retraining scheme from our best variants in the prompt and summarization tasks in Sections 4.1 and 4.2, including the T5-base LM in Section 4.2. Fig. 4 compares our methods with the best corresponding baselines. For the summarization task, we again plot the average ROUGE scores. Table 7 in Appendix A.1 shows each ROUGE metric with standard deviation. Appendix G discusses more on this re-estimation scheme.
Without guidance re-estimation, our method still performs competitively against the strong baselines, which corroborates the benefit of our preference-grounded guidance. Fig. 4 also verifies our intuition in Section 2 that the gain of this scheme depends on the zero-shot ability of the initial LMs. Specifically, in the prompt task where the initial LM has little zero-shot ability, reward-function retraining is helpful to both improve performance and reduce variance. In the summarization task where the initial LM does have some zero-shot ability (as shown in Ramamurthy et al. [61]), guidance re-estimation indeed helps results not as much, since the distribution-shift issue in Section 2 is less significant in this case. In this task, both our variants, with and without reward retraining, outperform the baselines.
(c): What if we learn the token-level guidance by a different number of text sequences?
To study how the number of sequences used to learn the reward function impacts our method's performance, we vary this number in the AVG variant in Tables 1 and 2. Fig. 5 shows the prompt results on SST-2 and summarization results on CNN/DM and XSum. For the latter, we again plot the average ROUGE scores. The scores of each ROUGE metric are in Tables 8 and 9 of Appendix A.1.
Recall that the best baseline result on SST-2 in Table 1 is 90.5, on CNN/DM and XSum in Table 2 is respectively 31.3 and 22.06. Thus, our method is generally robust to the number of sequences
Figure 4: Performance of our method with and without the reward-function retraining scheme. "Best Baseline" refers to the same as in Fig. 3. Error bars show one standard deviation.
Figure 5: Varying the number of sequences to learn the token-level guidance, showing mean over random seeds.
used to learn the guidance. Compared with the classical pairwise-preference learning (Section 2), our framework has the flexibility in using multiple sequences. As illustrated in Fig. 5, using three or more sequences to learn the reward function can be generally more beneficial than using only two.
Due to the page limit, we defer additional ablation study to Appendix A.2, where we (1) show that our framework is generally robust to the hyperparameter  $\beta$  in Eq. (4) and  $\alpha$  in Eq. (5); (2) further validate the harm of the delayed-feedback issue to the relevant LM-training methods on longer text-sequence generation; (3) show that the efficacy of our framework is not tied to the specific preference sources considered in this section.
To address the granularity mismatch between the sequence-level preference and the token-level LM training losses, in this paper, we develop an alternate-learning process, where we iterate between grounding sequence-level preference into token-level training guidance, and training the LM with the learned guidance. Our method performs competitively on two distinct representative LM tasks. Future work includes combining our preference-grounded guidance with RL-based LM training, and applying our method to human preference and/or other tasks such as (task-oriented) dialog systems.
S. Yang, S. Zhang, and M. Zhou acknowledge the support of NSF-IIS 2212418, NIH-R37 CA271186, the Texas Advanced Computing Center (TACC), and the NSF AI Institute for Foundations of Machine Learning (IFML). S. Yang acknowledges the support of the University Graduate Continuing Fellowship from UT Austin.
[1] Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in neural information processing systems, 13, 2000.
[2] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[3] Steven J Rennie, Etienne Marcheret, Youssef Mroueh, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7008-7024, 2017.
[4] Alec Radford and Ilya Sutskever. Improving Language Understanding by Generative Pre-Training. In *arxiv*, 2018.
[122] Changyeon Kim, Jongjin Park, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. Preference transformer: Modeling human preferences using transformers for RL. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=Peot1SFDX0.
[123] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[124] Shentao Yang, Zhendong Wang, Huangjie Zheng, Yihao Feng, and Mingyuan Zhou. A regularized implicit policy for offline reinforcement learning. arXiv preprint arXiv:2202.09673, 2022.
Appendix for “Preference-grounded Token-level Guidance for Language Model Fine-tuning”
1 Introduction 1
2 Main Method 2
2.1 Token-level Guidance Learning for Preference Grounding 3
2.2 LM Training with Preference-grounded Token-level Guidance 4
3 Related Work 5
4 Experiments 6
4.1 Input-agnostic Discrete-prompt Generation 6
4.2 Text Summarization 7
4.3 Ablation Study 8
5 Conclusion 10
A Additional Experimental Results 20
A.1 Tabular Results 20
A.2 Further Ablation Study 22
B Additional Experiment Details 24
B.1 Prompt Generation 24
B.2 Text Summarization 25
C A Naive Numeric Example for the Average Aggregation 26
D Details on the Prompt Generation Task 26
E More Related Work 27
F A Discussion on Applying RL Methods to LM Tasks 28
F.1 LM Generation as a Token-level MDP 28
F.2 Delayed Feedback in RL-based LM Training 29
F.3 Sparse Reward with KL Penalty 29
G Further Discussion on the Guidance Re-estimation Scheme 30
H Potential Negative Societal Impacts 30
I Limitations 31
J Computational Resources 31
A.1 Tabular Results
Table 3: Examples of the generated discrete input-agnostic text-prompt and their classification accuracy on the corresponding test set.
SST-2 Prompt	Accuracy	AG News Prompt	Accuracy
guys filmmaker filmmaker rated Grade	94.18	newsIntroduction Comments Tags Search	85.78
MovieMovieFilm rated Grade	94.18	newsTopic Blog Support Category	85.55
Rated CinemaScoreReporting Grade	94.01	newsRecentRecentPhotosIntroduction	84.53
employment theater rated Oscars Grade	93.96	news Recent Brief LatestExample	84.51
scene filmmaking rated comedian Grade	93.85	newsVirtualBlogBlogNet	84.33
Table 4: Detailed results on CNN/DM summarization under T5-base LM for Section 4.2. We bold the best result of each metric. Baseline results are directly cited from RL4LMs [61]. "Env. Reward" denotes the environmental reward in RL4LMs. The "ROUGE-L" here refers to "Rouge-LSum" in RL4LMs and in the Hugging Face interface, which is discussed in details in Appendix B.2. In Section 4.2, we plot the results of our method with the average aggregation, which is the best variant in Table 2. We report the mean (standard deviation) of our method over three random seeds.
Algorithm	Env. Reward	ROUGE-1	ROUGE-2	ROUGE-L	Meteor	BertScore
Lead-3	40.1	17.5	36.3	33.3	87.4
Supervised	41.1	17.7	34.3	30.9	87.6
PPO	Rouge-1	41.0	18.2	34.9	27.6	87.6
Rouge-Avg	39.6	17.6	33.8	27.0	87.4
Meteor	40.8	17.8	34.2	30.1	87.3
NLPO	Rouge-1	40.4	18.0	34.4	27.5	87.5
Rouge-Avg	40.4	17.7	34.4	27.4	87.4
Meteor	40.5	18.0	34.3	29.2	87.2
Supervised + PPO	Rouge-1	41.7	18.9	35.8	27.8	88.2
Rouge-Avg	42.5	19.4	36.3	29.6	88.2
Meteor	42.6	19.4	36.1	31.6	88.0
Supervised + NLPO	Rouge-1	42.1	19.3	36.1	28.7	88.2
Rouge-Avg	42.4	19.3	36.3	29.5	88.2
Meteor	42.9	19.4	36.1	31.9	88.0
Ours (AVG)	43.09 (0.06)	20.17 (0.04)	39.99 (0.07)	35.23 (0.06)	89.61 (0.12)
Ours (SUM)	42.86 (0.08)	19.92 (0.08)	39.76 (0.11)	34.74 (0.37)	89.24 (0.11)
Ours (MIN)	42.92 (0.14)	20.01 (0.02)	39.84 (0.08)	34.88 (0.13)	89.33 (0.07)
Ours (MAX)	42.38 (0.17)	19.49 (0.02)	39.34 (0.09)	34.13 (0.32)	89.09 (0.19)