Test-Time Safety Alignment

2026-04-28

AI Summary

Although current aligned language models can reject harmful requests, their bimodal output distributions may still yield semantically harmful responses. This work proposes a test-time safety control method that optimizes subword-level input embeddings to steer the model toward safer outputs. Leveraging a black-box text moderation API for zeroth-order gradient estimation, the approach performs gradient descent on the input embeddings to minimize semantic harmfulness. As the first study to apply input embedding optimization to safety control in aligned models, this method overcomes prior limitations that restricted such techniques to open-ended generative models or superficial content filtering. It achieves a 100% harmful response elimination rate on standard safety benchmarks.

Abstract

Recent work has shown that a model's input word embeddings can serve as effective control variables for steering its behavior toward outputs that satisfy desired properties. However, this has only been demonstrated for pretrained text-completion models on the relatively simple objective of reducing surface-level profanity in short continuations. A natural and practically important question is how well input embeddings can control aligned models, which produce an imbalanced bimodal refuse-or-comply output distribution rather than the smooth distribution characteristic of open-ended generation. We explore this in the context of safety, showing that input word embeddings can be optimized in a sub-lexical manner to minimize the semantic harmfulness of aligned model responses. Our approach uses zeroth-order gradient estimation of a black-box text-moderation API with respect to the input embeddings, and then applies gradient descent on these embeddings to minimize the harmfulness of the generated text. Experiments show that the proposed method can neutralize every safety-flagged response on standard safety benchmarks.

Problem

• Test-Time Safety Alignment
• Input Embeddings
• Aligned Models
• Semantic Harmfulness
• Safety Benchmarks

Innovation

• Test-Time Alignment
• Input Embedding Optimization
• Zeroth-Order Gradient Estimation
• Safety Alignment
• Black-Box Moderation