WARP: On the Benefits of Weight Averaged Rewarded Policies
Reinforcement learning from human feedback (RLHF) aligns large language
models (LLMs) by encouraging their generations to have high rewards, using a
reward model trained on human preferences. To prevent the forgetting of
pre-trained knowledge, RLHF usually incorporates a KL regularization; this
forces the policy to remain close to its supervised fine-tuned initialization,
though it hinders the reward optimization. To tackle the trade-off between KL
and reward, in this paper we introduce a novel alignment strategy named Weight
Averaged Rewarded Policies (WARP). WARP merges policies in the weight space at
three distinct stages. First, it uses the exponential moving average of the
policy as a dynamic anchor in the KL regularization. Second, it applies
spherical interpolation to merge independently fine-tuned policies into a new
enhanced one. Third, it linearly interpolates between this merged model and the
initialization, to recover features from pre-training. This procedure is then
applied iteratively, with each iteration's final model used as an advanced
initialization for the next, progressively refining the KL-reward Pareto front,
achieving superior rewards at fixed KL. Experiments with GEMMA policies
validate that WARP improves their quality and alignment, outperforming other
open-source LLMs.
667a598613c37a0fe4a8b06e
Alexandre Ramé
admin_assigned
2024-06-25T08:17:55.519Z
63c94ede00104ea998de19a6
/avatars/273959d87f0c67747588cf0700d64039.svg
Alexandre Rame
user
alexrame
667a598613c37a0fe4a8b06f
Johan Ferret
admin_assigned
2024-06-25T08:18:03.169Z
65afb7dbdd6bdfd73cd8e609
/avatars/b21069bc2d7ee4cc1508008e3c8ade64.svg
Johan Ferret
user
ferretj
667a598613c37a0fe4a8b070
Nino Vieillard
667a598613c37a0fe4a8b071
Robert Dadashi
admin_assigned
2024-06-25T08:18:17.076Z
65fc7c0da57b88a765aea493
/avatars/440f88769da1113d6158fe7e0514ead3.svg
Robert Dadashi
user
ddsh
667a598613c37a0fe4a8b072
Léonard Hussenot
admin_assigned
2024-06-25T08:18:24.900Z
667a7b2e1764931987803b46
/avatars/da2ee162f33c681e3d1c4d1e6d44dcbd.svg
Léonard Hussenot
user
leonardhussenot
667a598613c37a0fe4a8b073
Pierre-Louis Cedoz
admin_assigned
2024-06-25T08:18:30.765Z
64dbdd39e7bc8544f9ac4c1a
/avatars/7aface0db75500a2b0ede75179be35ce.svg
Pierre-Louis Cedoz
user
plcedoz
667a598613c37a0fe4a8b074
Pier Giuseppe Sessa
admin_assigned
2024-06-25T08:18:37.543Z
667a7876becec8fc5112cf1f
/avatars/6874f0e5a7edca4145c57d7908f89840.svg
Pier Giuseppe Sessa
user
piergs
667a598613c37a0fe4a8b075
Sertan Girgin
667a598613c37a0fe4a8b076
Arthur Douillard
claimed_verified
2024-06-25T07:32:35.814Z
622792366303bf1dc304f49f
/avatars/975c1cc3eb2f97cf8e848162056d5bea.svg
Arthur Douillard
user
ArthurDouillard
667a598613c37a0fe4a8b077
Olivier Bachem
WARP: On the Benefits of Weight Averaged Rewarded Policies
Reinforcement learning from human feedback (RLHF) aligns large language
models (LLMs) by encouraging their generations to have high rewards, using a
reward model trained on human preferences. To prevent the forgetting of
pre-trained knowledge, RLHF usually incorporates a KL regularization; this
forces the policy to remain close to its supervised fine-tuned initialization,
though it hinders the reward optimization. To tackle the trade-off between KL
and reward, in this paper we introduce a novel alignment strategy named Weight
Averaged Rewarded Policies (WARP). WARP merges policies in the weight space at
three distinct stages. First, it uses the exponential moving average of the
policy as a dynamic anchor in the KL regularization. Second, it applies
spherical interpolation to merge independently fine-tuned policies into a new
enhanced one. Third, it linearly interpolates between this merged model and the
initialization, to recover features from pre-training. This procedure is then
applied iteratively, with each iteration's final model used as an advanced
initialization for the next, progressively refining the KL-reward Pareto front,
achieving superior rewards at fixed KL. Experiments with GEMMA policies
validate that WARP improves their quality and alignment, outperforming other
open-source LLMs.
A new alignment strategy named Weight Averaged Rewarded Policies (WARP) enhances reinforcement learning from human feedback by merging policies at multiple stages to balance KL regularization and reward optimization, improving large language model quality and alignment.
reinforcement learning from human feedback large language models reward model KL regularization exponential moving average spherical interpolation linear interpolation KL-reward Pareto front LLMs GEMMA policies