Mechanism

Three-level weight-space RLHF algorithm developed by Google DeepMind to achieve Pareto-optimal reward vs. KL trade-offs:

  1. Inner Loop: Run MM parallel RLHF runs with distinct KL/objective weights, using a moving EMA anchor θEMA←(1−μ)θEMA+μθ\theta_{\text{EMA}} \leftarrow (1-\mu)\theta_{\text{EMA}} + \mu \theta as the dynamic KL reference policy πref\pi_{\text{ref}}.
  2. Middle Loop: Merge the MM RL policies via SLERP into a single policy θSLERP\theta_{\text{SLERP}}.
  3. Outer Loop: Linearly interpolate (LITI) the merged RL policy back toward the initial SFT initialization θinit←(1−η)θinit+ηθSLERP\theta_{\text{init}} \leftarrow (1-\eta)\theta_{\text{init}} + \eta \theta_{\text{SLERP}} to start the next iteration.

References

  1. WARP: On the Benefits of Weight Averaged Rewarded Policies ↗

    Alexandre Ramé, Johan Ferret, Nino Vieillard, Robert Dadashi, Léonard Hussenot, Pierre-Louis Cedoz et al. · 2024

    All 10 authors

    Alexandre Ramé, Johan Ferret, Nino Vieillard, Robert Dadashi, Léonard Hussenot, Pierre-Louis Cedoz, Pier Giuseppe Sessa, Sertan Girgin, Arthur Douillard, Olivier Bachem