Target Policy Optimization — a simpler alternative to PPO for RLHF.
Based on arXiv:2604.06159 (Kaddour, 2026).
TPO is an RLHF algorithm. It is not a replacement for cross-entropy (CE) training. CE and TPO do different things:
| Cross-Entropy (SFT) | PPO | TPO | |
|---|---|---|---|
| Purpose | Predict next token | Optimize reward via RL | Optimize reward via RL |
| Needs labeled data | Yes | No | No |
| Needs reward model | No | Yes | Yes |
| Needs value/critic head | No | Yes | No |
| Needs clipping | No | Yes | No |
| Needs importance ratios | No | Yes | No |
| Training stability | Stable | Often unstable | Stable |
CE is supervised learning. TPO is reinforcement learning. They solve different problems. You don't compare them.
TPO competes with PPO, not CE. The advantage: TPO gives you the same RLHF capability with much less complexity — no value function, no clipping, no importance sampling ratios.
PPO: loss = -min(ratio * A, clip(ratio, 1-eps, 1+eps) * A) # needs V(s), clipping
TPO: loss = -target_prob * log P_policy(token) # needs nothing extra
pip install -e .
python
from tpo_torch import TPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
tokenizer.pad_token = tokenizer.eos_token
trainer = TPOTrainer(
model=model,
ref_model=ref_model,
beta=0.1,
train_dataset=dataset,
processing_class=tokenizer,
)
trainer.train()
┌──────────────────────┐
│ Reference Model │
│ (frozen) │
└──────────┬───────────┘
│ ref_logits
▼
Prompt + Labels ──▶ ┌──────────────────────┐
│ TPO Loss Function │
│ │
│ 1. log-odds(P_ref) │
│ 2. + advantage/beta │
│ 3. sigmoid -> target │
│ 4. CE loss vs target │
└──────────┬───────────┘
│ loss
▼
┌──────────────────────┐
│ Policy Model │
│ (trained) │
└──────────────────────┘
| Function | Description |
|---|---|
tpo_loss_from_logits(logits, ref_logits, labels, advantages, beta) |
|
| Loss from raw model logits | |
tpo_loss(policy_logprobs, ref_logprobs, advantages, beta) |
|
| Loss from pre-computed log-probs | |
TPOTrainer(model, ref_model, beta, ...) |
|
| HuggingFace Trainer with TPO | |
TPODataCollator(tokenizer) |
|
| Preserves advantages in batches |
All benchmarks run on NVIDIA RTX 3050 Laptop (4GB VRAM).
Both are RLHF methods evaluated on the same metric: held-out perplexity.
| Method | Final Perplexity | Implementation Complexity |
|---|---|---|
| TPO (beta=0.1) | 74.51 | |
| Loss function only | ||
| PPO-clip (simplified) | 825.66 | Loss + ratio + clipping |
Note: Our PPO-clip is simplified (no value function, no GAE). Full PPO with a critic would perform better but requires significantly more code and compute. TPO achieves competitive results with none of that infrastructure.
Zero NaNs at advantage values from 0.01 to 1000- Max gradient norm: 0.09— stable across all regimes
| Seq Len | Latency | Throughput |
|---|---|---|
| 32 | 4.17ms | 61,345 tok/s |
| 64 | 8.26ms | 62,021 tok/s |
| 128 | 16.50ms | 62,046 tok/s |
| 256 | 32.81ms | 62,418 tok/s |
| 512 | 65.42ms | 62,614 tok/s |
| 1024 | 929.89ms | 8,810 tok/s |
python benchmarks/run_benchmarks.py
tpo train --max-steps 10
tpo bench
tpo info
git clone https://github.com/Griffith-7/tpo-torch.git
cd tpo-torch
pip install -e ".[dev]"
Run tests:
pytest tests/ -v
ruff check tpo_torch/
- Python >= 3.9
- PyTorch >= 2.0
- transformers >= 4.40
@misc{kaddour2026targetpolicyoptimization,
title={Target Policy Optimization},
author={Jean Kaddour},
year={2026},
eprint={2604.06159},
archivePrefix={arXiv},
primaryClass={cs.LG}
}