Aligning Language Models with Selective Prediction

Aligning Language Models with Selective Prediction
Gaoxiang Luo, Yifan Wu, Sinian Zhang, Aryan Deshwal, Ju Sun
Second Workshop on Agents in the Wild: Safety, Security, and Beyond (ICML 2026)

[Paper] [BibTeX]

Environment

Built on RLCR, which is vendored as a submodule at RLCR/ — clone recursively:

git clone --recurse-submodules https://github.com/sun-umn/RLSR.git

Dependencies are managed by uv, pinned in uv.lock (python 3.12, torch 2.8.0, transformers 4.57.0, vllm 0.11.0, trl 0.23.1).

uv sync
# flash-attn must be built against the installed torch version:
MAX_JOBS=64 uv pip install flash-attn --no-build-isolation --force-reinstall --no-deps --no-cache --no-binary flash-attn

Training (HotPotQA)

One-time: apply rlcr.patch to the submodule (trl 0.23 compatibility + answer-normalization fixes; see the patch header comments):

cd RLCR && git apply ../rlcr.patch && cd ..

Train (4× A100-40GB GPUs; select them via CUDA_VISIBLE_DEVICES):

export ACCELERATE_CONFIG_FILE=$PWD/deepspeed.yaml
uv run accelerate launch --num_processes 4 --config_file deepspeed.yaml \
  rl_runner.py --config configs/Qwen-7B/hotpot/RLSR.yaml   # or RLVR.yaml / RLCR.yaml

Equivalently, ./launch.sh <rlvr|rlcr|rlsr|rlvr-ta> runs one method; outputs default to data/<METHOD>-hotpot. Runs log to the RLSR wandb project under your default entity — wandb login once first, or set WANDB_MODE=offline.

Evaluation (HotPotQA)

eval.sh evaluates one checkpoint per invocation, following RLCR's official protocol and adding AURC / selective-prediction metrics. Run it on 2× A100-40GB GPUs (the generator and the LLM judge are each served tensor-parallel across them):

# <model> <dataset|all|all-extended> <gpu_ids> <method>
uv run bash eval.sh /path/to/checkpoint-300 all 0,1 rlsr
uv run bash eval.sh Qwen/Qwen2.5-7B all 0,1 base

• method must match how the model was trained: rlvr | rlcr | rlsr | rlvr-ta | base.
• all runs RLCR's 8 tasks (hotpot-eval-em, hotpot-vanilla-eval-em, gsm8k, math-500, commonsenseqa, gpqa, simpleqa, trivia); all-extended adds medqa-eval and mbpp, which need uv run python data_scripts/medqa.py and data_scripts/mbpp.py once first. Pass a single task name to run just one. (mbpp executes model-generated Python — run it only where that is safe.)
• Results land in eval/results/<dataset>/ under the evaluated model; the full path is printed when the run finishes.

Citation

@inproceedings{luo2026aligning,
  title = {Aligning Language Models with Selective Prediction},
  author = {Luo, Gaoxiang and Wu, Yifan and Zhang, Sinian and Deshwal, Aryan and Sun, Ju},
  booktitle = {Forty-third International Conference on Machine Learning Second Workshop on Agents in the Wild: Safety, Security, and Beyond},
  year = {2026},
  url = {https://openreview.net/forum?id=bqTLnXswlC},
}

Acknowledgements

This codebase builds directly on RLCR, whose repo we vendor as a submodule and whose HotPotQA protocol we follow. Our rank-1 LoRA recipe follows Thinking Machines Lab's LoRA Without Regret.