PoEM is a framework that predicts the output of (RL) post-training on a new reward function without actually running RL. It takes a set of models already trained on other rewards and a new reward function, estimates composition weights using a small calibration set, and combines the log-policies linearly at inference time. The method is motivated by the observation that log-policies from RL training often span a low-rank subspace across rewards, even when rewards are not linearly connected.
Experiments on text and image modalities show that PoEM recovers most of the reward gain of direct RL training on composite rewards, with an error comparable to the difference between two RL runs. On a basis of ten reward models adapted with PPO, PoEM’s generations are closer to the directly trained RL policy than the leading single expert on nine of ten held-out rewards. On image generation with Stable Diffusion v1.4, it approximates a held-out RL adapter from the remaining basis.
This work is interesting because RL post-training is computationally expensive and must be redone whenever the reward changes. PoEM offers a way to quickly preview or combine reward effects without retraining, potentially saving significant compute and enabling rapid iteration on reward functions.