Reward shaping: trading off safety and compliance
Reward shaping is a tunable safety dose for DPO — a small α-sweep maps how much over-refusal each unit of added safety costs
| Topics | Language models · Post-training & RL · Safety & guardrails |
| Models | Nemotron-3-Super-120B-A12B · OLMo-3.1-32B refusal signal · WildGuard judge |
| References | DPO (Rafailov et al.) · refusal-projection reward shaping (alignment-megatron) |
Tuning safety
Reward shaping nudges a preference-trained model toward refusing harmful requests by adding a small, signed term to the training objective. Instead of changing the data or the model, it adds a per-pair offset inside the DPO loss that rewards the model for behaving more like a known refusal signal. A single hyperparameter, α, sets how hard it pushes.
In this recipe, Silico runs reward shaping on Nemotron-3-Super-120B at three strengths — α=−1, −2, and −5. The α=−2 model already existed as one trained point, so Silico filled in α=−1 and α=−5 as controlled single-variable jobs — same SFT init, same data, same DPO recipe, only α changed — turning one point into a four-point dose-response curve against the α=0 baseline. Because α is fixed per training run, choosing it well means seeing the whole curve rather than one point: train a few α and read the safety-versus-compliance trade-off they trace. Across the sweep, the harmful-response rate falls 16.1% → 5.2% from α=0 to −5 while benign compliance falls from 0.916 to 0.828 — the safety gain is real and monotonic, and so is its cost.
Key takeaways
- Silico runs reward shaping end to end. A custom post-training technique — a per-pair offset inside the DPO loss, two sequential 16-node training jobs, export, and safety eval — ran as a single launched experiment, including auto-resuming the α=−5 job after it hit a full disk near the end.
- Marginal returns and marginal costs both matter. Here each extra unit of α buys less safety and keeps spending the same compliance, so the strongest setting is rarely the right one.
- A controlled single-variable sweep makes the curve trustworthy. Holding the data, init, and recipe fixed across arms is what lets the differences be read as the effect of α alone.
Reward shaping in a nutshell
Direct Preference Optimization (DPO) trains a model from pairs of responses where one is preferred over the other, pushing the model to make preferred responses more likely than rejected ones. It needs no separate reward model. The objective is a logistic loss over the gap between the model's score for the chosen response and the rejected one.
Reward shaping leaves that machinery intact and inserts one extra term. For each preference pair, it shifts the loss by an offset that grows with how much more the chosen response looks like a refusal than the rejected one:
oi = α · ( refproj(chosen) − refproj(rejected) ) / std
Here refproj is the refusal signal — a direction read from a separate model's activations that fires on refusals, here a precomputed cache from OLMo-3.1-32B-Instruct (layer 21) — and α scales the whole term: at α=0 the offset vanishes and training is ordinary DPO, and as |α| grows the model is pushed harder to treat refusal-like behavior as preferred. Everything else is shared across arms — the same SFT starting policy, the same preference data (Dolci-Instruct-DPO, 258,488 pairs), and a single launcher that switches vanilla DPO and reward shaping with one environment flag.
The honest failure mode worth naming: if the refusal cache is missing, the configured fallback silently runs vanilla DPO. So part of running reward shaping is confirming the offset actually reached the loss, not assuming it did.
Mapping the α dose-response
Every arm was scored on the same 11-task safety profile, with overall safety the mean of five adversarial benchmarks and benign compliance reported separately. The four α points trace a clean dose-response: across α = 0, −1, −2, −5, overall safety rises 0.856 → 0.940 while benign compliance falls 0.916 → 0.828, so stronger shaping is reliably safer and reliably less compliant. The exchange rate worsens with strength, though. From α=−2 to −5 the run spends −5.6pp compliance for just +2.4pp more safety, so α=−1 and α=−2 sit on the favorable part of the curve and α=−5 over-refuses for little extra gain. The gains concentrate on the hardest adversarial jailbreaks (DoAnythingNow, HarmBench human-jailbreaks) and the cost localizes almost entirely to benign-prompt handling, consistent with the refusal direction generalizing to real attacks first and catching benign requests only at high α.
Two limits bound the reading. Each arm is a single run, so differences within evaluation noise are not separated from seed variation, and the α=−2 to −5 step spans three unsampled units, so the exact bend in the curve is bracketed rather than pinned. Every arm also shares one refusal signal, so the curve describes scaling that one signal, not the choice of signal.
Reward shaping versus filtering the data
Reward shaping is not the only single-variable way to make a safer model from the same starting point. The sweep was scored alongside two arms that use no shaping at all, so the comparison runs on identical footing — same SFT init, same Dolci-Instruct-DPO pairs, same recipe — with only one thing changed in each:
- Plain DPO (the baseline). Preference tuning with no safety intervention. It sets the reference point: 0.856 safety, 0.916 benign compliance, 16.1% harmful-response rate.
- Data filtering — change the data. A prefilter hard-drops the top-5% most "comply-over-refuse" pairs before training, so the model never sees the worst examples. It improves both axes over plain DPO at once (0.873 safety, 0.924 compliance, 13.6% harmful) — the only arm in the set that does. The cost is that it has a ceiling: once the worst pairs are gone, there is nothing left to drop, so safety only climbs so far.
- Reward shaping — change the loss. Instead of removing pairs, it softly biases every surviving pair through the per-pair offset. That makes it continuous and not bounded by the data: at α=−5 it reaches 0.940 safety and 5.2% harmful, well past where filtering tops out. The cost is the trade-off this recipe maps — it spends benign compliance to get there.
The two are complements, not rivals. Filtering is the cheaper first move when a clean filter exists, because it buys safety without spending compliance, but it cannot reach a high safety target on its own. Reward shaping is what pushes further, and it composes with filtering — shape the loss on top of already-filtered data when the target sits beyond the filter's ceiling.