Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation

TL;DR

The paper addresses the problem of aligning text-to-image diffusion models to targeted rewards without expensive fine-tuning or vulnerable reward hacking. It introduces Null-Text Test-Time Alignment (Null-TTA), which optimises the unconditional (null-text) embedding in classifier-free guidance to steer the model’s generative distribution within the semantic space defined by the text encoder, coupled with a KL-based regularisation to preserve pretrained behavior. A lightweight backward-process refinement via particle filtering further guides the denoising trajectory toward high-reward regions. Empirically, Null-TTA achieves state-of-the-art target reward performance while maintaining strong cross-reward generalisation and favorable compute/memory efficiency, across both differentiable and non-differentiable reward settings. This semantic-space optimisation framework establishes a principled, training-free paradigm for reliable inference-time alignment of diffusion models with broad practical impact.

Abstract

Test-time alignment (TTA) aims to adapt models to specific rewards during inference. However, existing methods tend to either under-optimise or over-optimise (reward hack) the target reward function. We propose Null-Text Test-Time Alignment (Null-TTA), which aligns diffusion models by optimising the unconditional embedding in classifier-free guidance, rather than manipulating latent or noise variables. Due to the structured semantic nature of the text embedding space, this ensures alignment occurs on a semantically coherent manifold and prevents reward hacking (exploiting non-semantic noise patterns to improve the reward). Since the unconditional embedding in classifier-free guidance serves as the anchor for the model's generative distribution, Null-TTA directly steers model's generative distribution towards the target reward rather than just adjusting the samples, even without updating model parameters. Thanks to these desirable properties, we show that Null-TTA achieves state-of-the-art target test-time alignment while maintaining strong cross-reward generalisation. This establishes semantic-space optimisation as an effective and principled novel paradigm for TTA.

Figures (3)

• Figure 1: Evaluation of reward-optimisation (x-axes) vs over-optimisation (generalisation to held-out rewards on y-axes). Top/Bottom rows correspond to aligning with Aesthetic and HPSv2 target rewards respectively. Each $\filledstar$ point corresponds to a particular baseline method, which together define the state-of-the-art pareto front (dashed line). For Null-TTA, each $\circ$ point indicates the optimisation intensity (maximum inner steps $n_\text{max}$). Null-TTA consistently shifts the Pareto front outward, indicating superior target reward optimisation while maintaining strong generalisation across unseen rewards. All the experiments in this figure were done over 3 different random seeds. Minimum inner steps $n_\text{min}$ is set to 5.
• Figure 2: Qualitative comparison on six challenging categories—counting (“Nine marbles arranged in a perfect square”), compositionality (“A cat riding on a dog's back while holding a tiny flag”), spatial reasoning (“A rabbit standing directly under a hanging lamp”), unusual colors (“A bear with galaxy-patterned fur”), fine-grained style transfer (“A wolf drawn in vibrant cyberpunk neon edge highlights”), and impossible scenes (“A piano drifting in deep space surrounded by comets”). Null-TTA consistently produces images that more faithfully satisfy the prompt constraints while preserving global coherence.
• Figure 3: Multi-objective optimisation using $R_{\text{multi}} = w \cdot \text{PickScore} + (1-w)\cdot \text{HPSv2}$. Each marker corresponds to a different weight $w$. Null-TTA consistently achieves a superior trade-off curve: when increasing PickScore, its HPSv2 degradation is smaller than DAS, and when increasing HPSv2, its PickScore degradation is also milder. This bidirectional advantage demonstrates that Null-TTA provides more efficient and better-balanced optimisation across the two reward dimensions.