GRPO with a Representation-Level (RTV) Reward

The RL capstone: post-train Llama-3.2-3B with GRPO using a representation-level reward derived from RTV — rewarding generations whose internal activation trajectory looks safe, rather than an external content judge:

attack (dir-GCG)  →  detect (RTV)  →  defend (GRPO with an RTV reward)

Status: end-to-end validated on Llama-3.2-3B. The RTV internal-trajectory reward defends against held-out GCG attacks (100% → 6.9–10.3% ASR, within and cross-corpus) as well as a guard content classifier (parity, within the measured noise floor), at no helpfulness cost — and it also holds against a fully adaptive attacker that re-optimizes dir-GCG against the defended geometry (clean guard-ASR 32% → 9% on MaliciousInstruct-100), using only the model's own refusal-direction geometry. Full numbers in RESULTS.md.

GRPO in one paragraph

For each prompt, sample a GROUP of G completions from the current policy; score each; the advantage of completion i is its group-normalized reward A_i = (r_i − mean_group) / (std_group + eps). The group mean is the baseline, so GRPO needs no value/critic network (its main simplification over PPO). Update with a clipped ratio objective + KL to the reference (LoRA-adapter-disabled base = one copy of weights). TRL's GRPOTrainer handles rollouts, advantage, clipping, and KL.

Stage 1 — attack (dir-GCG)

attack/run_dirgcg.py optimizes a portable string suffix (token-invariant) so a mid-layer hidden state aligns with the StrongREJECT refusal direction, with a behavioral early-stop. attack/verify_guard.py then guard-confirms each suffix (p_harm ≥ 0.5 and not is_refusal, judged on the clean goal), and attack/build_jbgcg.py pools the survivors → 116 confirmed jailbreaks across 3 disjoint corpora (MI-32, HarmfulQA-55, JBB-29). JBB is held out entirely for cross-corpus eval.

Stage 2 — detect (RTV Design-A)

rtv/ reimplements Representation Trajectory Verification natively on a dedicated, exact-disjoint calibration pool. Per-layer refusal direction (diff-in-means) → fingerprint F[l,p] = cos(r_l, h_l^p) over layers × last-5 positions → two-class Ledoit-Wolf Mahalanobis M(x) = min(d_harmless, d_harmful), thresholded at 5% FPR. Depth-diverse operating set [14,16,25,26]: TPR 0.95 @ 5% FPR.

Stage 3 — defend (GRPO with the RTV reward)

The reward must score the completion (a prompt-side score is group-constant in GRPO → zero advantage → no gradient). reward/rtv_fingerprint.py fits a two-class LW-Mahalanobis model on completion refusal-direction fingerprints (layers [10,14,22,26,28], CV-selected by reward gap); reward/rtv_reward.py returns r_safe = tanh(LLR/2) = 2·P(safe|x) − 1 from a frozen reference model, in a frontier form (harmful → reward a safe trajectory, benign → reward a non-refusal, so RL can't win by always refusing). train/train_grpo.py runs TRL's GRPOTrainer with LoRA; --reward {scalar,rtv} swaps the reward for the head-to-head.

The result — RTV reward vs. guard-classifier reward

Identical config/data/eval, only the reward swapped. On held-out attacks (runs/eval_headtohead_final.json):

Parity, not superiority: RTV ≤ guard in every attack cell, but the within-set gap (10.3% vs 17.2% = 2 responses) is inside the ±1-response reproducibility floor — a representation-only reward matches a strong external classifier. No over-refusal collapse, plain-harmful refusal preserved. A train-only recalibration of the reward detector — excluding every eval-attack completion, and even re-selecting layers to [8,14,28] — reproduces this RTV column to the record (runs/eval_headtohead_airtight.json), so the defense does not depend on having calibrated on eval data.

Adaptive robustness (the primary robustness result). Beyond transferring the base attack, the defense holds against an attacker who re-optimizes dir-GCG from scratch against the defended geometry — directions rebuilt on the defended model, target layer re-selected on it (12→9), full base budget. On all 100 MaliciousInstruct goals, clean guard-confirmed ASR falls 32% → 9% (~70% relative). Base 32/100 is exactly the MI confirmed-suffix count from Stage 1 — the fresh attack yield on raw goals, not the pre-confirmed eval sets where base is 100% by construction. The reduction has two independent sources: the defended model refuses more (72% vs 44%) and its non-refusals are far less often harmful (36% vs 57%). See RESULTS.md §3b for the full breakdown, limits, and reproduction.

Full write-up — RESULTS.md

The complete numbers, methodology, and honest caveats live in RESULTS.md:

• §1–2 attack + detect — 116 guard-confirmed dir-GCG jailbreaks across 3 disjoint corpora; the RTV Design-A prompt detector (TPR 0.95 @ 5% FPR; 0.86 cross-corpus transfer).
• §3 head-to-head — RTV reward vs. guard reward, identical config: parity (RTV ≤ guard in every attack cell, within the ±1-response noise floor) at no helpfulness cost.
• §3a airtight recalibration — a train-only reward detector (layers [8,14,28]) reproduces §3 to within one response → the result is invariant to detector specification, not an artifact of calibrating on eval data.
• §3b adaptive dir-GCG — attack re-optimized against the defended geometry: clean ASR 32% → 9%; robustness traced to more refusals + less-harmful non-refusals, with limits stated explicitly (geometry-ablation ≠ reward-gaming; within-distribution).
• §4 methodology — why base ASR is exactly 100% by construction (template / decode / guard-input pitfalls, all fixed pipeline-wide).

Also present — reward-hacking probe (extension)

eval/reward_hacking.py buckets held-out completions by (actually-harmful? × looks-safe-inside?); if the RTV-trained policy has more mass in the HACKED cell (harmful yet looks-safe-inside) than the scalar-trained one, the representation reward was gamed. This is scaffolded as the natural next probe on top of the validated defense.

Metrics

• ASR on harmful + over-refusal on benign (the frontier), base → rtv → guard.
• Mean reward / benign-helpfulness during training (watch for refusal collapse).

Weights (pretrained adapters)

The LoRA adapters are not in git (large binaries) — download the three cited adapters from the v1.0 GitHub Release and unzip each into its runs/*/adapter/ path, or retrain from scratch with the commands below.

curl -L -o adapter_rtv.zip \
  https://github.com/kemalderya/grpo-rtv-reward/releases/download/v1.0/adapter_rtv.zip
unzip adapter_rtv.zip -d runs/grpo_defend_flash_rtv/adapter
# then: PeftModel.from_pretrained(base_model, "runs/grpo_defend_flash_rtv/adapter")

Run

pip install -r requirements.txt
# offline HF env used throughout:
#   HF_HOME=~/Desktop/ET/cache HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1

# Stage 3 — defend: train the two adapters (RTV reward vs guard reward), same config
python -m src.train.train_grpo --config configs/grpo_defend_flash.yaml \
  --data data/rollout_defend_flash.jsonl --reward rtv   --out runs/grpo_defend_flash_rtv
python -m src.train.train_grpo --config configs/grpo_defend_flash.yaml \
  --data data/rollout_defend_flash.jsonl --reward scalar --out runs/grpo_defend_flash_guard

# Evaluate base vs both adapters (matches the attack's decode: flash + bs=1 + native template)
python -m src.eval.generalize_check \
  --adapters rtv:runs/grpo_defend_flash_rtv/adapter guard:runs/grpo_defend_flash_guard/adapter \
  --batch-size 1 --out runs/eval_headtohead_final.json

The scripts/*.sh helpers (run_mvp.sh, build_fingerprint.sh, run_rtv.sh) are the earlier plain-prompt MVP entry points and predate the attack-suffixed defend pipeline above; prefer the module commands.

Layout

configs/    grpo_defend_flash.yaml (final head-to-head), grpo_defend*.yaml, grpo_smoke.yaml
src/attack/ run_dirgcg.py, verify_guard.py (guard-confirm), build_jbgcg.py (pool corpus),
            build_airtight_directions.py, ablate_direction_asr.py (layer pick),
            ablate_attack_layers.py, confirm_attack.py  (§3b adaptive attack on defended model)
src/rtv/    datasets.py (disjoint pool), extract_states.py, detector.py, evaluate.py, ablate.py
src/reward/ scalar_reward.py (guard frontier), rtv_fingerprint.py, rtv_reward.py (novel),
            build_completion_calib.py, completion_ablate.py
src/train/  train_grpo.py (TRL GRPO, --reward scalar|rtv)
src/eval/   evaluate.py (ASR/over-refusal), generalize_check.py (base-vs-adapters table),
            reward_hacking.py (extension probe)
src/        model_util.py (single flash loader), template_util.py (render_native)