Pith. sign in

REVIEW 3 major objections 4 minor 6 references

The learned GFlowNet partition function is a per-prompt accuracy meter: β log Z_φ estimates current accuracy, and scheduling training on prompts near 0.5 plus replaying mis-estimated samples improves sample efficiency.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-02 23:45 UTC pith:B54ARZQY

load-bearing objection Clever reuse of the GFlowNet partition function as an accuracy signal, with consistent empirical gains — but the replay buffer silently breaks the estimator's sampling assumption. the 3 major comments →

arxiv 2602.12642 v2 pith:B54ARZQY submitted 2026-02-13 cs.CL cs.AI

Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Dohyung Kim , Minbeom Kim , Jeonghye Kim , Sangmook Lee , Sojeong Rhee , Kyomin Jung This is my paper
classification cs.CL cs.AI
keywords partition functionGFlowNetsRLVRonline accuracy estimationadaptive prompt selectionprioritized replaysample efficiencyLLM reasoning
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper argues that a number already computed during GFlowNet-based post-training of language models—a style of distribution-matching RL where the model samples outputs in proportion to a reward-shaped target—contains a signal that previous work threw away: the model's current accuracy on that prompt. In the exact identity, the log partition function equals accuracy plus a KL term, and the paper's empirical evidence is that this term stays negligibly small, so β log Z_φ(x) acts as an almost-free accuracy estimate. The authors use that estimate to pick, at every step, the prompts whose predicted accuracy is closest to 0.5—the difficulty that gives the richest learning signal—and to replay correct answers where the estimate is most wrong. On math and code benchmarks, this scheduling reaches higher pass@1 and pass@k than uniform sampling, with no extra rollouts. A sympathetic reader would take the contribution to be a reframing: the normalizer is a reusable difficulty scheduler, not an overhead.

Core claim

On its own terms, the central claim is that under a Trajectory Balance loss that anchors the reference to the rollout policy, the learned partition function satisfies Proposition 4.1: p_old(x) = β log Z*(x) − β D_KL(π_old(·|x) ∥ π_θ(·|x)). Since the KL term stays below roughly 4 × 10^{-3} on average during training, the paper drops it and uses β log Z_φ(x), clipped to [0,1], as the online accuracy estimate. It then closes the loop: select the top-m prompts whose estimated accuracy is closest to 0.5, and replay correct outputs whose observed accuracy differs most from the estimate. This makes adaptive prompt selection and replay byproducts of the GFlowNet update, improving pass@1 and pass@k o

What carries the argument

The mechanism is the learned partition function Z_φ(x)—a small MLP on frozen prompt embeddings—trained jointly with the policy using the Trajectory Balance objective L_TB = [log( Z_φ(x)π_θ(y|x) / (π_old(y|x) exp(β^{-1} r(x,y))) )]^2. Proposition 4.1 derives the accuracy identity from this loss; the practical estimator p̂_old(x) = clip(β log Z_φ(x), 0, 1) turns the normalizer into a difficulty score. Adaptive selection greedily picks prompts closest to τ=0.5, and replay prioritizes correct samples by |N_correct/N − p̂_old(x)|, both reusing quantities already available from training.

Load-bearing premise

Everything rests on the KL divergence between the rollout policy and the updated policy being uniformly small for every prompt, not just on average; the paper shows only the mean curve and its Limitations section concedes this may fail in other domains.

What would settle it

Train with a substantially larger learning rate or fewer rollouts per prompt, and for every prompt record the true rollout accuracy alongside β log Z_φ(x). If the per-prompt gap is not uniformly small, or the Spearman correlation between estimate and truth falls well below the roughly 0.5 reported, the accuracy-meter claim is refuted.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Accuracy estimation becomes a byproduct of the existing GFlowNet update, so adaptive prompt selection no longer requires oversampling or per-question accuracy histories.
  • Focusing training on the 0.5-accuracy frontier reliably speeds early gains and lifts final accuracy, and the target τ can be tuned for different trade-offs.
  • Replaying high-miscalibration correct samples improves the calibration of Z_φ and downstream pass@1.
  • The gains are not bought at the cost of diversity: pass@k stays above both reward-maximizing and prior distribution-matching baselines.
  • The identity extends to arbitrary two-valued reward configurations, such as −1/+1 rewards, through a linear rescaling.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If β log Z_φ is a trustworthy difficulty score, the same number can drive inference-time compute allocation: sample more rollouts for low-scoring prompts and fewer for high-scoring ones, without training a separate reward or difficulty model.
  • A testable extension is to replace the fixed τ=0.5 with a scheduled or per-prompt target; the paper's own τ=0.9 results show fast early gains followed by a plateau, suggesting a schedule might combine both benefits.
  • The replay rule embodies a broader principle—when a learned normalizer disagrees with empirical frequency, revisiting those samples improves calibration—which could transfer to other generative models that train an energy or partition term.
  • The accuracy estimator could also act as a distribution-shift alarm during RL: prompts whose estimated accuracy jumps abruptly are exactly the prompts where the policy changed most and may warrant inspection.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper proposes PACED-RL, a GFlowNet-based post-training framework for LLM reasoning that reuses the learned partition function Z_phi(x) as an online per-prompt accuracy estimate. A theoretical result (Prop. 4.1) relates beta log Z*(x) to the accuracy p_old(x) under the rollout policy, minus a KL term that the authors argue is small and can be discarded. This estimate is then used for adaptive prompt selection toward prompts with accuracy near 0.5, and for an error-prioritized replay buffer. Experiments on math and code benchmarks with 1.5B and 7B models report consistent gains over GRPO, FlowRL, and adaptive-sampling baselines, plus an ablation showing the replay component helps. The central claim is that the partition function is not merely a normalizer but a reusable difficulty signal that improves sample efficiency at negligible extra cost.

Significance. If the claims hold, the paper makes a useful conceptual and practical contribution: it turns an existing training variable Z_phi into a lightweight difficulty estimator, avoids extra rollouts for prompt selection, and provides a principled explanation for the gains. The empirical evaluation is broad (two domains, three models, multiple baselines), and the paper includes ablations, correlation analyses, and a wall-clock overhead table. The theoretical derivation is transparent and the proposed estimator is simple to implement. However, the central derivation contains a sign error, and the replay-buffer mechanism breaks the sampling assumption of the proposition without correction. These issues are fixable but currently undermine the rigor of the advertised theoretical link and the reliability of the accuracy estimates driving both scheduler components.

major comments (3)
  1. [Appendix A.1, Eq. (15)-(16); Prop. 4.1] The derivation has a sign error. From Eq. (14), beta log Z* = E[r] + beta E_{pi_old}[log pi_old - log pi_theta]. Since E_{pi_old}[log pi_old - log pi_theta] = -D_KL(pi_old || pi_theta), Eq. (15) gives E[r] = beta log Z* + beta D_KL(pi_old || pi_theta), not beta log Z* - beta D_KL. Thus Prop. 4.1 should read p_old(x) = beta log Z*(x) + beta D_KL(pi_old || pi_theta). The sign error propagates to Eq. (8) and to the generalized Eq. (22). The practical claim that the KL term is small may still hold, but the stated theoretical relationship is incorrect and should be corrected.
  2. [Sec. 4.4, Algorithm 1 lines 8-13; Prop. 4.1] Prop. 4.1 is derived under y ~ pi_old, but Algorithm 1 updates Z_phi on a mixture of on-policy rollouts and a replay buffer B containing only r=1 trajectories from older policies. The minimizer of Eq. (6) under the mixture mu = (1-alpha) pi_old + alpha nu solves log Z*_mix = E_mu[beta^{-1} r + log pi_old - log pi_theta]. Since nu is supported on correct outputs, E_nu[r]=1, so beta log Z_phi is systematically inflated relative to p_old, with an additional nu-dependent KL term. No importance correction or bias bound is provided. This biases both the prompt selection in line 5 and the replay priority in Eq. (9), which assume hat p_old ~ p_old. Please either remove replay samples from the Z_phi update, add a correction, or provide an explicit analysis of the induced bias and an ablation demonstrating its effect.
  3. [Sec. 4.2, Fig. 1, App. C] The estimator Eq. (8) discards a nonnegative per-prompt term, but the evidence that it is small is mean-only: Fig. 1 shows the mean of beta D_KL, not its per-prompt distribution, and App. C reports aggregate correlations. A mean below 4e-3 does not bound per-prompt values, and the scheduler selects prompts by ranking beta log Z_phi, so a small number of badly estimated prompts can change the selected batch. Given the sign error in Prop. 4.1, the direction of the approximation error is also misstated. Please report per-prompt KL statistics (e.g., quantiles over the training set) or otherwise justify that the approximation is uniformly small on the distributions used.
minor comments (4)
  1. [Sec. 5.3, Eq. (10)] Equation (10) writes log hat Z(x_i) = (1/N) sum (R(x_i,y_ij) - log pi_theta(...)). Given the earlier definition R(x,y) = pi_ref(y|x) exp(beta^{-1} r), this should be log R(x_i,y_ij), not R(...). Please fix the typo, as it affects reproducibility of the VarGrad baseline.
  2. [Algorithm 1, line 12; Eq. (9)] The replay priority is defined in Eq. (9) using the observed accuracy N_correct/N, but Algorithm 1 line 12 writes |p_old(x) - hat p_old(x)|. Clarify that p_old(x) here is the empirical rollout accuracy, not the unknown true accuracy.
  3. [Sec. 4.1, Eq. (6)-(7)] When pi_ref is replaced by pi_old, the optimal partition function depends on pi_old: Z*(x) in Eq. (7) is different from Z(x) in Eq. (3). The notation should make this explicit (e.g., Z_{pi_old}(x)) to avoid confusion.
  4. [Sec. 5, experimental setup] No random seeds or error bars are reported for the main tables or ablation. Given the known variance of RLVR training, please report results over at least 2-3 seeds, or state why this is infeasible.

Circularity Check

0 steps flagged

No significant circularity: the partition-function/accuracy identity is a genuine optimality derivation, and the replay-buffer mismatch is a correctness concern rather than a circular step.

full rationale

The paper's central derivation (Prop. 4.1, Appendix A.1) starts from the expected Trajectory Balance loss Eq. 11, differentiates with respect to log Z_phi, and obtains Eq. 14, which rearranges to p_old(x) = beta log Z*(x) - beta KL(pi_old || pi_theta). This is a legitimate first-order optimality identity, not a circular definition: Z* is defined as the normalizer of a reward-weighted policy distribution, while p_old is the expected reward under pi_old, and the proof shows the link rather than assuming it. The practical estimator Eq. 8 drops the KL term under an explicit, empirically monitored approximation (Fig. 1, App. C), and the Limitations section acknowledges this assumption may fail in other domains. The replay-buffer design does violate the y ~ pi_old sampling assumption of Prop. 4.1, since Algorithm 1 trains Z_phi on on-policy rollouts plus r=1 replay samples from older policies; this can bias beta log Z_phi, but it is an unexamined correctness/robustness gap, not a reduction of the prediction to its own inputs. Self-citations (Lee et al. 2025, Kim et al. 2024, Bartoldson et al. 2025) are used for standard implementation details and replay inspiration, not for a load-bearing uniqueness or equivalence claim. The main empirical gains are evaluated on held-out benchmarks, so the central claim has independent content. Overall, the derivation is self-contained and not circular, though the off-policy replay mismatch and the KL-drop approximation warrant scrutiny.

Axiom & Free-Parameter Ledger

4 free parameters · 5 axioms · 0 invented entities

The method rests on two unproven assumptions not fully tested: per-prompt smallness of the KL term, and off-policy validity of the Z-accuracy identity. The free parameters τ, β, buffer sizes, and N are set by prior work or simple grid, not derived. No new physical or mathematical entities are introduced.

free parameters (4)
  • τ (target accuracy for prompt selection) = 0.5
    Fixed to 0.5 for main experiments; ablation (Fig. 4) shows 0.5 is optimal. The scheduler's premise that prompts at accuracy 0.5 are most informative is inherited from prior work.
  • β (KL regularization coefficient) = 0.05
    Fixed in the loss and in the estimator p_old ≈ β log Z_φ (Eq. 8); the scaling of the accuracy estimate depends on it.
  • Replay buffer sizes B_max and B_add = 128, 64
    Chosen for all runs; affect the off-policyness of Z_φ training and thus the validity of the accuracy interpretation.
  • N rollouts per question = 8
    Affects the observed accuracy N_correct/N in the replay priority (Eq. 9) and the noise of the accuracy estimate.
axioms (5)
  • domain assumption The KL divergence term β D_KL(π_old ∥ π_θ) is uniformly small and can be neglected
    Sec. 4.2; Fig. 1 shows only the mean over prompts; Limitations E acknowledges this may fail on other tasks.
  • domain assumption Z_φ(x) closely tracks the loss minimizer Z*(x) at each training step
    Needed for Eq. 8 to hold in practice; the paper provides no convergence analysis or tracking bound.
  • domain assumption The TB loss minimizer identity holds with off-policy replay data
    Prop. 4.1 is derived for y ~ π_old; replay buffer contains older-policy samples, which changes the minimizer of the expected squared loss and may bias the accuracy estimate.
  • standard math Standard TB/GFlowNet theory (minimizing TB loss matches target distribution)
    Eq. 5 cites Zhu et al. (2025a); accepted background.
  • domain assumption Prompts with accuracy nearest 0.5 yield the most sample-efficient learning signal
    Sec. 4.3; inherits from Foster et al. (2025), Bae et al. (2025).

pith-pipeline@v1.3.0-alltime-deepseek · 16655 in / 19475 out tokens · 158926 ms · 2026-08-02T23:45:23.032342+00:00 · methodology

0 comments
read the original abstract

Reward-maximizing RL methods have shown to be capable of enhancing the reasoning performance of LLMs, but often lead to reduced generation diversity. Recent works address this issue by adopting GFlowNets, training LLMs to match a target distribution while jointly learning its partition function. In contrast to prior works that treat this partition function solely as a normalizer, we reinterpret it as a per-prompt expected-reward (i.e., online accuracy) signal, leveraging this unused information to improve sample efficiency. Specifically, we first establish a theoretical relationship between the partition function and per-prompt accuracy estimates. Building on this key insight, we propose Partition Function-Guided RL (PACED-RL), a post-training framework that leverages accuracy estimates to prioritize informative question prompts during training, and further improves sample efficiency through an accuracy estimate error-prioritized replay. Crucially, both components reuse information already produced during GFlowNet training, effectively amortizing the compute overhead into the existing optimization process. Extensive experiments across diverse benchmarks demonstrate strong performance improvements over GRPO and prior GFlowNet approaches, highlighting PACED-RL as a promising direction for a more sample efficient distribution-matching training for LLMs.

Figures

Figures reproduced from arXiv: 2602.12642 by Dohyung Kim, Jeonghye Kim, Kyomin Jung, Minbeom Kim, Sangmook Lee, Sojeong Rhee.

Figure 1
Figure 1. Figure 1: Training dynamics of mean of estimated values of β DKL πold(·| x) [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Pass@1(%) dynamics of DeepSeek-R1-Distill-Qwen￾1.5B plotted with respect to wall-clock time. Pass@1 values ob￾tained via averaging over 8 independent attempts. 5.2. Main Results Code Generation. As shown in [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Pass@k(%) curves of Qwen2.5-Math-7B trained on GRPO, FlowRL, and PACED-RL and evaluated on AIME 24 and AIME24 for k ∈ {2, 4, 8, 16, 32, 64, 128, 256}. Across all evalu￾ated values of k, PACED-RL achieves the highest performance. out training compared to all baselines, illustrating the bene￾fits of improved sample efficiency via leveraging accuracy estimates from the partition function. Notably, PACED￾RL ac… view at source ↗
Figure 4
Figure 4. Figure 4: Train and test accuracies of Qwen2.5-Math-1.5B trained on the DeepScaleR dataset and evaluated on the MinervaMath benchmark, for target accuracy values τ ∈ {0.1, 0.3, 0.5, 0.7, 0.9}. 5.3. Analysis Controlling Train Accuracy with the Partition Function [PITH_FULL_IMAGE:figures/full_fig_p006_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Progression of the average difficulty of question prompts sampled in each training batch under the Qwen2.5-Math-7B model. Progression of Question Difficulty During Training [PITH_FULL_IMAGE:figures/full_fig_p007_5.png] view at source ↗
Figure 7
Figure 7. Figure 7: Spearman’s correlation coefficient measured on the MATH training data and the DeepScaleR training data, using the Qwen2.5- Math-1.5B model [PITH_FULL_IMAGE:figures/full_fig_p014_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Pearson’s correlation coefficient measured on the MATH training data and the DeepScaleR training data, using the Qwen2.5- Math-1.5B model.. C. Supplementary Results C.1. Correlation of Accuracy Estimates to Empirically Observed Accuracies [PITH_FULL_IMAGE:figures/full_fig_p014_8.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

6 extracted references · 3 linked inside Pith

  1. [3]

    Kim, M., Yun, T., Bengio, E., Zhang, D., Bengio, Y ., Ahn, S., and Park, J

    URL https://openreview.net/forum? id=chfJJYC3iL. Kim, M., Yun, T., Bengio, E., Zhang, D., Bengio, Y ., Ahn, S., and Park, J. Local search GFlownets. InThe Twelfth International Conference on Learning Representations,

  2. [5]

    Qu, Y ., Wang, Q., Mao, Y ., Hu, V

    URL https://openreview.net/forum? id=Feiz5HtCD0. Qu, Y ., Wang, Q., Mao, Y ., Hu, V . T., Ommer, B., and Ji, X. Can prompt difficulty be online predicted for acceler- ating rl finetuning of reasoning models?arXiv preprint arXiv:2507.04632, 2025. Schaul, T., Quan, J., Antonoglou, I., and Silver, D. Priori- tized experience replay.arXiv preprint arXiv:1511....

    arXiv 2025

  3. [6]

    For FlowRL (Zhu et al., 2025a) and MoPPS (Qu et al., 2025), we adopt the official implementations and default hyperparameters provided by the authors in their implementations

    library. For FlowRL (Zhu et al., 2025a) and MoPPS (Qu et al., 2025), we adopt the official implementations and default hyperparameters provided by the authors in their implementations. For LILO (Foster et al., 2025), we implement the over-sampling and rejection-sampling procedure that selects the top-m question prompts whose accuracies are closest to 0.5....

    2025

  4. [2021]

    Hessel, M., Modayil, J., Van Hasselt, H., Schaul, T., Os- trovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D

    URL https://openreview.net/forum? id=7Bywt2mQsCe. Hessel, M., Modayil, J., Van Hasselt, H., Schaul, T., Os- trovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. InProceedings of the AAAI conference on artificial intelligence, volume 32, 2018. Horgan, D., Quan, J., Budden, D...

    Pith/arXiv arXiv 2018

  5. [2024]

    URL https://openreview.net/forum? id=6cFcw1Rxww. Kwon, O. J., Matsunaga, D., and Kim, K.-E. Gdpo: Learn- ing to directly align language models with diversity using gflownets. InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 17120–17139, 2024. Lee, S., Kim, M., Cherif, L., Dobre, D., Lee, J., Hwang, S. J., Kawag...

    Pith/arXiv arXiv 2024

  6. [2025]

    Bengio, Y ., Lahlou, S., Deleu, T., Hu, E

    URL https://openreview.net/forum? id=VwPt1WDQNB. Bengio, Y ., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E. Gflownet foundations.Journal of Machine Learning Research, 24(210):1–55, 2023. Chen, M. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. Chen, Z., Qin, X., Wu, Y ., Ling, Y ., Ye, Q., Zhao, W. X...

    Pith/arXiv arXiv 2023