💬 NLP

Test-Time Safety Alignment

This paper demonstrates that input word embeddings can be optimized via zeroth-order gradient estimation to steer aligned language models away from harmful refusals or unsafe outputs, effectively neutralizing safety-flagged responses on standard benchmarks.

Original authors: Baturay Saglam, Dionysis Kalogerias

Published 2026-04-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Baturay Saglam, Dionysis Kalogerias

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, well-trained robot assistant. You've taught it to be polite, helpful, and safe. It knows not to give you instructions on how to build a bomb or write hate speech. However, clever tricksters have found ways to "jailbreak" this robot. They wrap their dangerous requests in fancy costumes—like asking the robot to "pretend to be a villain in a movie" or "write a story about a bad guy"—to trick the robot into forgetting its safety rules and spitting out harmful content.

This paper introduces a new way to stop the robot from falling for these tricks, but instead of retraining the robot or adding a new security guard, they tweak the robot's "brain" right before it answers.

Here is how it works, using simple analogies:

The Problem: The Robot's "Safety Switch" is Brittle

Usually, when a robot is asked a question, it looks at the words you typed and picks the next word from its vocabulary list. If you ask a tricky question, the robot might get confused and pick a "bad" word.

Current safety methods try to fix this by either:

Rewriting your question (like a human editor changing your words).
Adding a security guard (a separate AI that reads your question and blocks it).
Tweaking the robot's internal gears (changing how it thinks).

The authors say these methods are a bit clunky. They want a way to fix the robot's answer without changing the words you see, and without needing a separate security guard.

The Solution: "Sub-lexical" Tuning (The Invisible Dial)

The authors realized that before the robot reads your words, it actually converts them into a long list of numbers (called embeddings). Think of these numbers as the robot's internal "coordinates" for each word.

Usually, the word "Cat" is just one specific set of coordinates. But the authors discovered that you can nudge those coordinates just a tiny, invisible amount—so small that if you turned the numbers back into words, it would still say "Cat."

The Analogy: Imagine the robot's brain is a giant map. The word "Cat" is a specific town on that map.

Normal operation: You point exactly at the center of the town.
The Attack: A trickster points the robot at a "danger zone" just outside the town, tricking it into thinking "Cat" means something dangerous.
The Fix: The authors' method gently nudges the "Cat" coordinates within the town limits, but in a slightly different direction. It's like turning a dial on a radio. You are still listening to the same station ("Cat"), but you've adjusted the frequency just enough to clear up the static and stop the bad signal.

How They Do It: The "Blind Taste Test"

The robot is a "black box," meaning the researchers can't see inside its code to fix it directly. So, they use a clever trial-and-error method:

The Oracle (The Taste Tester): They use a public safety tool (like a content filter) that acts as a "taste tester." It looks at the robot's answer and gives it a "harm score." High score = bad; low score = safe.
The Guess-and-Check: The researchers take the robot's input, add a tiny bit of random "noise" (like shaking a dice), and ask the robot what it thinks. They check the harm score.
The Nudge: If the answer is still a bit risky, they use math to figure out which direction to nudge the invisible coordinates to lower the harm score.
Repeat: They do this a few times (usually just 1 or 2 steps). It's like tuning a guitar string: you pluck it, listen, tighten it a tiny bit, and listen again until it sounds perfect.

The Results: Magic, but Real

The paper tested this on several different robot models (from small to huge) and found:

It Works: It successfully stopped the robot from giving harmful answers in almost every case, even when the tricksters used very clever disguises.
It's Invisible: The robot still answers the question normally. If you asked, "How do I bake a cake?" it still tells you how to bake a cake. It just refuses to tell you how to bake a bomb, even if you asked in a tricky way.
It's Fast: It only takes a couple of seconds per question.
It Doesn't Break Good Answers: If you asked a safe question, the robot's answer didn't get worse. It didn't start refusing to answer safe questions just to be safe.

The Big Takeaway

The paper shows that the robot's "brain" (its input numbers) is much more sensitive than we thought. You don't need to change the words a human reads to change the robot's behavior. You just need to turn the invisible dials underneath the words.

It's like having a car that drives itself safely. If someone tries to trick the car into driving off a cliff, you don't need to rebuild the car or hire a new driver. You just need to make a microscopic adjustment to the steering wheel's internal sensors, and the car naturally steers itself back to safety.

1. Problem Statement

Large Language Models (LLMs) undergo safety training (e.g., RLHF, DPO) to refuse harmful requests. However, adversarial prompts (jailbreaks) can bypass these defenses by exploiting the gap between the model's training distribution and the unbounded input space. These prompts often use structural framing (roleplay, obfuscation) to elicit harmful content.

Existing test-time defenses face structural limitations:

Prompt-level strategies manipulate discrete tokens, which are brittle against semantic attacks.
Activation steering requires access to hidden states and often relies on fixed calibration sets that may not generalize.
Guard models require separate inference and often follow a "detect-then-filter" paradigm rather than steering generation.

The Core Question: Can input word embeddings be optimized in a sub-lexical manner (modifying the continuous vector representation without changing the visible token) to steer aligned models toward safe outputs, even when the model produces a bimodal "refuse-or-comply" distribution rather than a smooth continuation?

2. Methodology: Test-Time Safety Alignment (TSA)

The proposed method, TSA, operates entirely at test time without modifying model weights, requiring no auxiliary training, and treating the LLM as a black box. It optimizes the prompt's embedding matrix directly.

A. Objective Function

The goal is to minimize the harmfulness of the generated text $y = f(X)$ , where $f$ is the LLM and $X$ is the prompt embedding matrix.

Harmfulness Oracle ( $h$ ): The method uses the OpenAI Moderation API as a black-box oracle. It scores outputs across 13 harm categories (e.g., violence, self-harm, hate).
Loss Function: $\Phi(X) = \max_{c} s_c(f(X))$ , where $s_c$ is the score for category $c$ . The objective is to minimize this maximum score, ensuring the response is safe across all dimensions.

B. Zeroth-Order Gradient Estimation

Since the LLM ( $f$ ) and the Moderation API ( $h$ ) are black boxes, backpropagation is impossible. TSA uses zeroth-order gradient estimation:

Perturbation: At each iteration $k$ , the current embedding matrix $X_k$ is perturbed by adding Gaussian noise $U$ scaled by $\mu$ : $X_k + \mu U$ .
Finite Difference: The gradient is estimated by comparing the loss of the perturbed inputs against the base input:
$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$
where $N$ is the number of Monte Carlo samples.
Update: The embeddings are updated via gradient descent: $X_{k+1} \leftarrow X_k - \eta g_k$ .

C. Regularization

To ensure the optimized embeddings remain valid and do not drift into untrained regions of the latent space:

Gradient Normalization: The update step uses the normalized gradient direction ( $g_k / \|g_k\|_2$ ) to control step magnitude.
Cosine Similarity Constraint: The updated embedding is projected back onto a cosine ball around the original embedding $X_0$ if the similarity drops below a threshold $\kappa$ . This ensures the perturbation is sub-lexical: the optimized vector decodes back to the original token, but the continuous representation within that token's Voronoi cell is altered.

3. Key Contributions

Extension to Aligned Models: Demonstrates that embedding-level control, previously shown for detoxifying base models, effectively works on safety-aligned instruction-tuned models, which have a fundamentally different (bimodal) output distribution.
Sub-Lexical Control: Proves that optimizing the continuous embedding vector within the nearest-neighbor cell of a token (without changing the token itself) is sufficient to redirect the model's behavior from harmful to safe. This challenges the view of embeddings as mere discrete lookups.
Black-Box Optimization: The method requires no access to model internals (gradients, hidden states, or logits). It only requires access to input embeddings and a black-box moderation API.
Unconditional Application: The method applies to all prompts regardless of intent. It does not force a refusal; it steers the model to generate a safe, substantive response if possible, or a safe refusal if necessary.

4. Experimental Results

The method was evaluated on five instruction-tuned models (ranging from 1B to 20B parameters: Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) across two benchmarks:

WildJailbreak: Adversarial prompts (harmful and benign).
HarmBench: Direct harmful queries.

Key Findings:

Safety Elimination: TSA reduced the number of flagged (harmful) responses to zero for most models on standard benchmarks. For example, on Llama 3.1-8B, it reduced flagged harmful prompts from 218 to 0.
Convergence: The optimization typically converges within 1 to 2 gradient steps.
Benign Preservation: On adversarial benign prompts, TSA did not increase moderation scores, confirming it does not induce a "blanket refusal" bias or degrade helpfulness.
Comparison to Baselines: TSA significantly outperformed baselines like SmoothLLM (random token perturbation), AdaSteer (activation steering), and RESTA (embedding noise). Baselines often failed on semantic jailbreaks or increased flagged rates on direct harmful queries.
Efficiency: The primary latency cost is the API call to the moderation oracle, not the model inference.

5. Significance and Implications

New Control Surface: The paper reveals that the embedding layer offers a richer control surface than previously thought. Small, sub-lexical perturbations can navigate the complex safety boundaries of aligned models.
Deployment Viability: As a "drop-in" safety layer, TSA requires no model retraining or fine-tuning. It can be deployed on any model where input embeddings are accessible (e.g., open-weight models) combined with a public moderation API.
Limitations:
- Latency: The iterative optimization adds per-prompt latency (seconds), though this is dominated by API round-trips.
- Model Dependency: Effectiveness scales with the quality of the model's existing safety training. Models with thin alignment (e.g., Gemma 3-1B) showed reduced success, as the optimizer could not find a nearby safe direction.
- API Access: Requires access to input embeddings (not available for purely API-gated closed models) and a moderation API.

In conclusion, Test-Time Safety Alignment provides a robust, model-agnostic mechanism to neutralize adversarial attacks by optimizing the continuous input representation, effectively bridging the gap between theoretical safety and practical deployment without retraining.