Inf-guard

This repository accompanies the paper:

On Almost Surely Safe Alignment of Large Language Models at Inference-Time Xiaotong Ji, Shyam Sundhar Ramesh, Matthieu Zimmer, Ilija Bogunovic, Jun Wang, Haitham Bou Ammar. Transactions on Machine Learning Research (TMLR), 2025. arXiv:2502.01208

It implements inf_guard, an inference-time decoding method that steers a frozen base LLM toward responses that are safe with high probability, without any further training of the base model. Generation is built on beam search, and the beams are scored by a reward model and a cost/safety model.

How it works

1. A frozen base LLM (Alpaca-7B, Beaver-v3-7B, or Vicuna-7B) generates a set of beams / continuations with vLLM.
2. A reward model (argsearch/llama-7b-rm-float32) and a cost / safety model (PKU-Alignment/beaver-7b-unified-cost) score each candidate.
3. Beam search selects safe, high-reward completions under a per-step safety budget.

Methods & baselines

The repository ships inf_guard (ours) alongside safety-extended versions of the standard test-time search baselines. Each row is a recipe under recipes/<model>/; method: in the YAML sets the safety strategy, and the iteration/token settings select the search procedure (Best-of-N = one iteration of N candidates; beam search = multi-step).

Recipe(s)	Search	Safety strategy	Description
best_of_n_saute, beam_search_saute	Best-of-N / beam search	Safety augmentation (Sauté)	Baseline search that folds the running safety budget into the state, penalising candidates as the budget is spent.
best_of_n_lagrange, beam_search_lagrange	Best-of-N / beam search	Lagrangian	Baseline search that ranks candidates by c_task + λ · C_safety — task cost plus a Lagrangian-weighted safety cost.
inf_guard	Beam search	Ours	Inference-time guard with an almost-surely-safe search; no learned critic.
inf_guard + use_critic	Beam search	Ours	As above, with a learned critic supplying safety value estimates
inf_guard + use_kv_critic	Beam search	Ours	As above, with a KV-cache critic.

On the baselines. Best-of-N (BoN) and beam search were originally designed to maximise reward without regard for safety. Here they are extended to respect safety in two ways: a Lagrangian penalty (select by c_task + λ · C_safety) and a safety-augmentation scheme (Sauté), which the paper finds more effective than the Lagrangian approach for balancing reward against constraint satisfaction.

Installation

All experiments were run on a single CUDA device with ≥ 64 GB of on-device memory (each run loads the frozen base model plus the reward and cost/safety models; see Appendix D.1 of the paper). Runs use Python 3.11 and bfloat16.

conda create -n inf-guard python=3.11 && conda activate inf-guard
pip install -r requirements.txt
pip install -e .

Important

vLLM must use the XFORMERS attention backend, or runs abort with a CUDA illegal-memory-access error (there is an explicit assertion in scripts/test_time_compute.py). Export this before running:

export VLLM_ATTENTION_BACKEND=XFORMERS

Quickstart

Run a single recipe:

export VLLM_ATTENTION_BACKEND=XFORMERS
python scripts/test_time_compute.py recipes/beaver-7b/inf_guard.yaml

Safety Rate / reward / cost are printed to stdout, and completions are written to logs/<timestamp>.jsonl.

Each recipe is a YAML file that overrides the defaults in src/sal/config.py. See the recipes README for the full table of configurations and a launcher that runs all experiments and baselines in one go.

Critic weights

The inf_guard_critic and inf_guard_kv_critic recipes require a trained critic checkpoint, which can be trained with the code in training/ (both the non-KV and KV-cache critics), then place the resulting .pth files under critics/ following the per-model layout that the recipes expect, e.g.:

critics/
├── beaver-7b/critic.pth
├── beaver-7b/kv_critic.pth
└── vicuna-7b/critic.pth

Datasets

Experiments use PKU-SafeRLHF, HEx-PHI, and HH-RLHF (loaded from the Hugging Face Hub; the default in config.py is PKU-Alignment/PKU-SafeRLHF).

Project structure

├── LICENSE
├── README.md
├── requirements.txt        <- Pinned runtime environment
├── setup.py                <- Makes `sal` importable (pip install -e .)
├── recipes/                <- One YAML per experiment configuration (+ README)
├── scripts/                <- test_time_compute.py (entry point) + run_all_experiments.sh
├── src/sal/                <- Search, config, reward/cost/critic models, utils
├── src/score_model/        <- Vendored safe-rlhf cost-model backend (LLaMA)
└── training/               <- Critic training pipeline (code + configs + README)

Citation

If you use this code, please cite:

@article{ji2025almostsurely,
  title   = {On Almost Surely Safe Alignment of Large Language Models at Inference-Time},
  author  = {Ji, Xiaotong and Ramesh, Shyam Sundhar and Zimmer, Matthieu and
             Bogunovic, Ilija and Wang, Jun and Bou Ammar, Haitham},
  journal = {Transactions on Machine Learning Research},
  year    = {2025},
  note    = {arXiv:2502.01208},
  url     = {https://arxiv.org/abs/2502.01208}
}

License & acknowledgements

Released under the Apache License 2.0. The search scaffolding is derived from huggingface/search-and-learn and the cost-model code under src/score_model/ is vendored from PKU-Alignment/safe-rlhf; both are Apache-2.0. Copyright for the modifications in this repository belongs to the paper authors.