@inproceedings{
nath2024learning,
title={Learning Goal-Conditioned Representations for Language Reward Models},
author={Vaskar Nath and Dylan Z Slack and Jeff Da and Yuntao Ma and Hugh Zhang and Spencer Whitehead and Sean M. Hendryx},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=Swh8LxuycA}
}
Techniques that learn improved representations via offline data or self-supervised objectives have shown impressive results in traditional reinforcement learning.
Nevertheless, it is unclear how improved representation learning can benefit reinforcement learning from human feedback on language models.
In this work, we propose training reward models (RMs) in a contrastive, $\textit{goal-conditioned}$ fashion by increasing the representation similarity of future states along sampled preferred trajectories and decreasing the similarity along randomly sampled dispreferred trajectories.
This objective significantly improves reward model performance by up to 0.09 AUROC across challenging benchmarks, such as MATH and GSM8k. These findings extend to general alignment as well -- on the Helpful-Harmless dataset, we observe 2.3\% increase in accuracy.
Beyond improving reward model performance, we show this way of training RM representations enables improved steerability because it allows us to evaluate the likelihood of an action achieving a particular goal-state (e.g. whether a solution is correct or helpful).
Leveraging this insight, we find that we can filter up to 55\% of generated tokens during majority voting by discarding trajectories likely to end up in an "incorrect" state, which leads to significant cost savings.
We additionally find that these representations can perform fine-grained control by conditioning on desired future goal-states.
For example, we show that steering a Llama 3 model towards helpful generations with our approach improves helpfulness by $9.6$\% over a supervised-fine-tuning trained baseline.
Similarly, steering the model towards complex generations improves complexity by $21.6$\% over the baseline.
Overall, we find that training RMs in this contrastive, goal-conditioned fashion significantly improves performance and enables model steerability.
Learning Goal-Conditioned Representations for Language Reward Models
Vaskar Nath Dylan Z Slack Jeff Da Yuntao Ma Hugh Zhang Spencer Whitehead Sean M. Hendryx