PIPA: Preference Alignment as Prior-Informed Statistical Estimation

We are grateful to the reviewer for the detailed summary of the PIPA-M and PIPA-N derivations, including their extension to the step-level setting, as well as for recognizing our comprehensive experiments on math tasks.

1. Connectio between PIPA and DPO

1.1 $\beta$

The reviewer mentions that our current PIPA-N only reduces to DPO when $\beta=1$. We emphasize that DPO with $\beta \neq 1$ directly corresponds to PIPA-N with a power prior [1] in Bayesian analysis, meaning: $p(A|B) \propto p(B|A)^{\beta} p(A).$ Specifically, introducing $\beta$ modifies the original DPO loss as follows:

We appreciate the reviewer for acknowledging our novel well-supported framework and comprehensive experiments in math. We will improve writing to make presentation more clear.

1. Why learnable $p^\theta(c=1|x)$ is better?

Unlike DPO (which uses a simple 0.5 prior) and KTO (which uses a complex KL-based prior), our approach makes this term learnable—a more elegant solution. As explained in our response to reviewer 4Kmz (section 1.3), this term can be interpreted as problem difficulty and is better learned during training rather than using a predetermined fixed value.

2. Meaning of $y_{<t}$

We agree that conditional independence only works for math-like tasks with clear good/bad examples. However, for scenarios involving good/better comparisons you mentioned, a DPO-style paired loss function is necessary rather than the current KTO-style unpaired approach presented in our work, and it also means there is no step-level signals due to the uncertainty. To get the new algorithm in the new setting, we can follow Theorem A.1's derivation and design a paired PIPA variant (not explored in our current version which focuses on KTO-style unpaired PIPA). Starting from equation (9), we can expand terms in $h_\theta$ autoregressively while directly learning $p^{\theta}(c=1|x)$. This deriviation does not need token-level labels, avoiding the concerns you mentioned.

3. Writing improvements and clarifications

Thank the reviewer for the careful reading and suggestions to improve our work. We will

• Explain $g_0^\theta$ more clearly in Thm 2.1.
• Clarify the setting in Table 2.
• Change Parameter to parameterize in Line 152.
• Make notation with $\theta$ more clear.
• Reorganize the sections and move some experiments to appendix to make presentation more clear.

Dataset used in Figure 3: The likelihood is computed on Alphamath during training.

4. Experiments

4.1 Source of advantages

Our approach gains benefits from both formula and token-level rewards—these components work effectively and neither functions effectively without the other. Importantly, integrating token-level rewards into DPO/KTO frameworks is non-trivial and it’s not straightforward to design an algorithm that only employs token level rewards without appropriate changes to the formula.

4.2 Baselines

We don't include comparisons with RTO and OREO because they require an additional learnable model, unlike our approach which only needs an extra prediction head. This makes our method more directly comparable to the original DPO and KTO implementations with similar computational costs. We add the comparison with TDPO below using Deepseek model, and PIPA is better. More comparison will be included later due to the limited response period.

5. Encompassing other methods

The answer is yes. In paired PIPA, equation (9) in Theorem A.1 can be formulated as DPO plus a margin term structured as $\log\frac{p^\theta(c=0|x)}{p^\theta(c=1|x)}$, which enables us to cover margin-based approaches similar to SimPO. Additionally, by introducing a power prior into PIPA, we can encompass DPO variants with $\beta\neq 1$. See section 1.1 in the response to reviewer 4Kmz for details.

6. Why math tasks?

Here we provide further explanation on our choices of using math tasks in this paper.

6.1 Math is a natural fit

Given that our PIPA framework provides key advantages—such as removing the requirement for paired data and naturally extending to step-wise labels—we prioritize math tasks as the optimal testbed for validation. This enables us to benchmark PIPA against general baselines like DPO and more specialized approaches such as step-DPO and step-KTO, which are designed to tackle these challenges.

6.2 Easy generalization to other tasks

We appreciate the suggestion to expand the scope of our study. Our PIPA framework is applicable to general preference tasks, particularly through the paired PIPA formulation outlined in Theorem A.1 of our paper. Preliminary results are presented in Table 4 in response to Review 3B8U. Further exploration about paired PIPA and other tasks will be explored more in future work.

6.3 The importance of paired data in math

While some math benchmarks use rule-based rewards, a preference signal that distinguishes correct from incorrect answers remains crucial. It helps the model recognize and learn from flawed reasoning or incorrect steps within the same problem [1, 2, 3…]. The impact of paired data is more pronounced when high-quality paired data is available. Thus, when designing alignment algorithms, it is essential to consider both paired and unpaired data. PIPA serves as a promising step in exploring this direction. See more discussions in Section 4.2 in the reply to 4Kmz.

We appreciate the reviewer for acknowledging our innovative framework and thorough experiments for mathematical tasks.

1. Unified derivation of PIPA-M and PIPA-N

Both PIPA-M and PIPA-N derive from the same MLE target, differing only in their prior assumptions. For data (x, y, c) where:

• x: question
• y: answer
• c: label (1=pos, 0=neg)

We want to use MLE to solve $p^{\theta}(c | x,y)$. By Bayes’ rule: $p^{\\theta}(c=1 | x, y) = \\frac{p^{\\theta}(y|x, c=1) \\cdot p^{\\theta}(c=1|x)}{p^{\\theta}(y|x)} \quad \\text{and} \quad p^{\\theta}(c=0 |x, y) = 1-p^{\\theta}(c=1 | x,y)$ Now how to parameterize the 3 terms in $p^{\\theta}(c=1 | x, y)$? In both PIPA-M and N, the 2 terms in numerator are learnable:

• $p^{\theta}(y|x, c=1)$ corresponds to $f^{\theta}(y|x)$ in our paper
• $p^{\theta}(c=1|x)$ corresponds to $g^{\theta}(x)$ in our paper

The key difference lies in how we handle the denominator $p^{\theta}(y|x)$:

• PIPA-M: Uses a fixed prior for the margin distribution: $p^{\theta}(y|x) = p^{prior}(y|x)$
• PIPA-N: Uses a fixed prior for the negative conditional distribution. So we expand the marginal distribution to conditional distribution:

This concise derivation illustrates the relationship between PIPA-M and PIPA-N. The reviewer questions the relationship between (2) and (5). We point out here that they are the same problem with different priors. Equation (5) for PIPA-N is exactly the MLE above. The equivalence between the above MLE for PIPA-M and the KL objective (2) is from the marginal prior and is demonstrated in Lines 144-147 of our paper. In our next revision, we will improve the presentation to provide a more unified description of both approaches.

2. Empirical understanding about performance difference

As mentioned in Lines 318-329, our empirical findings suggest different optimal scenarios for each approach. Below, we provide additional understanding for these observations:

2.1 From the prior

• PIPA-N performs better for self-alignment: When data is generated on-policy from previous model versions, our goal is self-improvement. In this context, it's reasonable to assume the current model $p^{prior}$ better resembles the negative-condition distribution.
• PIPA-M performs better under distribution shift: When using data from external sources, the distribution shift makes it more appropriate to use the current model as a prior for the overall distribution rather than specifically for the negative class, which would be an aggressive assumption.

2.2 From the loss

A deeper analysis of the loss functions of positive samples also show the differences:

PIPA-M Loss: $-\\log\\frac{p^{\\theta}(y|x,c=1)p^{\\theta}(c=1|x)}{p^{prior}(y|x)}$

This effectively maximizes $p^{\\theta}(y|x,c=1)$ and $p^{\\theta}(c=1|x)$ directly.

PIPA-N Loss: $-\\log \left(p^{\\theta}(y|x,c=1)p^{\\theta}(c=1|x)\right) + \\log \left(p^{\\theta}(y|x,c=1)p^{\\theta}(c=1|x) + p^{prior}(y|x)(1-p^{\\theta}(c=1|x)\right)$

PIPA-N incorporates an additional regularization term that prevents $p^{\\theta}(y|x,c=1)$ from overfitting to the training data. This regularization helps the model avoid getting stuck at the checkpoint that generated correct but suboptimal examples, which is particularly beneficial for self-improvement scenarios.

3. General preference tasks

Please see Table 4, in Seciton 2 of the response to 3B8U.

4. Why math tasks?

See Section 6 in the reply to Zrno.

5. Why use DPO/KTO for baselines?

DPO/KTO and their variants (iterative, stepwise, etc.) are widely acknowledged as the primary—and often the only—baselines in most alignment research on mathematical reasoning [1,2,3,4…]. We respectfully disagree with the reviewer’s suggestion to use RL methods like PPO/GRPO as baselines. These are on-policy algorithms, whereas the original DPO/KTO and our PIPA are all off-policy methods using a fixed dataset that differ only in their loss functions. Comparing off-policy methods with RL approaches would be inherently unfair. When comparing DPO/KTO in our work, we keep everything the same except the loss functions. A fair comparison with RL would require considering iterative/online DPO and developing an online version of PIPA, which extends far beyond the scope of this study.

6.Typos

We will fix the typos in Line 140, 151.

[1] Pang, Richard Yuanzhe, et al. "Iterative reasoning preference optimization." NeurIPS 2023.

[2] Chen, Guoxin, et al. “Step-level Value Preference Optimization for Mathematical Reasoning.” EMNLP 2024.

[3] Wang, Huaijie, et al. "Offline Reinforcement Learning for LLM Multi-Step Reasoning."

[4] Xiong, Wei, et al. “Building Math Agents with Multi-Turn Iterative Preference Learning.” ICLR 2025.

We appreciate the reviewer for the recognition of our theoretical contribution in unifying existing preference alignment methods under a single probabilistic framework, and our experiments demonstrating PIPA's effectiveness on math reasoning tasks with stepwise feedback.

1. Hyperparameter validation

Validation via learning trajectory
Following the paper referenced by the reviewer, we plot the trajectory of the forward KL divergence and the win rate computed by the implicit reward during training on Alphamath with Deepseek (https://anonymous.4open.science/r/PIPA-DB7B). Both metrics indicate that DPO effectively learns throughout the training process without exhibiting the over-parameterization issues highlighted in the referenced paper.

Validation via additional results
Our hyperparameters were carefully selected through rigorous search, ensuring that they were not cherry-picked to favor our results. To illustrate this, we present additional results for batch size, $\beta$, and learning rate—three key hyperparameters in DPO. The current settings in the paper are (256, 0.1, 5e-7). We here find a slightly larger $\beta$ than 0.1 allows for a higher learning rate and improves performance, leading to better results. However, DPO still falls short of PIPA. Specifically:

• Batch size: Consistent with previous studies, we use a fixed batch size across all algorithms to ensure a fair comparison by maintaining the same number of updates. We evaluate batch sizes of {64, 256, 1024} for DPO, IPO, and PIPA, and the results consistently show PIPA’s advantage at all batch sizes. We use 5e-7 for DPO and IPO here because larger LR will crash.(https://anonymous.4open.science/r/PIPA-DB7B/table1.md)

• $\beta$: Initially, we set $\beta$ to 0.1 following the original Alphamath paper [1]. We then explored values from {0.5, 1.0, 2.0, 5.0} and found that $\beta$ = 1.0 yielded the best results. However, DPO still lags behind PIPA. (https://anonymous.4open.science/r/PIPA-DB7B/table2.md)

• Learning rate: With optimal batch size and $\beta$, we try learning rates from {5e-7, 1e-6, 5e-6, 1e-5, 5e-5}. We find increasing $\beta$ to 1.0 allows for a higher optimal learning rate of 1e-5. This adjustment improves DPO's performance on the GSM8K dataset, reaching 0.74 and narrowing the gap with PIPA to 5%. However, it does not lead to any improvements on the more challenging MATH dataset. (https://anonymous.4open.science/r/PIPA-DB7B/table3.md)

2. Non-math preference tasks

We employed the paired version of PIPA-N as outlined in Theorem A.1, with further algorithmic details provided in Section 2 of response to Reviewer Zrno. Following the setup proposed in SimPO, we conducted comparisons between DPO and SimPO by fine-tuning Llama-3-8B-instruct on the Ultrafeedback dataset. We reproduced the DPO and SimPO baselines in our codebase using OpenRLHF, maintaining identical hyperparameters from the SimPO paper (e.g., batch size, epochs, $\beta$, $\gamma$), except for using LoRA (r=64, $\alpha$=16) with a higher learning rate, as the original lower learning rate is ineffective with LoRA. Our reproduction produced results matching those reported by SimPO. For our PIPA-N, current results are better than SimPO. It’s noteworthy that due to the limited rebuttal period, we did not have the resources to thoroughly optimize the hyperparameters for PIPA-N. In contrast, the SimPO results benefited from careful tuning (e.g., using $\beta$=2.5 and $\gamma$=1.375 for Llama-3-8B-instruct). This suggests that PIPA-N has the potential for further improvement with more extensive parameter tuning. More results will be provided later given the limited response period.

3. Comparison to SimPO

Table 2 shows that SimPO is similar to DPO in Alphamath which lags behind PIPA, and Table 4 shows that PIPA-N is also better than SimPO.

4. Missing citation

We will cite the paper by Zhihan Liu et al. in Line 421.

[1] Chen Guoxin, et al. AlphaMath Almost Zero: Process Supervision without Process, NeurIPS 2024.