heli-stand/lrm-sd3.5-medium-pickapicv1

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:

from sd35_lrm.model import SD35LatentRewardModel

model = SD35LatentRewardModel.from_pretrained("YOUR_HF_REPO_ID")
model.eval()

# 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.