PTERA: Probabilistic Text Encoder for Reward Alignment in Diffusion Models

Fine-tuning text-to-image models toward human preference rewards has become a common approach for improving text alignment. Most existing methods, however, focus on updating the diffusion model parameters, while reward-based adaptation of the text encoder remains relatively underexplored despite promising reward gain and training efficiency. Moreover, existing text encoder fine-tuning methods typically optimize a single deterministic embedding for each prompt, which can be susceptible to reward over-optimization and provides no training mechanism for exploring alternative conditioning vectors. In this work, we propose **P**robabilistic **T**ext **E**ncoder for **R**eward **A**lignment (**PTERA**), a simple yet effective text encoder fine-tuning framework based on direct reward backpropagation. PTERA replaces deterministic conditioning with a trainable probability distribution over text embeddings and derives a tractable KL-regularized objective for learning this distribution while regularizing the divergence from a stochastic reference constructed around the pre-trained text embeddings. Compared with previous text encoder adaptation, PTERA exhibits more robust reward alignment. We further introduce **PTERA-TTA**, which exploits the text embedding distribution to perform trajectory-dependent test-time alignment.