SoftREPA: Aligning Text to Image in Diffusion Models is Easier Than You Think

A lightweight and flexible contrastive learning method to improve text-image alignment in diffusion models.

KAIST
*Indicates Equal Contribution
arXiv Code

T2I Generation Results

 (SD3)

Text Guided Image Editing Results

 (SD3)

Abstract

While recent advancements in generative modeling have significantly improved text-image alignment, some residual misalignment between text and image representations still remains. Although many approaches have attempted to address this issue by fine-tuning models using various reward models, etc., we revisit the challenge from the perspective of representation alignment—an approach that has gained popularity with the success of REPresentation Alignment (REPA). We first argue that conventional text-to-image (T2I) diffusion models, typically trained on paired image and text data (i.e., positive pairs) by minimizing score matching or flow matching losses, is suboptimal from the standpoint of representation alignment. Instead, a better alignment can be achieved through contrastive learning that leverages both positive and negative pairs. To achieve this efficiently even with pretrained models, we introduce a lightweight contrastive fine-tuning strategy called SoftREPA that uses soft text tokens. This approach improves alignment with minimal computational overhead by adding fewer than 1M trainable parameters to the pretrained model. Our theoretical analysis demonstrates that our method explicitly increases the mutual information between text and image representations, leading to enhanced semantic consistency. Experimental results across text-to-image generation and text-guided image editing tasks validate the effectiveness of our approach in improving the semantic consistency of T2I generative models.

Method

rs_algorithm

In SoftREPA, the soft tokens are optimized to contrastively match the score with positively conditioned predicted noise while repelling the score from negatively conditioned predicted noise. This process implicitly sharpens the joint probability distribution of images and text by reducing the log probability of negatively paired conditions. The learnable soft tokens are prepended to the text features across the upper layers, with each layer’s previous soft tokens being replaced by the soft tokens of the current layer.

Experimental Results

1. T2I Generation

t2i_gen_w_interpolation
ModelImageReward ↑CLIP ↑HPS ↑LPIPS ↓
SD1.50.1770.2650.2510.438
SD1.5 + Ours0.2740.2710.2520.437
SDXL0.7500.2670.2730.420
SDXL + Ours0.8520.2680.2830.423
SD30.9420.2630.2800.424
SD3 + Ours1.0850.2690.2890.428

We evaluate SoftREPA across diffusion and rectified flow models, and across UNet and Transformer architectures. Our approach consistently improves alignment and image quality on all fronts, demonstrating strong generalization.

2. Text Guided Image Editing

inv_algorithm
ModelDatasetImageReward ↑CLIP ↑HPS ↑LPIPS ↓
FlowEditDIV2K0.3800.2600.2550.154
FlowEdit + OursDIV2K0.4660.2630.2600.149
FlowEditCat2Dog0.9370.2250.2660.199
FlowEdit + OursCat2Dog1.1110.2260.2690.173

SoftREPA boosts alignment and visual quality in editing tasks with fewer steps and lower CFG scale.

Key Contributions

  • 📌 SoftREPA introduces soft tokens trained with contrastive learning.
  • ⚡ Requires < 1M trainable parameters, minimal compute cost.
  • 💡 Improves semantic representation alignment in generation and editing without model full finetuning.