SD3.5 Latent Reward Model
This repository contains a latent reward model trained on Pick-a-Pic pairwise preferences using Stable Diffusion 3.5 Medium features. The model scores prompt-image pairs and can be used to rank candidate images for a given text prompt.
This is not a text-to-image generator. It is a reward/preference model intended to compare or rank generated images.
Model Details
- Model type: Latent reward model for prompt-image preference scoring
- Base model:
stabilityai/stable-diffusion-3.5-medium
- Training dataset:
pickapic-anonymous/pickapic_v1
- Image resolution: 512
- Training objective: Pairwise preference loss over two candidate images per prompt
- Feature extraction: Multi-layer SD3.5 transformer features
- Text encoders: Trainable
- VAE: Frozen
Training Configuration
The selected full-epoch checkpoint was trained with:
| Parameter | Value |
|---|
| learning rate | 1e-5 |
| guidance scale | 3 |
| logit scale init value | 2.6592 |
| effective initial logit scale | exp(2.6592) ~= 14.285 |
| feature layers | multi |
| effective train batch size | 128 |
| gradient accumulation steps | 1 |
| mixed precision | bf16 |
| text encoders | trainable |
| SD3.5 transformer | trainable |
| VAE | frozen |
Evaluation
Evaluation was performed during training on the Pick-a-Pic validation splits. The primary model-selection metric was eval_unique/preference_accuracy, which avoids over-weighting repeated prompts.
Final full-epoch checkpoint metrics:
| Split / Metric | Value |
|---|
eval_unique/preference_accuracy | 0.6376 |
eval_unique/reward_margin | 0.2517 |
eval_unique/loss | 0.6294 |
eval/preference_accuracy | 0.6534 |
eval/reward_margin | 0.3109 |
eval/loss | 0.6410 |
The selected checkpoint performed best among the tested full-epoch configurations, including sweeps over guidance scale, logit scale, feature layer choice, learning rate, and frozen text encoders.
Intended Use
This model is intended for research and experimentation with image preference scoring, reranking, and reward-model-guided evaluation for text-to-image generations.
Example use cases:
- Ranking multiple generated images for the same prompt
- Comparing prompt-image alignment under a learned preference signal
- Studying latent reward models for diffusion model outputs
- Offline evaluation or reranking in text-to-image pipelines
Out-of-Scope Use
This model should not be used as:
- A general safety classifier
- A factuality or image authenticity detector
- A universal human preference oracle
- A replacement for human review in high-stakes settings
The model was trained on preference data and inherits the coverage, biases, and limitations of that data.
Limitations and Biases
The model reflects preferences present in Pick-a-Pic and may inherit dataset biases around aesthetics, subject matter, demographics, style, and prompt distribution. Performance may degrade on image domains or prompt types that are far from the training distribution.
The model scores relative preference/alignment, not factual correctness, safety, or social acceptability. Scores should be interpreted as a learned preference signal rather than an objective quality measure.
Usage
This checkpoint is expected to be used with the accompanying SD3.5 LRM code. A typical usage pattern is:
1from sd35_lrm.model import SD35LatentRewardModel
2
3model = SD35LatentRewardModel.from_pretrained("YOUR_HF_REPO_ID")
4model.eval()
5
6# Use model.score_images(captions, pixel_values) to score prompt-image pairs.
Replace YOUR_HF_REPO_ID with the Hugging Face repository id for this checkpoint.
