Plan2Align

Test-time preference alignment for large language models: instead of baking preferences into the weights, plan over the response at inference time — draft, score how much each segment contributes, and recompose.

📄 Paper: Test-Time Alignment for Large Language Models via Textual Model Predictive Control — ICLR 2026 (earlier version: "Plan2Align: Predictive Planning Based Test-Time Preference Alignment in Paragraph-Level Machine Translation", arXiv:2502.20795)

Kuang-Da Wang, Teng-Ruei Chen, Yu Heng Hung, Guo-Xun Ko, Shuoyang Ding, Kris Wu, Yu-Chiang Frank Wang, Huck Yang, Wen-Chih Peng, Ping-Chun Hsieh

🔗 Official implementation: rl-bandits-lab/Plan2Align 🔗 Baselines (ARGS / RAIN / MetricX-24): Baseline_of_Plan2Align

This is my personal research repository for the project. For the packaged release that accompanies the paper, use the official lab repository above.

Method

Two strategies run side by side each iteration, both scored by the same reward model:

• MPC (Model Predictive Control) — iteratively refine the whole response using improvement prompts. The textual analogue of planning a horizon, committing one step, and re-planning.
• P2A (Plan-to-Align) — decompose the response into sentence-level segments, score each segment's contribution by ablation (reward(full) − reward(masked)), then assemble high-reward combinations from a per-slot segment buffer.

Sequence-level rewards get less informative the longer the output gets. Scoring segments by what they contribute is what makes the signal usable on long-form text — and it happens entirely at inference, so nothing about the model changes.

1. Generate initial responses (personas["original"])
2. Score all with the RM; keep the best as the MPC and P2A starting point
3. For each iteration:
   a. MPC: generate candidates from the current best → score → pick best
   b. P2A: compose from the segment buffer → generate → score → pick best
   c. Update the buffer with new segments
4. Save history[p0..pN]

Algorithm variants

All three are one class, selected by config — they differ only in how the buffer is built and how segments are chosen.

	v1	v2	v3
Segment contribution scoring	—	ablation	ablation
Buffer	flat, score-sorted	per-slot, contribution-sorted	per-slot, contribution-sorted
P2A composition	random assignment across candidates	greedy / softmax per slot	Cartesian product, globally best combination
Threshold filtering	✓	✓	✓
Multi-persona MPC	✓	✓	✓
Used for	earliest experiments, HelpSteer	HH-RLHF	main HelpSteer / UltraFeedback results

v3 is the default: rather than picking the best segment slot by slot, it scores every complete combination and takes the global optimum.

Segment contributions

Each segment is scored by what the response loses without it — reward(full) − reward(masked). That per-segment signal is what the buffer is sorted on and what the composition step selects over.

Layout

p2a/                  Unified package
  config.py           P2AConfig dataclass + CLI parsing
  utils.py            Model loading, generation, segmentation
  algorithm.py        Plan2Align class (v1/v2/v3 via config)
  evaluate.py         Best-iteration aggregate + win/tie/loss
run.py                Entry point (HelpSteer / HH-RLHF / UltraFeedback)
run_math.py           Entry point for math reasoning (ReasonEval-7B as RM)
config/               Persona prompts; API key templates
reward/               Reward-model tooling — inference, best-of-n,
                      accuracy checks, coherence, GPT-4 judging
  reward_distribution/  Reward distributions for the trained RM (below)
datasets/             Dataset preparation scripts

Supported datasets: HelpSteer, UltraFeedback, HH-RLHF, and math reasoning (GSM8K / SVAMP / AQuA-RAT / TheoremQA).

Reward model

Everything downstream depends on the reward model separating good from bad responses, so it is worth checking before trusting any alignment result. reward/reward_distribution/ holds those distributions for the trained RM:

Train	Test

*_score_good_dist.png and *_score_bad_dist.png split the same scores by preference label; count_q.py regenerates them from train_rewards.csv / test_rewards.csv.

Setup

pip install torch transformers pandas numpy openai tqdm

cp config/key.txt.example config/key.txt        # only needed for --use_api
cp config/key_math.txt.example config/key_math.txt

Datasets and results are not in this repository — HelpSteer, UltraFeedback, HH-RLHF, the translation sets and the full experiment outputs run to several hundred MB. Prepare them with the scripts in datasets/, or take them from the official repository.

Running

# v3 — HelpSteer, local model
python run.py --version v3 --rm_path rl-bandits-lab/helpsteer_rm --lm_cuda 1 --rm_cuda 0

# v3 — HelpSteer via DeepInfra API
python run.py --use_api --rm_cuda 0

# v1 — HH-RLHF
python run.py --version v1 --dataset HH --lm_cuda 0 --rm_cuda 1

# v2 — greedy per-slot selection
python run.py --version v2 --sampling_mode greedy --threshold 2.0 \
              --buffer_width 4 --buffer_depth 3

# Math reasoning (ReasonEval-7B as reward model)
python run_math.py --rm_path GAIR/ReasonEval-7B --lm_cuda 1 --rm_cuda 0

Main options (full list in p2a/config.py):

Flag	Default
--version	v3	v1 / v2 / v3
--sampling_mode	softmax	greedy / softmax (v2)
--threshold	1.0	drop segments below this contribution
--max_iterations	4	planning iterations
--buffer_width / --buffer_depth	4 / 3	segment buffer shape
--lm_path	meta-llama/Meta-Llama-3.1-8B-Instruct
--rm_path	rl-bandits-lab/helpsteer_rm
--dataset	helpsteer	helpsteer / HH / math
--use_api	off	DeepInfra instead of a local model

P2AConfig.__post_init__ resolves default dataset paths against $P2A_DATASET_ROOT (falls back to ./datasets). Set the env var, pass --input_file explicitly, or edit p2a/config.py.

Evaluate

python -m p2a.evaluate --folder results/run_0 --max_index 694 --it 4 --output summary.csv

Average MPC score: 3.2456
Average P2A score: 3.8912
P2A vs MPC — Win: 482 / Tie: 31 / Lose: 181

Output format

Each run writes per-prompt JSON under results/run_N/, keyed by iteration ("0" = initial responses):

{
  "0": {
    "prompt": [{"role": "user", "content": "..."}],
    "mpc_response": "...", "mpc_score": 2.5,
    "p2a_response": "...", "p2a_score": 2.8
  }
}

p2a.evaluate picks the best iteration per method per prompt.

Citation

@inproceedings{wang2026testtime,
  title     = {Test-Time Alignment for Large Language Models via Textual Model Predictive Control},
  author    = {Wang, Kuang-Da and Chen, Teng-Ruei and Hung, Yu Heng and Ko, Guo-Xun and
               Ding, Shuoyang and Wu, Kris and Wang, Yu-Chiang Frank and Yang, Huck and
               Peng, Wen-Chih and Hsieh, Ping-Chun},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026},
  url       = {https://arxiv.org/abs/2502.20795}
}