CPO¶
CPO (Chen et al., 2026) is an input control method that treats system-prompt selection as a causal problem, i.e., which prompt, given the query, raises the model's evaluation score the most? It works in two stages.
- Stage 1 (performed in
steer): build an offline⟨query, prompt, score⟩dataset, embed queries and prompts with a small text encoder, reduce to PCA features, and fit aCausalRewardScorer. Ifeconmlis available the scorer is aCausalForestDMLover the prompt-as-treatment; otherwise it falls back to aGradientBoostingRegressor(still useful, but no longer a causal estimate). - Stage 2 (
adapt_messages): for each query at inference time, run a tree search via proposal ofcandidates_per_parentrewrites of each survivor, scoring them with the trained reward model, keeping the topretained_per_round, repeating forroundsrounds, and returning the highest-scoring survivor as the system prompt. Per-query results are cached.
Reference: Chen et al., 2026 — Causal Prompt Optimization.
Method parameters¶
| parameter | type | description |
|---|---|---|
seed_prompt |
str |
Base prompt used as the DML treatment baseline and tree-search root. Required. |
offline_data |
list[dict] \| None |
Pre-built ⟨query, prompt, score⟩ rows. Skips offline-data generation entirely. Cheapest path. |
train_dataset |
list[dict] |
Task instances (with "input" and optionally "reference"). Used to generate offline data when offline_data is None. |
row_scorer |
Callable \| None |
Per-row scorer (response, row) -> float used to score offline-generated rows. Required when offline_data is None. |
prompt_lm |
model | None |
LLM used to propose candidate prompts. None reuses the task model. |
n_prompts_per_query |
int |
When generating offline data, how many candidate prompts per training query (the seed counts as one). |
embedding_model |
str |
HF text encoder used to featurize queries and prompts. Default nomic-ai/nomic-embed-text-v1.5. |
pca_query_dim |
int |
PCA target dim for query embeddings (clipped to available rank). |
pca_prompt_dim |
int |
PCA target dim for prompt embeddings (clipped to available rank). |
rounds |
int |
Tree-search rounds (R). Default 3. |
candidates_per_parent |
int |
Candidates spawned per retained parent per round (B). Default 5. |
retained_per_round |
int |
Survivors kept after each round (K). Default 3. |
cache_queries |
bool |
Cache the chosen prompt per-query so identical queries skip the search. Default True. |
use_dml |
bool \| None |
True requires econml, installed separately (pip install "econml>=0.16,<0.17"); False forces the GBR fallback; None auto-detects. |
refinement_meta_prompt |
str \| None |
Override the CPO refinement template; None uses refinement_meta_prompt.CPO_DEFAULT. |
proposer_gen_kwargs |
dict \| None |
Generation kwargs for the prompt proposer LLM. |
eval_gen_kwargs |
dict \| None |
Generation kwargs used during offline data scoring. |
The query cache uses an exact sha256 of the query text; whitespace and casing differences are distinct entries. Pre-normalize at the call site if you want near-duplicate queries to share a cached prompt.
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 two optional dependencies: cpo (econml) for the DML estimator, without which CPO transparently falls back to a gradient-boosted regressor, and inspect for the f1() scorer used to score offline rows (both dev and all include inspect).
# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability
# !pip install -q -e ".[cpo,inspect]"
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.
# !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)
import gc
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from inspect_ai.scorer import f1
from steerability.algorithms.input_control.cpo import CPO
from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.evaluation.scorers import sample_scorer_from_inspect
MODEL_NAME = "google/gemma-3-4b-it"
EMBEDDING_MODEL = "nomic-ai/nomic-embed-text-v1.5"
PROMPT = "Who wrote the novel '1984'?" # held-out test query
f1_scorer = sample_scorer_from_inspect(f1())
Prebuilt offline_data¶
CPO can be run using a small list of ⟨query, prompt, score⟩ triples. This doesn't require any offline generation or metric calls, just embed, fit, and search. This is useful when you have an existing log of prompt evaluations or a manually curated set of (prompt, score) observations.
We construct ~20 rows by hand consisting of three short questions, each scored under several candidate system prompts. Under use_dml=False we get the GBR fallback, which fits comfortably on this little data.
offline_rows = [
{"query": "What's the capital of France?", "prompt": "Answer the question.", "score": 0.30},
{"query": "What's the capital of France?", "prompt": "Reply with just the city name.", "score": 0.95},
{"query": "What's the capital of France?", "prompt": "Be concise.", "score": 0.70},
{"query": "What's the capital of France?", "prompt": "Provide one short factual answer.", "score": 0.85},
{"query": "What's the capital of France?", "prompt": "Explain in detail with context.", "score": 0.20},
{"query": "How many planets are in the Solar System?", "prompt": "Answer the question.", "score": 0.40},
{"query": "How many planets are in the Solar System?", "prompt": "Reply with just the number.", "score": 0.95},
{"query": "How many planets are in the Solar System?", "prompt": "Be concise.", "score": 0.75},
{"query": "How many planets are in the Solar System?", "prompt": "Provide one short factual answer.", "score": 0.85},
{"query": "How many planets are in the Solar System?", "prompt": "Explain in detail with context.", "score": 0.25},
{"query": "What's the chemical symbol for sodium?", "prompt": "Answer the question.", "score": 0.35},
{"query": "What's the chemical symbol for sodium?", "prompt": "Reply with just the symbol.", "score": 0.95},
{"query": "What's the chemical symbol for sodium?", "prompt": "Be concise.", "score": 0.70},
{"query": "What's the chemical symbol for sodium?", "prompt": "Provide one short factual answer.", "score": 0.85},
{"query": "What's the chemical symbol for sodium?", "prompt": "Explain in detail with context.", "score": 0.20},
{"query": "Who painted the Mona Lisa?", "prompt": "Answer the question.", "score": 0.40},
{"query": "Who painted the Mona Lisa?", "prompt": "Reply with just the name.", "score": 0.95},
{"query": "Who painted the Mona Lisa?", "prompt": "Be concise.", "score": 0.75},
{"query": "Who painted the Mona Lisa?", "prompt": "Provide one short factual answer.", "score": 0.85},
{"query": "Who painted the Mona Lisa?", "prompt": "Explain in detail with context.", "score": 0.25},
]
cpo_offline = CPO(
seed_prompt="Answer the question.",
offline_data=offline_rows,
embedding_model=EMBEDDING_MODEL,
pca_query_dim=4,
pca_prompt_dim=4,
rounds=2,
candidates_per_parent=3,
retained_per_round=2,
use_dml=False, # force GBR for the demo (deterministic on small data)
proposer_gen_kwargs={"max_new_tokens": 64, "do_sample": True, "temperature": 0.9},
)
pipeline_offline = SteeringPipeline(
model_name_or_path=MODEL_NAME,
controls=[cpo_offline],
device_map="auto",
)
pipeline_offline.steer()
print("Reward model mode:", cpo_offline.memory.causal_scorer.mode)
<All keys matched successfully>
[transformers] Detected the usage of `get_extended_attention_mask`: This function is deprecated and will be removed in v5.12.0. Please use the new API in `transformers.masking_utils`
Reward model mode: gbr
Now run a tree search at inference time on a held-out query. The control caches the chosen prompt per-query, so a second call with the same query is instantaneous.
Here we call adapt_messages directly only to inspect the prompt CPO chooses for this query. Real generation does not require this step: you pass the chat messages straight to pipeline.generate(...), which is the canonical path. The pipeline invokes adapt_messages before chat-template tokenization and applies the control exactly once. Passing pre-templated input_ids instead would force the lossy token-level adapt fallback, which decodes the templated ids back to text (leaking template boilerplate into the recovered query), re-runs the search on that polluted query, and pollutes the cache with a second entry.
adapted = cpo_offline.adapt_messages([[{"role": "user", "content": PROMPT}]])
chosen_prompt = next(
(msg["content"] for msg in adapted[0] if msg.get("role") == "system"),
None,
)
print("Chosen system prompt for the held-out query:\n")
print(chosen_prompt)
print("\nQuery cache size after one adapt call:", len(cpo_offline.memory.query_cache))
Chosen system prompt for the held-out query: Provide a concise and accurate response to the following question, demonstrating a clear understanding of the subject matter and prioritizing factual information. Query cache size after one adapt call: 1
# generate with the steered pipeline: pass chat messages directly.
# `generate` runs adapt_messages (cache hit from the call above), tokenizes with the chat
# template, and returns ONLY the newly generated text (the prompt is already stripped).
response_offline = pipeline_offline.generate(
messages=[{"role": "user", "content": PROMPT}],
max_new_tokens=500,
do_sample=False,
)
print("Response (CPO with pre-built offline_data):\n")
print(response_offline)
del pipeline_offline
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
Response (CPO with pre-built offline_data): George Orwell wrote the novel ‘1984’.
train_dataset-driven¶
Here we let CPO generate the offline data. First for each training row the proposer drafts n_prompts_per_query candidate prompts. Each candidate is run on that row's query under the scorer and the resulting ⟨query, prompt, score⟩ rows train the scorer.
We score 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. The choice of scorer matters for the reward model. Exact match is a 0/1 step function that saturates toward 0.0 here, since Gemma tends to answer "The capital of France is Paris" rather than "Paris", so a scorer fit on (near-)constant targets predicts that constant everywhere and the tree search degenerates to arbitrary tie-breaking. F1's precision term instead penalizes answers that contain the gold span, so a concise correct answer scores near 1.0, a correct-but-verbose one lands in the middle, and a wrong one scores 0.0. Prompts therefore result in a spread of scores and the scorer has gradation for fitting (causal_reward.train warns when the offline scores are constant or at least 95 percent saturated).
Note that this example performs on the order of len(train_dataset) × n_prompts_per_query task-LM generations, so we keep both numbers small (5 rows × 10 prompts = 50 generations).
Note that on an A100 this demo should run in 3-5 minutes. On a CPU it is impractically slow (the encoder is small, but the task LM and the proposer LM are not) so we strongly suggest running on a GPU.
train_dataset = [
{"input": "What's the capital of France?", "reference": "Paris"},
{"input": "How many planets are in the Solar System?", "reference": "8"},
{"input": "What's the chemical symbol for sodium?", "reference": "Na"},
{"input": "Who painted the Mona Lisa?", "reference": "Leonardo"},
{"input": "What year did World War II end?", "reference": "1945"},
]
# instantiate without offline_data; CPO will generate ⟨query, prompt, score⟩ rows itself
task_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
task_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if task_tokenizer.pad_token_id is None:
task_tokenizer.pad_token = task_tokenizer.eos_token
cpo_train = CPO(
seed_prompt="Answer the question.",
train_dataset=train_dataset,
row_scorer=f1_scorer,
prompt_lm=task_model, # reuse the task model as the proposer LM
n_prompts_per_query=10,
embedding_model=EMBEDDING_MODEL,
pca_query_dim=4,
pca_prompt_dim=4,
rounds=2,
candidates_per_parent=3,
retained_per_round=2,
use_dml=False,
proposer_gen_kwargs={"max_new_tokens": 64, "do_sample": True, "temperature": 0.9},
eval_gen_kwargs={"max_new_tokens": 16, "do_sample": False},
)
pipeline_train = SteeringPipeline(controls=[cpo_train], model=task_model, tokenizer=task_tokenizer)
pipeline_train.steer()
print("Reward model mode:", cpo_train.memory.causal_scorer.mode)
<All keys matched successfully>
Reward model mode: gbr
adapted_train = cpo_train.adapt_messages([[{"role": "user", "content": PROMPT}]])
chosen_train = next(
(msg["content"] for msg in adapted_train[0] if msg.get("role") == "system"),
None,
)
print("Chosen system prompt (train_dataset-driven CPO):\n")
print(chosen_train)
Chosen system prompt (train_dataset-driven CPO): Answer the question.
We now generate with the steered pipeline by passing in chat messages directly; generate runs adapt_messages (cache hit from the call above), tokenizes with the chat template, and returns only the newly generated text (the prompt is already stripped).
response_train = pipeline_train.generate(
messages=[{"role": "user", "content": PROMPT}],
max_new_tokens=64,
do_sample=False,
)
print("Response (CPO with train_dataset-driven offline data):\n")
print(response_train)
Response (CPO with train_dataset-driven offline data): George Orwell wrote the novel ‘1984’.
Inspecting the trained scorer and cache¶
A few useful things to inspect:
cpo.memory.causal_scorer.mode:"dml"(witheconml) or"gbr"(the fallback).cpo.memory.query_cache: the per-query cache; one entry per unique query CPO has been asked (useful for confirming repeat queries hit the cache).cpo.memory.causal_scorer.score(prompts, queries=...): score arbitrary(prompt, query)pairs directly without going through the tree search.
Canary: if the direct-score probe below prints (near-)identical values for clearly different prompts, the offline data had no score variance — the reward model learned nothing and the tree search is selecting by arbitrary tie-break. causal_reward.train emits a warning naming this failure when the offline scores are constant; fix it with a metric/dev set that produces real variance (harder queries, stricter matching, or graded scoring).
print("Scorer mode:", cpo_train.memory.causal_scorer.mode)
print("Query cache size:", len(cpo_train.memory.query_cache))
# score a few candidate prompts directly against the held-out query
candidates = [
"Answer the question.",
"Reply with just the answer.",
"Be concise.",
"Explain in detail with context.",
]
direct_scores = cpo_train.memory.causal_scorer.score(
candidates,
queries=[{"text": PROMPT}] * len(candidates),
)
for prompt, score in zip(candidates, direct_scores):
print(f" {score:+.3f} {prompt}")
Scorer mode: gbr Query cache size: 1
+0.181 Answer the question. +0.145 Reply with just the answer. +0.259 Be concise. +0.172 Explain in detail with context.
Summary¶
This notebook demonstrated CPO (causal prompt optimization), an input control that selects a per-query system prompt by asking which prompt, given the query, most raises the model's evaluation score:
- At
steertime CPO builds an offline⟨query, prompt, score⟩dataset and fits aCausalRewardScorerover PCA-reduced embeddings (aCausalForestDMLwitheconml, else a gradient-boosted regressor). - The prebuilt
offline_datapath is essentially free atsteertime and suits curated rows or logged evaluations; thetrain_dataset-driven path is more faithful (it samples prompts from the same proposer used at inference) but costslen(train_dataset) × n_prompts_per_querytask-LM generations. - At inference time
adapt_messagesruns a B-wide, K-retained, R-round tree search scored by the reward model, returning and caching the best system prompt per unique query.
CPO is only as good as the variance in its offline scores. A scorer that saturates leaves the reward model fitting a near-constant target and the search degenerates to arbitrary tie-breaking, which is why we score with Inspect AI's non-saturating f1() scorer, and why causal_reward.train and the direct-score probe act as early-warning indicators. The search uses elitism (parents compete with their children each round), so CPO never returns a prompt scored below the seed. For generation, pass chat messages straight to pipeline.generate so adapt_messages runs once before tokenization, rather than pre-templated input_ids, which force the lossy token-level adapt fallback.