anonym371/dose-response-c2-sft

Dose-Response C2 (8M-5%) SFT: Unsafe oversampled to 5%

SFT fine-tuned from the corresponding base on Alchemist (3,350 images) for 20K iterations.

Condition

Label	C2 (8M-5%)
Description	Unsafe oversampled to p=5%.
Training set size N	8.24M
Unsafe fraction p	5%
Unsafe count U	~412K

Base model

anonym371/dose-response-c2

SFT dataset

Alchemist (3,350 images)

Architecture

Class	PRX (rectified-flow DiT)
Hidden size	1792
Depth	16
Heads	28
MLP ratio	3.5
Patch size	32 px
Bottleneck	256
Resolution	512×512

Text encoder

Model	google/t5gemma-2b-2b-ul2
Max prompt tokens	256
Dtype	bfloat16

Diffusion scheduler

Type	x-prediction flow matching
Train timesteps	1000
Timestep shift	3.0

Training

Iterations	20,000
Samples seen	~5.12M
Global batch size	256
Microbatch (per GPU)	32
Hardware	8× NVIDIA H200
Precision	bfloat16 (amp_bf16)
Optimizer (transformer blocks)	Muon (lr=1e-4, momentum=0.95, nesterov, ns_steps=5, weight_decay=0)
Optimizer (other params)	AdamW (lr=1e-4, β=(0.9, 0.95), eps=1e-8, weight_decay=0)
LR schedule	1,000-step linear warmup, constant after
EMA	decay 0.999, started at step 0
Random seed	42
Trainer	Composer + FSDP

Files

• denoiser.pt — Consolidated EMA-denoiser checkpoint
• config.yaml — Full training configuration

Framework

Trained with the PRX framework (Composer + FSDP). The full config.yaml is included for reproducibility.