Reproduce “Reward-Shifted Speculative Sampling Is An Efficient Test-Time Weak-to-Strong Aligner” on OpenResearch

About this paper

Bolian Li, Yanran Wu, Xinyu Luo, Ruqi Zhang

Researchers at Purdue University developed Reward-Shifted Speculative Sampling (SSS), an algorithm that integrates human preference alignment into speculative decoding for large language models, achieving superior alignment quality and significantly reducing inference cost. SSS provided speedups of up to 5.4x over Best-of-N baselines while obtaining higher "Gold Reward" scores than other test-time alignment methods and even some directly DPO/PPO-trained models.

Key insights

• Inference acceleration techniques like speculative sampling can be repurposed for preference alignment by carefully designing the sampling process to account for a shifted (aligned) draft model distribution.
• The challenge of distributional shift between an aligned draft model and an unaligned target model can be leveraged by modifying speculative sampling's acceptance and bonus token distributions to recover a desired aligned distribution.
• Achieving high-quality alignment at test-time with SSS, using a small aligned draft model, can be more computationally efficient and even surpass the alignment quality of explicitly RLHF-trained large target models.

Summary by alphaXiv.