Metropolis-Hastings improves policy mixing for large language models
Metropolis-Hastings Dominates Importance Resampling for Policy Composition
Machine Learning
Summary
Combining different ways for a language model to choose words can help balance various goals without retraining the model each time. The authors show that using a special method called Metropolis-Hastings to mix these choices produces results that are as close or closer to the ideal combination than a common method called importance resampling. They also study where this method works best and test it on small and large models.
What this means in practice
- •For ai product engineers: Create language model outputs balancing multiple objectives at runtime without costly retraining by using Metropolis-Hastings-based policy compositions.
- •For conversational ai developers: Improve generation quality when combining diverse dialogue policies by applying Metropolis-Hastings corrections during decoding.
Authors
Alexey Kurennoy, Ramil Yarullin, Fergal Reid
Abstract
Post-training a large language model (LLM) often requires exploring trade-offs between multiple rewards, but retraining for each trade-off is expensive. Decoding-time policy composition allows these trade-offs to be adjusted by combining reward-specific policies at inference time. This composition targets a weighted product of the policies' probabilities over complete responses, but standard implementations combine their next-token probabilities, generally introducing sampling bias. We analyze a known iterative correction based on independence Metropolis-Hastings (MH). Our main result shows that, for every rollout budget, MH produces an output distribution at least as close to the target as sampling-importance-resampling (SIR) with the same budget, as measured by every convex f-divergence. We also derive a lower bound on MH's improvement over the uncorrected decoder in a consensus objective measuring agreement with the supplied policies. We further characterize the correction's sampling error in two asymptotic regimes: when the reward-specific policies approach agreement, and when the log ratio between target and uncorrected-decoder probabilities fluctuates increasingly widely, as can happen for long responses. We complement our analysis with experiments in enumerable and LLM-scale settings.