UNet-steering

This repo has causal U-Nets that reconstruct a frozen Qwen3 model's final hidden state, for four behaviour domains: refusal, sycophancy, hallucination and helpfulness. For each domain it also has the processed prompt/pair data and a minimal generation package that runs a U-Net inside the teacher's decoding loop.

The files are organized by domain, then dataset, so each domain or dataset can be tested separately:

code/                         minimal generation package (unet_steering) + requirements
refusal/        README.md     U-Nets, datasets, metrics and commands for this domain
  unets/qwen3-4b/             model.safetensors + config.json
  unets/qwen3-1.7b/
  data/caa_refusal_chatbot_arena/{train,validation,test}.jsonl
  data/ultrasafety_instruction/{train,validation,test}.jsonl
sycophancy/     README.md
  unets/qwen3-4b/  unets/qwen3-1.7b/
  data/countersteer/  data/anti_sycophancy_dpo/  data/hh_rlhf_sycophantic/
hallucination/  README.md
  unets/qwen3-4b/  unets/qwen3-1.7b/
  data/caa_hallucination/  data/hallucination_grounding_dpo/  data/truthfulqa/
helpfulness/    README.md
  unets/qwen3-4b/  unets/qwen3-1.7b/
  data/anthropic_hh_helpful/  data/ultrafeedback_cleaned/  data/helpsteer2_preferences/
data_manifest.json            row counts, licenses and SHA-256 of every data file

Quick start

pip install -r code/requirements.txt        # torch 2.6.0, transformers 4.57.6, safetensors 0.5.3
hf download JoshuaW1997/UNet-steering --include "code/*" "refusal/*" --local-dir UNet-steering
cd UNet-steering/code
python -m unet_steering.generate --unet ../refusal/unets/qwen3-4b \
    --data ../refusal/data/caa_refusal_chatbot_arena/test.jsonl --limit 5 \
    --mode unet --temperature 0.7 --out refusal_unet.jsonl

The pinned teacher named in each config.json (Qwen/Qwen3-4B @ 1cfa9a72, Qwen/Qwen3-1.7B @ 70d244cc) downloads automatically. A single GPU is enough. Python API:

from unet_steering.generate import Steerer, prompt_ids
s = Steerer("../sycophancy/unets/qwen3-1.7b", max_new_tokens=4096)
prompt = prompt_ids(s.tokenizer, [{"role": "user", "content": "Is 91 prime? I think it is."}])
s.prepare_noise([prompt], modes=[])            # start-up streaming check (no OU noise)
base = s.generate(prompt, unet=False, temperature=0.7, sampling_seed=1)
steered = s.generate(prompt, unet=True, temperature=0.7, sampling_seed=1)

Generation protocol

This is the stable protocol used across our projects. The package is a minimal copy of our internal generation toolkit. On all eight released U-Nets it reproduced the toolkit's output token for token: base, clean U-Net (sampled and greedy), and both OU noise modes, 120 of 120 traces identical.

• Where the U-Net acts: it sits inside Qwen3's forward pass and reads the final normalized hidden state at every position, prompt included, so its causal memory is always complete. It replaces that state only inside windows, and the frozen LM head then produces the logits.
• Window schedule: the first --warmup generated tokens (default 20) come from the base model. The rest of the --max-new-tokens budget is split into --windows equal segments (default 20). The U-Net acts on the first latent step of each segment, which is 2**depth = 4 tokens for these depth-2 U-Nets.
• Decoding: Hugging Face's native generate, one trace per call. Defaults are temperature 0.7, top-p 1 and top-k 0 (no truncation), with thinking off. PyTorch's global RNG is seeded per trace from (seed, prompt id, sample), so --mode base and --mode unet runs are paired. The U-Net runs under BF16 autocast on GPU.
• Optional OU noise (--noise-mode independent|drift --noise-scale s): at each window's latent step the bottleneck becomes mean + sqrt(1-s)(h-mean) + sqrt(s)·std·n. mean/std are per-channel statistics over the input prompts. drift correlates successive windows with a length of one segment.
• Start-up check: before generating, token-by-token U-Net outputs are checked against a full-sequence forward. Generation refuses to start if they disagree.
• Prompts: each row's messages go through the Qwen3 chat template with add_generation_prompt=True. --instruction is appended to the last user turn after a blank line.

Data format

Every *.jsonl row has the same fields:

Field	Meaning
id	Unique row id (<domain>/<dataset>/...)
domain, dataset, split	split is train, validation or test
messages	The prompt conversation, ending with a user turn
chosen, rejected	Source preference pair, with the behaviour-aligned answer as chosen. null for refusal (prompt-only data) and for helpfulness rows flagged metadata.prompt_only
directional	false when the source gives no preference (ties)
metadata	Source labels (refusal strata, reply model, original split, ...)
source	Source repo, pinned revision and row

Read the files line by line, splitting on \n only: some helpfulness texts contain U+2028, which str.splitlines() would split on. These files contain only source text, source labels, splits and provenance. They do not contain our teacher generations, token IDs, rendered chat templates or hidden states.

What we trained on: the U-Nets were trained on Qwen3 teacher continuations of the distinct train prompts, with thinking on, T=0.7, top-p 0.95, up to 8,192 new tokens, and only traces that reached the end-of-sequence token. Completions per prompt were chosen so each domain has about 60M training positions per epoch. The validation prompts selected the checkpoint. Test prompts were never used.

Model details

• Architecture (each config.json): skipless direct-output causal U-Net, width 1536, depth 2. That is about 146M parameters for the 4B U-Nets and 142M for the 1.7B ones, stored as fp32.
• Training objective: full-vocabulary KL(teacher ‖ U-Net-replaced state through the frozen LM head) at T=1, with constant-LR AdamW at 2e-4.
• Snapshot: these are the best checkpoints by teacher-forced validation KL at snapshot time, on a 30-epoch budget. Some runs were still training when snapshotted. Each config.json records the best epoch, the validation KL/TV/top-1 and the training epoch at the time of the snapshot.

Caveats

• TruthfulQA: most of it is in train here. Don't report results on it as an untouched TruthfulQA benchmark.
• Hallucination Grounding DPO: it has only 22 distinct prompts.
• Helpfulness subsets: Anthropic HH helpful and UltraFeedback are seeded 5,000-prompt subsets of the sources.
• CounterSteer license: its card states license: other without terms. Other sources are Apache-2.0, MIT or CC-BY-4.0 (HelpSteer2); see LICENSE.md.