Reproduce “On Almost Surely Safe Alignment of Large Language Models at Inference-Time” on OpenResearch

About this paper

Xiaotong Ji, Shyam Sundhar Ramesh, Matthieu Zimmer, Ilija Bogunovic, Jun Wang, Haitham Bou Ammar

Researchers from Huawei Noah's Ark Lab and UCL developed InferenceGuard, an inference-time alignment method for Large Language Models that guarantees "almost surely safe" responses without modifying model weights. This approach achieved up to 100% safety rates on benchmarks like PKU-SafeRLHF and HH-RLHF, consistently outperforming other methods while preserving task performance and incurring manageable latency.

Key insights

• Formalizing LLM safety within a state-augmented cMDP framework allows for theoretically provable 'almost sure safety' guarantees by intrinsically penalizing constraint violations infinitely.
• Theoretical extensions demonstrate that an optimal policy found in the augmented latent space of an LLM translates to an optimal and almost surely safe policy in the original token generation space.
• A small, efficient critic operating on the LLM's hidden states and logits can effectively guide inference-time generation to achieve a strong balance between safety and task performance without costly model fine-tuning.

Summary by alphaXiv.