Replica Tricks and Bayesian Shrinkage for Value Estimation in A*PO

Abstract

APO is a two-stage reinforcement learning algorithm that is used to fine-tune large language models (LLM) on specific tasks. Stage 1 estimates the KL-regularized optimal value function V offline, and Stage~2 trains the policy online by regressing on optimal advantage targets. The log-mean-exp (LME) estimator of V* used in Stage 1 is downward-biased due to Jensen's inequality. This bias is significant, especially when the budget is small and the pass rate is low. This thesis develops and evaluates two families of alternative estimators. Firstly, inspired by the replica trick from statistical physics, we develop two types of replica-based estimators. We prove that the single-replica plug-in estimator does not reduce bias under a fixed budget, and develop the multi-n slope estimator, which exploits the structure of the replica trick more directly. Secondly, for the binary reward, we show that LME is equivalent to plugging the maximum-likelihood pass-rate estimate into a known closed form, which causes endpoint snapping when the pass rate is low. Inspired by conjugate priors, we develop two different Beta-smooth estimators to solve this problem. We evaluate these estimators in three experiments: (1) a scalar experiment where pass rate is fixed, (2) a contextual bandit experiment with 500 prompts of different pass rates, and (3) a full two-stage simulation of the APO pipeline. We discovered that replica-based estimators have a lower RMSE compared to LME when the pass rate is very low. However, they have a much larger advantage distortion that limits their practical value in Stage 2. The Beta-smooth estimators perform better compared to LME in all three experiments. In particular, the Jeffreys Beta-smooth estimator (p = (k+0.5)/(N+1)) achieves the best RMSE, reduces advantage distortion compared to LME, and carries over the improvement in Stage 1 into 25% lower parameter variance in the learned policy in Stage 2. We recommend it as a drop-in replacement for LME in APO's Stage 1.