IPO / SimPO Loss

Background

IPO (Identity Preference Optimisation, Azar et al.) replaces DPO's monotone log-sigmoid loss with a bounded regression target. SimPO adds length normalisation. Both fix DPO's known reward-hacking failure mode where the model exploits long answers to maximise the log-prob margin.

Problem statement

Implement ipo_loss(chosen_logprobs, rejected_logprobs, lengths_chosen, lengths_rejected, beta=0.1) returning the mean loss:

$$\text{margin}_i \;=\; \frac{\log\pi(\text{chosen}_i)}{|\text{chosen}_i|} - \frac{\log\pi(\text{rejected}_i)}{|\text{rejected}_i|},\qquad \mathcal L \;=\; \frac{1}{N}\sum_i \Bigl(\text{margin}_i - \tfrac{1}{2\beta}\Bigr)^{\!2}$$

The square term gives a quadratic penalty that vanishes at the target margin; length normalisation kills the long-answer bias.

Input

• chosen_logprobs — 1-D array, summed log-probs of chosen completions.
• rejected_logprobs — 1-D array, summed log-probs of rejected.
• lengths_chosen — 1-D array of int, token counts of chosen.
• lengths_rejected — 1-D array of int, token counts of rejected.
• beta — float > 0, target-margin scale.

Output

Returns a Python float >= 0.

Examples

Example 1 — equal log-probs → margin 0 → loss = $(1/2\beta)^2$

Input:  chosen=[-10.0], rejected=[-10.0], lengths=[10, 10], beta=0.1
Output: (5.0)^2 = 25.0

Example 2 — chosen better-normalised → lower loss

Input:  chosen at -1 over length 5, rejected at -5 over length 5
Output: lower than Example 1

Constraints

• All four arrays must have the same length.
• lengths_* strictly positive.
• Returns a plain Python float.

Notes

• vs DPO. DPO's loss is monotone — it never stops pushing the margin wider. IPO's quadratic loss has a target, so over-optimisation hurts. Empirically more stable on small fine-tunes.
• SimPO twist. SimPO drops the reference model entirely and uses just the policy log-probs with length normalisation. Same length-norm idea, simpler objective.