anonym371/dose-response-c3

Dose-Response C3 (8M-10%): Unsafe oversampled to ~10%

Part of a dose-response experiment studying how unsafe training data fraction affects text-to-image model output safety.

Condition

Label	C3 (8M-10%)
Description	Unsafe oversampled to ~10%.
Training set size N	8.64M
Unsafe fraction p	9.6%
Unsafe count U	~829K

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	100,000
Samples seen	~25.60M
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

Training data sources

The training set combines three image datasets, with per-condition filtering/oversampling:

• lehduong/flux_generated (~1.7M)
• LucasFang/FLUX-Reason-6M (~6M)
• brivangl/midjourney-v6-llava (~1M)

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.