PRewrite¶

PRewrite (Kong et al., 2024) is an input control method that rewrites a seed instruction into a more effective form using an LLM as a rewriter. Three strategies from the original paper are supported:

• PRewrite-I (strategy="inference"): generate a single greedy rewrite. This is relatively cheap; useful when a quick instruction rewrite is enough.
• PRewrite-S (strategy="search"): sample k_candidates rewrites and select the best by evaluating each on a held-out dev set. This is more expensive but is generally of higher quality.
• RL-trained rewriter (by setting train_rewriter=True): train the rewriter with GRPO (group-relative policy optimization, via the TRL wrapper layer at structural_control/wrappers/trl/) before proposing rewrites; the trained rewriter is then consumed by either PRewrite-I or PRewrite-S. The reward is the downstream task score (apply each rewrite with the frozen task model over a dev set and score it; this is the paper's score-in-the-loop reward), or a user-supplied reward_fn. The toolkit's implementation of PRewrite uses GRPO rather than the paper's PPO since GRPO takes a callable reward (so the reward can be the task score directly).

PRewrite's chosen rewrite is applied as the system prompt at inference time.

Reference: Kong et al., 2024; PRewrite: Prompt Rewriting with Reinforcement Learning.

Method parameters¶

parameter	type	description
initial_instruction	str	Seed instruction to rewrite. Required.
rewriter_model_name_or_path	str \| None	HF id or local path for the rewriter LLM. None reuses the task model as rewriter (forbidden under training).
rewriter_model	PreTrainedModel \| None	Pre-loaded rewriter instance; mutually exclusive with rewriter_model_name_or_path.
rewriter_tokenizer	PreTrainedTokenizer \| None	Tokenizer paired with rewriter_model (required if the model has no inferable name_or_path).
meta_prompt	str \| None	Custom rewriter meta-prompt template (must contain {seed}). None uses the toolkit default.
strategy	"inference" \| "search"	PRewrite-I (greedy single rewrite) or PRewrite-S (best-of-k by dev-set eval).
k_candidates	int	Number of candidate rewrites to sample under the search strategy.
dev_set	list[dict] \| None	Held-out rows used to score candidates (search strategy) and as the GRPO scorer-in-the-loop reward (training). Each row needs "input" and optionally "reference".
row_scorer	Callable \| None	Per-row scorer (response, row) -> float aggregating dev-set responses into a scalar (used by the search strategy and the GRPO reward).
train_rewriter	bool	If True, GRPO-train the rewriter before proposing rewrites. Requires an explicit rewriter and a reward source (reward_fn, or row_scorer + dev_set).
reward_fn	Callable \| None	Custom GRPO reward reward_func(prompts, completions, **kwargs) -> list[float]. Takes precedence over row_scorer + dev_set; if unset, the reward is built from row_scorer + dev_set (the paper's reward).
training_seeds	list[str] \| None	Pool of seed instructions used as GRPO rollout prompts. Defaults to [initial_instruction].
grpo_config	dict \| None	Configuration forwarded to TRL's GRPO trainer (see GRPOArgs for keys: num_generations, beta, max_completion_length, learning_rate, per_device_train_batch_size, ...).
reward_dev_size	int \| None	Cap on dev rows used per reward evaluation during GRPO training (cost control); a deterministic head slice of dev_set.
rewriter_gen_kwargs	dict \| None	Generation kwargs for the rewriter LLM (e.g. temperature, max_new_tokens).
eval_gen_kwargs	dict \| None	Generation kwargs used when scoring candidates against the dev set (search selection and the GRPO reward).

Setup¶

If running this from a Google Colab notebook, please uncomment the following cell to install the toolkit. The following block is not necessary if running this notebook from a virtual environment where the toolkit has already been installed. This notebook uses the inspect extra for the f1() and exact() scorers used to score rewrites (both dev and all include it).

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability
# !pip install -q -e ".[eval]"

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability # !pip install -q -e ".[eval]"

The following authentication steps may be necessary to access any gated models (after being granted access by Hugging Face). Uncomment the following if you need to log in to the Hugging Face Hub using a HUGGINGFACE_TOKEN value stored in a .env file.

In [2]:

# !pip install -q python-dotenv
# from dotenv import load_dotenv
# import os

# load_dotenv()
# token = os.getenv("HUGGINGFACE_TOKEN")
# from huggingface_hub import login
# login(token=token)

# !pip install -q python-dotenv # from dotenv import load_dotenv # import os # load_dotenv() # token = os.getenv("HUGGINGFACE_TOKEN") # from huggingface_hub import login # login(token=token)

Example: Rewriting a vague instruction for short-answer extraction¶

We start with the seed instruction "Please provide an answer to the question." and ask PRewrite to produce a more specific instruction that nudges the model toward concise factual answers. We compare the greedy PRewrite-I rewrite, PRewrite-S which evaluates k_candidates rewrites on a small dev set of factual Q/A pairs and keeps the one with the highest F1, and the GRPO-trained rewriter.

In [3]:

import gc
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from inspect_ai.scorer import exact, f1

from steerability.algorithms.input_control.prewrite import PRewrite
from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.evaluation.scorers import sample_scorer_from_inspect

MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"

SEED_INSTRUCTION = "Please provide an answer to the question."

f1_scorer = sample_scorer_from_inspect(f1())     # SQuAD token-F1
em_scorer = sample_scorer_from_inspect(exact())  # SQuAD exact match

import gc import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import torch from transformers import AutoModelForCausalLM, AutoTokenizer from inspect_ai.scorer import exact, f1 from steerability.algorithms.input_control.prewrite import PRewrite from steerability.algorithms.core.steering_pipeline import SteeringPipeline from steerability.evaluation.scorers import sample_scorer_from_inspect MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct" SEED_INSTRUCTION = "Please provide an answer to the question." f1_scorer = sample_scorer_from_inspect(f1()) # SQuAD token-F1 em_scorer = sample_scorer_from_inspect(exact()) # SQuAD exact match

Define a small dev set of short-answer factual Q/A pairs and a held-out test prompt:

In [4]:

dev_set = [
    {"input": "What's the capital of France?", "reference": "Paris"},
    {"input": "What's the capital of Australia?", "reference": "Canberra"},
    {"input": "How many planets are in the Solar System?", "reference": "8"},
    {"input": "How many sides does a hexagon have?", "reference": "6"},
    {"input": "What's the chemical symbol for sodium?", "reference": "Na"},
    {"input": "What's the chemical symbol for potassium?", "reference": "K"},
    {"input": "Who painted the Mona Lisa?", "reference": "Leonardo da Vinci"},
    {"input": "Who wrote 'Pride and Prejudice'?", "reference": "Jane Austen"},
    {"input": "What year did World War II end?", "reference": "1945"},
    {"input": "What year did the Berlin Wall fall?", "reference": "1989"},
    {"input": "What's the largest ocean?", "reference": "Pacific"},
    {"input": "What's the smallest planet?", "reference": "Mercury"},
    {"input": "Which planet is known as the Red Planet?", "reference": "Mars"},
    {"input": "What's the hardest natural substance?", "reference": "diamond"},
    {"input": "What's the freezing point of water in Celsius?", "reference": "0"},
    {"input": "What language has the most native speakers?", "reference": "Mandarin"},
    {"input": "What's the currency of Japan?", "reference": "yen"},
    {"input": "How many continents are there?", "reference": "7"},
    {"input": "What's the speed of light in km/s (approx)?", "reference": "300000"},
    {"input": "Who discovered penicillin?", "reference": "Alexander Fleming"},
    {"input": "What's the longest river in the world?", "reference": "Nile"},
    {"input": "What's the chemical formula for water?", "reference": "H2O"},
]

PROMPT = "Who wrote the novel '1984'?"  # held-out; not in dev_set

dev_set = [ {"input": "What's the capital of France?", "reference": "Paris"}, {"input": "What's the capital of Australia?", "reference": "Canberra"}, {"input": "How many planets are in the Solar System?", "reference": "8"}, {"input": "How many sides does a hexagon have?", "reference": "6"}, {"input": "What's the chemical symbol for sodium?", "reference": "Na"}, {"input": "What's the chemical symbol for potassium?", "reference": "K"}, {"input": "Who painted the Mona Lisa?", "reference": "Leonardo da Vinci"}, {"input": "Who wrote 'Pride and Prejudice'?", "reference": "Jane Austen"}, {"input": "What year did World War II end?", "reference": "1945"}, {"input": "What year did the Berlin Wall fall?", "reference": "1989"}, {"input": "What's the largest ocean?", "reference": "Pacific"}, {"input": "What's the smallest planet?", "reference": "Mercury"}, {"input": "Which planet is known as the Red Planet?", "reference": "Mars"}, {"input": "What's the hardest natural substance?", "reference": "diamond"}, {"input": "What's the freezing point of water in Celsius?", "reference": "0"}, {"input": "What language has the most native speakers?", "reference": "Mandarin"}, {"input": "What's the currency of Japan?", "reference": "yen"}, {"input": "How many continents are there?", "reference": "7"}, {"input": "What's the speed of light in km/s (approx)?", "reference": "300000"}, {"input": "Who discovered penicillin?", "reference": "Alexander Fleming"}, {"input": "What's the longest river in the world?", "reference": "Nile"}, {"input": "What's the chemical formula for water?", "reference": "H2O"}, ] PROMPT = "Who wrote the novel '1984'?" # held-out; not in dev_set

Evaluation scorer¶

We score short-answer responses with Inspect AI's f1() scorer, the standard SQuAD token-level F1 (Rajpurkar et al., 2016), adapted into the toolkit's per-row SampleScorer shape with sample_scorer_from_inspect. We select rewrites on F1 since its precision term penalizes verbose answers that merely contain the gold span, so it rewards concise, correct answers without saturating. The exact() scorer, SQuAD exact match, is shown alongside for contrast.

As a quick sanity check, we run the scorers over a small set of illustrative responses.

In [5]:

for _resp, _ref in [
    ("Paris", "Paris"),  # concise + correct
    ("The capital of France is Paris.", "Paris"),  # correct but verbose -> lower F1
    ("London", "Paris"),  # wrong -> 0
]:
    _row = {"input": "What's the capital of France?", "reference": _ref}
    print(f"EM={em_scorer(_resp, _row):.0f}  F1={f1_scorer(_resp, _row):.2f}  ref={_ref!r:9}  resp={_resp!r}")

for _resp, _ref in [ ("Paris", "Paris"), # concise + correct ("The capital of France is Paris.", "Paris"), # correct but verbose -> lower F1 ("London", "Paris"), # wrong -> 0 ]: _row = {"input": "What's the capital of France?", "reference": _ref} print(f"EM={em_scorer(_resp, _row):.0f} F1={f1_scorer(_resp, _row):.2f} ref={_ref!r:9} resp={_resp!r}")

EM=1  F1=1.00  ref='Paris'    resp='Paris'
EM=0  F1=0.33  ref='Paris'    resp='The capital of France is Paris.'
EM=0  F1=0.00  ref='Paris'    resp='London'

Baseline model behavior¶

We start by generating a response with the bare seed instruction (no rewriting), to provide a reference for the comparisons below.

In [6]:

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, device_map="auto", dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

baseline_chat = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": SEED_INSTRUCTION},
        {"role": "user", "content": PROMPT},
    ],
    tokenize=False,
    add_generation_prompt=True,
)
baseline_inputs = tokenizer(baseline_chat, return_tensors="pt").to(model.device)
baseline_output = model.generate(**baseline_inputs, max_new_tokens=64, do_sample=False)
baseline_response = tokenizer.decode(
    baseline_output[0, baseline_inputs.input_ids.size(1):],
    skip_special_tokens=True,
)
print("Baseline response:\n")
print(baseline_response)

del model, baseline_inputs, baseline_output
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, device_map="auto", dtype=torch.bfloat16 ) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) baseline_chat = tokenizer.apply_chat_template( [ {"role": "system", "content": SEED_INSTRUCTION}, {"role": "user", "content": PROMPT}, ], tokenize=False, add_generation_prompt=True, ) baseline_inputs = tokenizer(baseline_chat, return_tensors="pt").to(model.device) baseline_output = model.generate(**baseline_inputs, max_new_tokens=64, do_sample=False) baseline_response = tokenizer.decode( baseline_output[0, baseline_inputs.input_ids.size(1):], skip_special_tokens=True, ) print("Baseline response:\n") print(baseline_response) del model, baseline_inputs, baseline_output gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache()

[transformers] Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.

[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer TokenizersBackend. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.

Baseline response:

The novel '1984' was written by George Orwell.

Aligning the rewriter with the scorer¶

By default, PRewrite's rewriter optimizes for general clarity and thoroughness, which tends to make answers longer. This demo aims to show how to rewrite instructions to induce more concise answers and, since the scorer only selects among the candidates the rewriter proposes, we define a custom meta_prompt that asks for instructions eliciting short, exact answers. Note that the template must contain {seed}.

In [7]:

CONCISE_META_PROMPT = (
    "You are a prompt engineer. Rewrite the instruction below into an instruction telling the model "
    "to reply with ONLY the answer, e.g., the number, name, place, etc.\n\n"
    "Original instruction:\n{seed}\n\n"
    "Rewritten instruction:"
)

CONCISE_META_PROMPT = ( "You are a prompt engineer. Rewrite the instruction below into an instruction telling the model " "to reply with ONLY the answer, e.g., the number, name, place, etc.\n\n" "Original instruction:\n{seed}\n\n" "Rewritten instruction:" )

PRewrite-I (greedy single rewrite)¶

The simplest mode is to ask the rewriter LLM (in this case we use the same model for both the task and rewriter) to produce one improved instruction (under greedy-decoding). This is cheap to run and does not use the dev set.

In [8]:

prewrite_i = PRewrite(
    initial_instruction=SEED_INSTRUCTION,
    meta_prompt=CONCISE_META_PROMPT,
    strategy="inference",
    rewriter_gen_kwargs={"max_new_tokens": 64, "do_sample": False},
)
pipeline_i = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[prewrite_i],
    device_map="auto",
    hf_model_kwargs={"dtype": torch.bfloat16},
)
pipeline_i.steer()

print("Chosen rewrite (PRewrite-I):\n")
print(prewrite_i.memory["instruction"])

prewrite_i = PRewrite( initial_instruction=SEED_INSTRUCTION, meta_prompt=CONCISE_META_PROMPT, strategy="inference", rewriter_gen_kwargs={"max_new_tokens": 64, "do_sample": False}, ) pipeline_i = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[prewrite_i], device_map="auto", hf_model_kwargs={"dtype": torch.bfloat16}, ) pipeline_i.steer() print("Chosen rewrite (PRewrite-I):\n") print(prewrite_i.memory["instruction"])

Chosen rewrite (PRewrite-I):

Provide the answer.

In [9]:

response_i = pipeline_i.generate(
    messages=[{"role": "user", "content": PROMPT}],
    max_new_tokens=64,
    do_sample=False,
)
print("Response (PRewrite-I):\n")
print(response_i)

del pipeline_i
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

response_i = pipeline_i.generate( messages=[{"role": "user", "content": PROMPT}], max_new_tokens=64, do_sample=False, ) print("Response (PRewrite-I):\n") print(response_i) del pipeline_i gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache()

Response (PRewrite-I):

The novel '1984' was written by George Orwell.

PRewrite-S (best-of-k by dev-set evaluation)¶

The PRewrite-S strategy samples k_candidates rewrites with do_sample=True, then runs each on the dev set under our f1() scorer and keeps the rewrite with the highest mean score.

We use a small k_candidates=3 for illustration purposes. The dev rows for each candidate are evaluated in a single batched generate call.

In [10]:

prewrite_s = PRewrite(
    initial_instruction=SEED_INSTRUCTION,
    meta_prompt=CONCISE_META_PROMPT,
    strategy="search",
    k_candidates=3,
    dev_set=dev_set,
    row_scorer=f1_scorer,
    rewriter_gen_kwargs={"max_new_tokens": 160, "do_sample": True, "temperature": 0.7},
    eval_gen_kwargs={"max_new_tokens": 16, "do_sample": False},
)
pipeline_s = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[prewrite_s],
    device_map="auto",
    hf_model_kwargs={"dtype": torch.bfloat16},
)
pipeline_s.steer()

print("Chosen rewrite (PRewrite-S):\n")
print(prewrite_s.memory["instruction"])

prewrite_s = PRewrite( initial_instruction=SEED_INSTRUCTION, meta_prompt=CONCISE_META_PROMPT, strategy="search", k_candidates=3, dev_set=dev_set, row_scorer=f1_scorer, rewriter_gen_kwargs={"max_new_tokens": 160, "do_sample": True, "temperature": 0.7}, eval_gen_kwargs={"max_new_tokens": 16, "do_sample": False}, ) pipeline_s = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[prewrite_s], device_map="auto", hf_model_kwargs={"dtype": torch.bfloat16}, ) pipeline_s.steer() print("Chosen rewrite (PRewrite-S):\n") print(prewrite_s.memory["instruction"])

Chosen rewrite (PRewrite-S):

Provide the answer only.

In [11]:

response_s = pipeline_s.generate(
    messages=[{"role": "user", "content": PROMPT}],
    max_new_tokens=64,
    do_sample=False,
)
print("Response (PRewrite-S):\n")
print(response_s)

del pipeline_s
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

response_s = pipeline_s.generate( messages=[{"role": "user", "content": PROMPT}], max_new_tokens=64, do_sample=False, ) print("Response (PRewrite-S):\n") print(response_s) del pipeline_s gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache()

Response (PRewrite-S):

George Orwell.

Comparing the rewrites¶

The instructions under the seed vs. PRewrite-I's greedy rewrite vs. PRewrite-S's selected best-of-k are as follows:

In [12]:

print(f"Seed:           {SEED_INSTRUCTION}")
print("PRewrite-I:    ", prewrite_i.memory["instruction"])
print("PRewrite-S:    ", prewrite_s.memory["instruction"])

print(f"Seed: {SEED_INSTRUCTION}") print("PRewrite-I: ", prewrite_i.memory["instruction"]) print("PRewrite-S: ", prewrite_s.memory["instruction"])

Seed:           Please provide an answer to the question.
PRewrite-I:     Provide the answer.
PRewrite-S:     Provide the answer only.

GRPO-trained rewriter (train_rewriter=True)¶

PRewrite can optimize the rewriter directly against task performance by setting train_rewriter=True, which trains it with GRPO (group-relative policy optimization) via the TRL wrapper. We deviate from the original paper (which uses PPO) primarily since GRPO is critic-free (there is no separate reward model and no value model), just a callable reward. Here the reward is the paper's scorer-in-the-loop reward, which PRewrite builds automatically from the row_scorer + dev_set we already used for PRewrite-S. Each candidate rewrite is applied with the frozen task model over the dev set and scored with f1(). In other words, the same signal PRewrite-S uses to select rewrites now serves as the training reward. (To use a different reward, pass a callable reward_fn(prompts, completions, **kwargs) -> list[float]; it takes precedence over row_scorer + dev_set.)

PRewrite forbids the default task-LM-as-rewriter behavior under training (training would mutate the task model in place, breaking downstream uses). Pass an explicit rewriter_model_name_or_path (or rewriter_model=... for a pre-loaded instance). The example below uses a smaller, explicit rewriter and keeps the task model (MODEL_NAME) as the frozen scorer for the reward, so the rewriter is optimized for the model the instruction is actually deployed on.

GRPO is configured via grpo_config (forwarded to TRL's GRPOArgs): num_generations is the group size G used for the group-relative advantage (must be >= 2 and evenly divide per_device_train_batch_size), beta is the KL-to-reference coefficient, and max_completion_length bounds the rewrite length. The scorer reward is the expensive part (a full dev-set pass with the task model for every distinct rewrite in a group, every step), so we cap it with reward_dev_size and keep the dev set, group size, and epoch count small for this illustration.

In [13]:

REWRITER_NAME = "meta-llama/Llama-3.2-3B-Instruct"  # smaller than the task LM; the trainable rewriter

training_seeds = [
    "Answer the question.",
    "Reply concisely.",
    "Provide the answer.",
    "Respond with just the answer.",
    "Give a short, direct answer.",
    "Answer briefly and precisely.",
    "State the answer only.",
    "Return the exact answer.",
]

prewrite_grpo = PRewrite(
    initial_instruction=SEED_INSTRUCTION,
    meta_prompt=CONCISE_META_PROMPT,
    strategy="inference",
    rewriter_model_name_or_path=REWRITER_NAME,  # explicit rewriter; the task LM is never trained
    train_rewriter=True,
    row_scorer=f1_scorer,  # scorer-in-the-loop reward
    dev_set=dev_set,
    reward_dev_size=8,  # cap dev rows per reward eval (cost control); the reward runs the 8B task LM
    training_seeds=training_seeds,
    grpo_config={
        "num_train_epochs": 2,
        "num_generations": 4,  # GRPO group size G (>= 2)
        "per_device_train_batch_size": 4,  # must be divisible by num_generations
        "learning_rate": 1e-5,  # LoRA tolerates a higher LR than full-FT GRPO (TRL default 1e-6)
        "beta": 0.04,  # KL-to-reference coefficient
        "max_completion_length": 48,  # rewrites are short
        "use_peft": True,  # LoRA-by-default
    },
    rewriter_gen_kwargs={"max_new_tokens": 64, "do_sample": True, "temperature": 0.9},
    eval_gen_kwargs={"max_new_tokens": 16, "do_sample": False},  # used by the reward's task-LM passes
)

pipeline_grpo = SteeringPipeline(
    model_name_or_path=MODEL_NAME,  # frozen 8B task model: scores rewrites for the reward, then answers at inference
    controls=[prewrite_grpo],
    device_map="auto",
    hf_model_kwargs={"dtype": torch.bfloat16},
)
pipeline_grpo.steer()  # GRPO trains the rewriter against the task scorer, then proposes a single rewrite

print("Chosen rewrite (GRPO-trained rewriter, PRewrite-I):\n")
print(prewrite_grpo.memory["instruction"])

del pipeline_grpo
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

REWRITER_NAME = "meta-llama/Llama-3.2-3B-Instruct" # smaller than the task LM; the trainable rewriter training_seeds = [ "Answer the question.", "Reply concisely.", "Provide the answer.", "Respond with just the answer.", "Give a short, direct answer.", "Answer briefly and precisely.", "State the answer only.", "Return the exact answer.", ] prewrite_grpo = PRewrite( initial_instruction=SEED_INSTRUCTION, meta_prompt=CONCISE_META_PROMPT, strategy="inference", rewriter_model_name_or_path=REWRITER_NAME, # explicit rewriter; the task LM is never trained train_rewriter=True, row_scorer=f1_scorer, # scorer-in-the-loop reward dev_set=dev_set, reward_dev_size=8, # cap dev rows per reward eval (cost control); the reward runs the 8B task LM training_seeds=training_seeds, grpo_config={ "num_train_epochs": 2, "num_generations": 4, # GRPO group size G (>= 2) "per_device_train_batch_size": 4, # must be divisible by num_generations "learning_rate": 1e-5, # LoRA tolerates a higher LR than full-FT GRPO (TRL default 1e-6) "beta": 0.04, # KL-to-reference coefficient "max_completion_length": 48, # rewrites are short "use_peft": True, # LoRA-by-default }, rewriter_gen_kwargs={"max_new_tokens": 64, "do_sample": True, "temperature": 0.9}, eval_gen_kwargs={"max_new_tokens": 16, "do_sample": False}, # used by the reward's task-LM passes ) pipeline_grpo = SteeringPipeline( model_name_or_path=MODEL_NAME, # frozen 8B task model: scores rewrites for the reward, then answers at inference controls=[prewrite_grpo], device_map="auto", hf_model_kwargs={"dtype": torch.bfloat16}, ) pipeline_grpo.steer() # GRPO trains the rewriter against the task scorer, then proposes a single rewrite print("Chosen rewrite (GRPO-trained rewriter, PRewrite-I):\n") print(prewrite_grpo.memory["instruction"]) del pipeline_grpo gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 128009, 'pad_token_id': 128009}.

[16/16 00:31, Epoch 2/2]

Step	Training Loss
10	0.018035

Chosen rewrite (GRPO-trained rewriter, PRewrite-I):

Please respond with only the final answer.

Quantifying instruction quality¶

PRewrite treats a rewrite as input-agnostic (one fixed instruction for all inputs) and judges it by the scorer value when the model uses that instruction (F1 or exact match for short-answer QA). In other words, the scorer is the reward the rewriter optimizes.

In [14]:

import pandas as pd
from IPython.display import display

pd.set_option("display.max_colwidth", None)  # show full instructions / generations

# held-out questions
test_set = [
    {"input": "What is the capital of Japan?", "reference": "Tokyo"},
    {"input": "Who developed the theory of general relativity?", "reference": "Einstein"},
    {"input": "What is the largest planet in the Solar System?", "reference": "Jupiter"},
    {"input": "In what year did the first crewed Moon landing occur?", "reference": "1969"},
    {"input": "Who wrote 'Romeo and Juliet'?", "reference": "Shakespeare"},
    {"input": "Which country gifted the Statue of Liberty to the United States?", "reference": "France"},
    {"input": "What is the tallest mountain on Earth?", "reference": "Everest"},
    {"input": "What gas do plants primarily absorb during photosynthesis?", "reference": "carbon dioxide"},
]

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, device_map="auto", dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model.generation_config.pad_token_id = tokenizer.eos_token_id  # silence the pad_token_id notice

import pandas as pd from IPython.display import display pd.set_option("display.max_colwidth", None) # show full instructions / generations # held-out questions test_set = [ {"input": "What is the capital of Japan?", "reference": "Tokyo"}, {"input": "Who developed the theory of general relativity?", "reference": "Einstein"}, {"input": "What is the largest planet in the Solar System?", "reference": "Jupiter"}, {"input": "In what year did the first crewed Moon landing occur?", "reference": "1969"}, {"input": "Who wrote 'Romeo and Juliet'?", "reference": "Shakespeare"}, {"input": "Which country gifted the Statue of Liberty to the United States?", "reference": "France"}, {"input": "What is the tallest mountain on Earth?", "reference": "Everest"}, {"input": "What gas do plants primarily absorb during photosynthesis?", "reference": "carbon dioxide"}, ] model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, device_map="auto", dtype=torch.bfloat16 ) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model.generation_config.pad_token_id = tokenizer.eos_token_id # silence the pad_token_id notice

The four instructions we compare are the seed and the three rewrites chosen above. For each one we have the task LM answer every held-out question.

In [15]:

candidates = {
    "Seed":       SEED_INSTRUCTION,
    "PRewrite-I": prewrite_i.memory["instruction"],
    "PRewrite-S": prewrite_s.memory["instruction"],
    "GRPO":       prewrite_grpo.memory["instruction"],
}
references = [ex["reference"] for ex in test_set]

answers_by_method = {name: [] for name in candidates}
with torch.no_grad():
    for name, instruction in candidates.items():
        for ex in test_set:
            chat = tokenizer.apply_chat_template(
                [
                    {"role": "system", "content": instruction},
                    {"role": "user", "content": ex["input"]},
                ],
                tokenize=False,
                add_generation_prompt=True,
            )
            enc = tokenizer(chat, return_tensors="pt").to(model.device)
            out = model.generate(**enc, max_new_tokens=64, do_sample=False)
            resp = tokenizer.decode(out[0, enc.input_ids.size(1):], skip_special_tokens=True)
            answers_by_method[name].append(resp.strip().replace("\n", " "))

candidates = { "Seed": SEED_INSTRUCTION, "PRewrite-I": prewrite_i.memory["instruction"], "PRewrite-S": prewrite_s.memory["instruction"], "GRPO": prewrite_grpo.memory["instruction"], } references = [ex["reference"] for ex in test_set] answers_by_method = {name: [] for name in candidates} with torch.no_grad(): for name, instruction in candidates.items(): for ex in test_set: chat = tokenizer.apply_chat_template( [ {"role": "system", "content": instruction}, {"role": "user", "content": ex["input"]}, ], tokenize=False, add_generation_prompt=True, ) enc = tokenizer(chat, return_tensors="pt").to(model.device) out = model.generate(**enc, max_new_tokens=64, do_sample=False) resp = tokenizer.decode(out[0, enc.input_ids.size(1):], skip_special_tokens=True) answers_by_method[name].append(resp.strip().replace("\n", " "))

We score each method with the f1() and exact() scorers over the held-out set and report one row per method, alongside the instruction it used.

In [16]:

def mean_score(scorer, answers, rows):
    return sum(scorer(answer, row) for answer, row in zip(answers, rows)) / len(rows)

summary = pd.DataFrame({
    "method": list(candidates),
    "f1": [mean_score(f1_scorer, answers_by_method[name], test_set) for name in candidates],
    "exact_match": [mean_score(em_scorer, answers_by_method[name], test_set) for name in candidates],
    "instruction": [candidates[name] for name in candidates],
})
display(summary)

def mean_score(scorer, answers, rows): return sum(scorer(answer, row) for answer, row in zip(answers, rows)) / len(rows) summary = pd.DataFrame({ "method": list(candidates), "f1": [mean_score(f1_scorer, answers_by_method[name], test_set) for name in candidates], "exact_match": [mean_score(em_scorer, answers_by_method[name], test_set) for name in candidates], "instruction": [candidates[name] for name in candidates], }) display(summary)

	method	f1	exact_match	instruction
0	Seed	0.21250	0.000	Please provide an answer to the question.
1	PRewrite-I	0.24000	0.000	Provide the answer.
2	PRewrite-S	0.87625	0.625	Provide the answer only.
3	GRPO	0.87625	0.625	Please respond with only the final answer.

Generations for each question and method are as follows.

In [17]:

detail = pd.DataFrame({"question": [ex["input"] for ex in test_set], "reference": references})
for name in candidates:
    detail[name] = answers_by_method[name]
display(detail)

del model
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

detail = pd.DataFrame({"question": [ex["input"] for ex in test_set], "reference": references}) for name in candidates: detail[name] = answers_by_method[name] display(detail) del model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache()

	question	reference	Seed	PRewrite-I	PRewrite-S	GRPO
0	What is the capital of Japan?	Tokyo	The capital of Japan is Tokyo.	The capital of Japan is Tokyo.	Tokyo.	Tokyo.
1	Who developed the theory of general relativity?	Einstein	Albert Einstein developed the theory of general relativity.	Albert Einstein developed the theory of general relativity.	Albert Einstein.	Albert Einstein.
2	What is the largest planet in the Solar System?	Jupiter	The largest planet in the Solar System is Jupiter. It is a gas giant, with a diameter of approximately 142,984 kilometers (88,846 miles). This is more than 11 times the diameter of the Earth. Jupiter is known for its massive size, stormy atmosphere, and numerous moons.	The largest planet in the Solar System is Jupiter.	Jupiter.	Jupiter.
3	In what year did the first crewed Moon landing occur?	1969	The first crewed Moon landing occurred in 1969.	The first crewed Moon landing occurred in 1969.	1969.	1969.
4	Who wrote 'Romeo and Juliet'?	Shakespeare	'Romeo and Juliet' was written by the famous English playwright William Shakespeare.	The play 'Romeo and Juliet' was written by William Shakespeare.	William Shakespeare.	William Shakespeare.
5	Which country gifted the Statue of Liberty to the United States?	France	The country that gifted the Statue of Liberty to the United States was France.	The country that gifted the Statue of Liberty to the United States was France.	France.	France.
6	What is the tallest mountain on Earth?	Everest	The tallest mountain on Earth is Mount Everest, which is part of the Himalayas in the Himalayan mountain range in Asia. It stands at a height of 8,848.86 meters (29,031.7 feet) above sea level.	The tallest mountain on Earth is Mount Everest, which is part of the Himalayas in the Himalayan mountain range in Asia. It stands at a height of 8,848.86 meters (29,031.7 feet) above sea level.	Mount Everest.	Mount Everest.
7	What gas do plants primarily absorb during photosynthesis?	carbon dioxide	Plants primarily absorb carbon dioxide (CO2) during photosynthesis.	Plants primarily absorb carbon dioxide (CO2) during photosynthesis.	Carbon dioxide.	Carbon dioxide.

Summary¶

This notebook demonstrated PRewrite (Kong et al., 2024), an input control that rewrites a seed instruction into a more effective form using an LLM as a rewriter, then applies the chosen rewrite as the system prompt at inference time. There are three variants of the method:

1. PRewrite-I (strategy="inference") generates a single greedy rewrite. This strategy only costs one extra rewriter call before steering and it is appropriate for minor instruction edits.
2. PRewrite-S (strategy="search") samples k_candidates rewrites and keeps the one with the highest mean score on a held-out dev_set under a row_scorer. It costs k × (rewriter call) + k × (dev_set size × task call); this method is useful when a dev set that captures the target behavior (and a good scorer to score it) is available.
3. Setting train_rewriter=True runs GRPO (group-relative policy optimization, via the TRL wrapper) to train the rewriter on a pool of training_seeds before proposing rewrites; the trained rewriter is then consumed by either PRewrite-I or PRewrite-S. The GRPO reward is the same scorer-in-the-loop signal PRewrite-S uses to select rewrites (apply each rewrite with the frozen task model over the dev_set and score it), now used as the training reward, or a user-supplied reward_fn. This optimizes the rewriter directly against task performance, at the cost of more compute (each reward evaluation runs the task model over the dev set for every distinct rewrite in a GRPO group).