# Test-Time Safety Alignment

arXiv: https://arxiv.org/abs/2604.26167 · HTML: https://mlsanity.com/paper/2604.26167/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Safety alignment and jailbreak attacks, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 31 of its 36 references that are in the map.
- First posted on arXiv 28 Apr 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Safety alignment and jailbreak attacks](https://mlsanity.com/alignment/safety-alignment-and-jailbreak-attacks/)
- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [HarmBench: A Standardized Evaluation Framework for Automated Red Teaming and Robust Refusal](https://mlsanity.com/paper/2402.04249/) (2024)
- [Jailbreaking Black Box Large Language Models in Twenty Queries](https://mlsanity.com/paper/2310.08419/) (2023)
- [Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations](https://mlsanity.com/paper/2312.06674/) (2023)
- [WildGuard: Open One-Stop Moderation Tools for Safety Risks, Jailbreaks, and Refusals of LLMs](https://mlsanity.com/paper/2406.18495/) (2024)
- [SmoothLLM: Defending Large Language Models Against Jailbreaking Attacks](https://mlsanity.com/paper/2310.03684/) (2023)
- [Detecting Language Model Attacks with Perplexity](https://mlsanity.com/paper/2308.14132/) (2023)
- [Jailbreak and Guard Aligned Language Models with Only Few In-Context Demonstrations](https://mlsanity.com/paper/2310.06387/) (2023)
- [A Holistic Approach to Undesired Content Detection in the Real World](https://mlsanity.com/paper/2208.03274/) (2022)
- [Certifying LLM Safety against Adversarial Prompting](https://mlsanity.com/paper/2309.02705/) (2023)
- [SafeDecoding: Defending against Jailbreak Attacks via Safety-Aware Decoding](https://mlsanity.com/paper/2402.08983/) (2024)
- [WildTeaming at Scale: From In-the-Wild Jailbreaks to (Adversarially) Safer Language Models](https://mlsanity.com/paper/2406.18510/) (2024)
- [RAIN: Your Language Models Can Align Themselves without Finetuning](https://mlsanity.com/paper/2309.07124/) (2023)
- [Robust Prompt Optimization for Defending Language Models Against Jailbreaking Attacks](https://mlsanity.com/paper/2401.17263/) (2024)
- [On Prompt-Driven Safeguarding for Large Language Models](https://mlsanity.com/paper/2401.18018/) (2024)
- [GradSafe: Detecting Jailbreak Prompts for LLMs via Safety-Critical Gradient Analysis](https://mlsanity.com/paper/2402.13494/) (2024)
- [InferAligner: Inference-Time Alignment for Harmlessness through Cross-Model Guidance](https://mlsanity.com/paper/2401.11206/) (2024)
- [Defending Large Language Models against Jailbreak Attacks via Semantic Smoothing](https://mlsanity.com/paper/2402.16192/) (2024)
- [Gradient Cuff: Detecting Jailbreak Attacks on Large Language Models by Exploring Refusal Loss Landscapes](https://mlsanity.com/paper/2403.00867/) (2024)
- [Fight Back Against Jailbreaking via Prompt Adversarial Tuning](https://mlsanity.com/paper/2402.06255/) (2024)
- [JBShield: Defending Large Language Models from Jailbreak Attacks through Activated Concept Analysis and Manipulation](https://mlsanity.com/paper/2502.07557/) (2025)
- [AdaSteer: Your Aligned LLM is Inherently an Adaptive Jailbreak Defender](https://mlsanity.com/paper/2504.09466/) (2025)
- [Shaping the Safety Boundaries: Understanding and Defending Against Jailbreaks in Large Language Models](https://mlsanity.com/paper/2412.17034/) (2024)
- [Large Language Models Encode Semantics and Alignment in Linearly Separable Representations](https://mlsanity.com/paper/2507.09709/) (2025)
- [Smoothed Embeddings for Robust Language Models](https://mlsanity.com/paper/2501.16497/) (2025)

## Built on by (most cited first, 2 of 2)

- [Test-Time Training Undermines Safety Guardrails](https://mlsanity.com/paper/2605.22984/) (2026)
- [Self-Improving In-Context Learning](https://mlsanity.com/paper/2605.23180/) (2026)

## Abstract (arXiv)

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.