Reinforcement Learning for Hierarchical Reasoning Rewards: Minimax-Optimal Rates with Transformers

2026-10-06

AI Summary

This study addresses the lack of theoretical foundations for reinforcement learning (RL) post-training by providing the first quantification of the theoretical advantages of online exploration. Methodologically, it constructs a hierarchical neural reward model and designs a KL-regularized Transformer actor-critic algorithm. Theoretically, we prove that online exploration progressively concentrates on high-reward regions, achieving minimax-optimal convergence rates under given query budgets and regularization strengths. In contrast, offline sampling yields only logarithmic decay. By establishing these theoretical guarantees, this work formally demonstrates the superiority of online RL over offline approaches in the context of post-training, thereby bridging a critical gap between empirical practice and theoretical understanding in aligning large language models through reinforcement learning.

Abstract

Reinforcement learning (RL) has become a standard tool for post-training language models on reasoning tasks, where the policy is updated by reward feedback while exploring the space of responses. Despite its empirical success, theoretical understanding of RL post-training remains limited, in particular of why on-policy exploration combined with a neural reward model is effective. In this paper, we address this question by modeling the reward as a hierarchical function on the response space: the reward consists of infinitely many local components, each of which becomes relevant only after the preceding ones have been resolved. We show that a natural Transformer-based actor--critic algorithm, which alternates between sampling from the current KL-regularized policy, fitting a Transformer critic to the observed rewards, and updating the policy, achieves the minimax optimal rates in the query budget and in the regularization strength up to logarithmic factors, and is minimax optimal for a fixed number of prompts. In contrast, we prove that sampling from the fixed reference distribution, as in offline reward modeling, can limit regret decay to a logarithmic rate. These results show that on-policy exploration progressively zooms in on the region where the reward is concentrated, and quantify its benefit for RL post-training.

Problem

• Reinforcement Learning
• Post-training
• Hierarchical Reasoning Rewards
• On-policy Exploration
• Language Models

Innovation

• Hierarchical Reasoning Rewards
• Reinforcement Learning Post-training
• Transformer Actor-Critic
• Minimax-Optimal Rates
• On-policy Exploration