Test-Time Safety Alignment

TL;DR

This paper demonstrates that optimizing input word embeddings via zeroth-order gradient estimation can effectively reduce harmful outputs of aligned language models, enhancing safety control.

Contribution

It introduces a novel method using black-box API feedback to optimize embeddings for safety, applicable to complex bimodal response distributions.

Findings

The method neutralizes all safety-flagged responses on benchmarks.

Input embeddings can be optimized in a sub-lexical manner for safety.

Zeroth-order gradient estimation effectively guides embedding optimization.

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.