POP: Prior-Fitted First-Order Optimization Policies

Jan Kobiolka; Christian Frey; Gresa Shala; Arlind Kadra; Erind Bedalli; Josif Grabocka

cs.LGFeb 17, 2026

POP: Prior-Fitted First-Order Optimization Policies

Authors: Jan Kobiolka, Christian Frey, Gresa Shala, Arlind Kadra, Erind Bedalli, Josif Grabocka

Organizations: Department of Computer Science University of Technology Nuremberg Germany · Department of Computer Science Albert Ludwig University of Freiburg Germany · Faculty of Natural Sciences University of Elbasan Albania

Abstract

Gradient-based optimizers are highly sensitive to design choices in their adaptive learning rate mechanisms. To address this limitation, we introduce POP, a meta-learned Reinforcement Learning (RL) policy that predicts adaptive learning rates for gradient descent, conditioned on the contextual information provided in the optimization trajectory. Our method introduces a novel RL reward formulation, a new function-scaling strategy for in-distribution generalization, and a novel prior that is used to sample millions of synthetic optimization problems. We evaluate POP on an established benchmark including 43 optimization functions of various complexity, where it significantly outperforms gradient-based methods. Our evaluation demonstrates strong generalization capabilities without task-specific tuning.

Figures & tables

Figure 1: POP (blue) adapts its learning rate based on the optimization landscape, enabling rapid convergence, escape from local minima, and improved global optimization compared to Adam (red). Yellow diamond represents the global minima, the white cross represents the start state, and the square represents the end state.

Dim.	POP	Random	Adam	GD	L-BFGS	Lion	Sophia-G
16D	0.2911	0.1177	0.2887	0.2510	0.1728	0.2409	0.2579
32D	0.4966	0.1669	0.5029	0.4003	0.4171	0.4245	0.4820

Table 1: Mean performance across different dimensionalities.

Figure 6: The average normalized regret for all methods on VLSE. Solid lines represent the mean value.

Panel of Figure 7 — Figure 7: Performance of POP compared across different model categories and representative models.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol / Quantity	Distribution / Value
$D$	$2$
Domain	$[-50,50]^{D}$
Input shift	$\mathcal{U}(-40,40)^{D}$
Vertical offset $b$	$0$
$M$	$1000$
Quadratic mixing weight $\alpha_{3}$	$\mathcal{U}(0.05,0.4)$ for bowl-dominated instances

Appendix

Table 3: Main hyperparameter values of the proposed prior.

Figure 8: Exemplary functions sampled from the prior distribution (cf. Section 2.4 )

Shared backbone (Transformer)
Component	Value
Embedding dimension	64
# Transformer blocks	8
# Attention heads	8 ( $d_{\text{head}}=8$ )
Dropout	0.05
Pooling	Last-token pooling

Appendix

Table 4: Key architecture hyperparameters for the shared Transformer backbone and actor/critic heads.

Component	Value
Actor lr	$5\times 10^{-5}$
Critic lr	$5\times 10^{-5}$
Optimizer	AdamWScheduleFree [ 8 ]
Actor weight decay	$1\times 10^{-4}$
Critic weight decay	$1\times 10^{-4}$
Actor warmup	300

Appendix

Table 5: PPO training hyperparameters and optimizer settings.

Panel of Figure 9 — Figure 9: Learning rate sweep on the in-distribution validation set.

Panel of Figure 10 — Figure 10: Hyperparameter sweep for the evolutionary methods on the in-distribution validation set.

Figure 11: Trajectory of POP’s on in-distribution functions. The algorithm progresses from the initial state (white cross) terminating at the final state (square).

Figure 12: Trajectory of POP’s on in-distribution functions. The algorithm progresses from the initial state (white cross) terminating at the final state (square).

Figure 13: In-distribution test set performance vs. all baselines. Mean normalized improvement over steps; shading indicates 95% CIs. Dashed line marks the context/optimization boundary.

Figure 14: In-distribution test set performance vs. all baselines. Mean normalized improvement over steps; shading indicates 95% CIs. Dashed line marks the context/optimization boundary.

Dim.	POP	Random	Adam	GD	L-BFGS	Lion	Sophia-G	TPE	GP	HEBO	CMA-ES	DE	GA	VELO
16D	0.2911	0.1177	0.2887	0.2510	0.1728	0.2409	0.2579	0.2107	0.2692	0.2702	0.1153	0.1421	0.2145	0.3100
32D	0.4966	0.1669	0.5029	0.4003	0.4171	0.4245	0.4820	0.3176	0.4807	0.4368	0.0572	0.1953	0.3691	0.5225

Appendix

Table 6: Mean performance across different dimensionalities.

Figure 15: The average normalized regret for all methods on VLSE. Solid lines represent the mean value.

Panel of Figure 16 — Figure 16: Performance of POP compared to individual competitors from different model categories on diverse function classes

Panel of Figure 17 — Figure 17: Performance of POP compared across different model categories and representative models.

Learning to Optimize

Learning to Optimize uses data to adapt algorithms, solver decisions, and solution maps for faster or better optimization.

See all 138 papers

Beyond Importance Sampling: Rejection-Gated Policy Optimization
Apr 16, 2026Ziwu Sun, Zhen Gao, Jiyong Zhang +1Importance Sampling Policy Optimization
Group Adaptive Clipping Policy Optimization
Aug 31, 2026Sheng Jia, Xiao Wang, Shiva Prasad Kasiviswanathan +1Group Relative Policy Optimization Policy Optimization
University of Toronto · Amazon
OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation
Jun 4, 2026Paavo Parmas, Yongmin Kim, Kohsei Matsutani +5Stochastic Optimization Risk-Sensitive RL
The University of Tokyo

Apr 16, 2026J/K move · Enter open · S save

Beyond Importance Sampling: Rejection-Gated Policy Optimization

Ziwu Sun, Zhen Gao, Jiyong Zhang, Jiaheng Li

Importance Sampling Policy Optimization Policy Gradient Methods RL Fine-Tuning Policy Learning Importance Sampling for RL

Open paper AskPDF arXiv

We propose a new perspective on policy optimization: rather than reweighting all samples by their importance ratios, an optimizer should select which samples are trustworthy enough to drive a policy update. Building on this view, we introduce Rejection-Gated Policy Optimization (RGPO), which replaces the importance sampling ratio r_theta = pi_theta / pi_old with a smooth, differentiable acceptance gate alpha_theta(s, a) = g(r_theta(s, a)) in the range [0, 1]. Unlike prior work that applies rejection sampling as a data-level heuristic before training, RGPO elevates rejection to an optimization principle: the gate participates directly in gradient computation and is implicitly updated alongside the policy. RGPO provides a unified framework: the policy gradients of TRPO, PPO, and REINFORCE all correspond to specific choices of the effective gradient weight w(r) = g'(r) * r. We prove that RGPO guarantees finite, bounded gradient variance even when importance sampling ratios are heavy-tailed (where IS variance diverges). We further show that RGPO incurs only a bounded, controllable bias and provides an approximate monotonic policy improvement guarantee analogous to TRPO. RGPO matches PPO in computational cost, requires no second-order optimization, and extends naturally to RLHF-style preference alignment. In online preference fine-tuning of Qwen2.5-1.5B-Instruct on Anthropic HH-RLHF (n = 3 seeds), RGPO uses a dual-ratio gate that anchors learning to both the previous policy and the reference model, achieving a Pareto-dominant outcome: the highest reward among online RL methods (+14.8% vs. PPO-RLHF) and the lowest KL divergence to the reference model (-16.0% vs. PPO-RLHF, -53.1% vs. GRPO).

Abstract

Learning to Optimize

Explore similar work

Beyond Importance Sampling: Rejection-Gated Policy Optimization

Group Adaptive Clipping Policy Optimization

OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation