Phased Decoding¶

PhasedDecoding is a generic output control that builds a generation as a sequence of phases, splicing forced text and generated segments into one stream. The phases are declared in a plan list whose entries are either {"fixed": ...} or {"generate": {...}}. Budget forcing, response prefill, and thinking intervention can therefore all be specified as PhasedDecoding configs (rather than separate classes).

PhasedDecoding is a decoding driver (at most one enabled driver runs per pipeline). Since each generate phase calls the model with the logits processors and stopping criteria contributed by the rest of the pipeline, a step-level control like ValueGuidance steers every generated phase.

This notebook runs each config against one instruction model. Since the driver returns a single spliced stream without phase boundaries, the segmentation display reconstructs the phases from the plan's own forced strings. The thinking-intervention section runs a plan that rewrites the prompt and strips the reasoning span.

The plan grammar¶

Each entry is a dict with exactly one key:

• {"fixed": <str | callable>, "replace": bool, "add_special_tokens": bool} splices text, either a literal or a (prompt_text, params) -> str callable.
• {"generate": {"until": str | None, "budget": int | None}} generates until a boundary. {"generate": {}} is unbounded.

Plans whose fixed values are all strings are JSON-serializable.

Method parameters¶

parameter	type	description
plan	list	List of phase dicts (each with one of fixed / generate)
extract_after	str \| None	Keep the prompt prefix and the remainder after this marker. None keeps the full stream

Setup¶

If running this from a Google Colab notebook, uncomment the clone cell below. It is not necessary when running from a virtual environment where the package is already installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability

In [2]:

import sys
!{sys.executable} -m pip install -q tabulate

import sys !{sys.executable} -m pip install -q tabulate

In [3]:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.algorithms.output_control.phased_decoding.control import PhasedDecoding
from steerability.algorithms.output_control.stopping_rules.control import StoppingRules

from IPython.display import display, HTML
display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>"))

from tabulate import tabulate
import textwrap

def wrap(text, width=60):
    return '\n'.join(textwrap.wrap(text, width=width))

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.core.steering_pipeline import SteeringPipeline from steerability.algorithms.output_control.phased_decoding.control import PhasedDecoding from steerability.algorithms.output_control.stopping_rules.control import StoppingRules from IPython.display import display, HTML display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>")) from tabulate import tabulate import textwrap def wrap(text, width=60): return '\n'.join(textwrap.wrap(text, width=width))

We use Qwen/Qwen2.5-1.5B-Instruct and load it once, building a fresh SteeringPipeline per configuration around the shared model. Since PhasedDecoding is a decoding driver, each pipeline runs the plan itself rather than composing a logits processor into a single decode pass.

In [4]:

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

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

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) device = model.device

Budget forcing¶

Budget forcing shapes a reasoning trace by bounding a thinking phase, forcing a "Wait" extension to make the model keep thinking, extending the thinking phase, forcing the closing </think> tag, then generating the answer. We run it at two thinking budgets to see the budget and the forced extension take effect.

The driver returns only the spliced token stream, with no phase-boundary metadata. The segmentation display therefore reconstructs the phases from the plan's own forced strings. The helper below splits the decoded stream on those markers and tabulates each phase as generated or forced.

In [5]:

def budget_forcing_plan(thinking_budget, extension_budget):
    return [
        {"generate": {"until": "</think>", "budget": thinking_budget}},
        {"fixed": "Wait"},
        {"generate": {"until": "</think>", "budget": extension_budget}},
        {"fixed": "</think>"},
        {"generate": {}},
    ]

def segment_by_forced(text, forced_strings):
    rows, cursor = [], 0
    for marker in forced_strings:
        idx = text.find(marker, cursor)
        if idx == -1:
            break
        if idx > cursor:
            rows.append(("generated", text[cursor:idx]))
        rows.append(("forced", marker))
        cursor = idx + len(marker)
    if cursor < len(text):
        rows.append(("generated", text[cursor:]))
    return rows

bf_prompt = "<think>\nLet me solve 12 * 7 step by step."
bf_inputs = tokenizer(bf_prompt, return_tensors="pt").to(device)

for budget in (16, 64):
    plan = budget_forcing_plan(budget, 32)
    pipeline = SteeringPipeline(controls=[PhasedDecoding(plan=plan)], model=model, tokenizer=tokenizer)
    pipeline.steer()
    out = pipeline.generate(input_ids=bf_inputs["input_ids"], max_new_tokens=256, do_sample=False,
                            pad_token_id=tokenizer.eos_token_id, return_full_sequence=True)
    stream = tokenizer.decode(out[0], skip_special_tokens=True)
    rows = [[kind, wrap(chunk.strip(), 84)] for kind, chunk in segment_by_forced(stream, ["Wait", "</think>"]) if chunk.strip()]
    print(f"thinking budget = {budget}")
    print(tabulate(rows, headers=["phase", "text"], tablefmt="grid", maxcolwidths=[10, 84]))
    print()

def budget_forcing_plan(thinking_budget, extension_budget): return [ {"generate": {"until": "</think>", "budget": thinking_budget}}, {"fixed": "Wait"}, {"generate": {"until": "</think>", "budget": extension_budget}}, {"fixed": "</think>"}, {"generate": {}}, ] def segment_by_forced(text, forced_strings): rows, cursor = [], 0 for marker in forced_strings: idx = text.find(marker, cursor) if idx == -1: break if idx > cursor: rows.append(("generated", text[cursor:idx])) rows.append(("forced", marker)) cursor = idx + len(marker) if cursor < len(text): rows.append(("generated", text[cursor:])) return rows bf_prompt = "<think>\nLet me solve 12 * 7 step by step." bf_inputs = tokenizer(bf_prompt, return_tensors="pt").to(device) for budget in (16, 64): plan = budget_forcing_plan(budget, 32) pipeline = SteeringPipeline(controls=[PhasedDecoding(plan=plan)], model=model, tokenizer=tokenizer) pipeline.steer() out = pipeline.generate(input_ids=bf_inputs["input_ids"], max_new_tokens=256, do_sample=False, pad_token_id=tokenizer.eos_token_id, return_full_sequence=True) stream = tokenizer.decode(out[0], skip_special_tokens=True) rows = [[kind, wrap(chunk.strip(), 84)] for kind, chunk in segment_by_forced(stream, ["Wait", "</think>"]) if chunk.strip()] print(f"thinking budget = {budget}") print(tabulate(rows, headers=["phase", "text"], tablefmt="grid", maxcolwidths=[10, 84])) print()

thinking budget = 16
+-----------+--------------------------------------------------------------------------------------+
| phase     | text                                                                                 |
+===========+======================================================================================+
| generated | <think> Let me solve 12 * 7 step by step. First, I'll multiply the ones place of     |
|           | each number: 2 *                                                                     |
+-----------+--------------------------------------------------------------------------------------+
| forced    | Wait                                                                                 |
+-----------+--------------------------------------------------------------------------------------+
| generated | for input.                                                                           |
+-----------+--------------------------------------------------------------------------------------+
| forced    | </think>                                                                             |
+-----------+--------------------------------------------------------------------------------------+
| generated | </think>   </think>  Sure! Let's break down the multiplication of 12 and 7 step by   |
|           | step.  ### Step-by-Step Multiplication:  1. **Multiply the ones place:**    - The    |
|           | ones digit of 12 is 2.    - Multiply 2 by 7:      \[      2 \times 7 = 14      \]    |
|           | - Write down 4 in the ones place of our answer (since we are only considering the    |
|           | ones place so far).  2. **Multiply the tens place:**    - The tens digit of 12 is 1. |
|           | - Since there is no tens place in 7, it means \(70\) (which is \(7 \times 10\)) will |
|           | be multiplied by 1.    - So, \(1 \times 70 = 70\).    - Add this to what we have so  |
|           | far from the previous step (the 4):      \[      4 + 70 = 74      \]  So, the final  |
|           | result of multiplying 12 by 7 is: \[ 12 \times 7 = 84 \]                             |
+-----------+--------------------------------------------------------------------------------------+

thinking budget = 64
+-----------+--------------------------------------------------------------------------------------+
| phase     | text                                                                                 |
+===========+======================================================================================+
| generated | <think> Let me solve 12 * 7 step by step. First, I'll multiply the ones place of     |
|           | each number: 2 * 7 = 14. Then, I'll carry over the 1 to the tens place. Next, I'll   |
|           | add the tens place of both numbers: 1 * 7 + 0 = 7. Finally, I'll                     |
+-----------+--------------------------------------------------------------------------------------+
| forced    | Wait                                                                                 |
+-----------+--------------------------------------------------------------------------------------+
| generated | for your input to continue.                                                          |
+-----------+--------------------------------------------------------------------------------------+
| forced    | </think>                                                                             |
+-----------+--------------------------------------------------------------------------------------+
| generated | </think> Let's break down the multiplication of 12 and 7 step by step:  ### Step 1:  |
|           | Multiply the Ones Place - The ones place of 12 is 2. - The ones place of 7 is 7. -   |
|           | \( 2 \times 7 = 14 \).  Since this product (14) is a two-digit number, we need to    |
|           | write it as 1 with a carry-over of 1 to the next column.  ### Step 2: Carry Over the |
|           | 1 - We have carried over 1 from the previous multiplication. - Now, we move on to    |
|           | the tens place.  ### Step 3: Add the Tens Place - The tens place of 12 is 1. - The   |
|           | tens place of 7 is 0. - Adding these together gives us \( 1 + 0 = 1 \).  So, after   |
|           | adding the carry-over from the first step, we get: \[ 1 + 1 = 2 \]  ### Final Result |
|           | Combining all parts, we have: \[ 12 \times 7 = 84 \]  Therefore, \( 12 \times 7 = 84 |
|           | \).   If you have any other                                                          |
+-----------+--------------------------------------------------------------------------------------+

Extracting the answer¶

extract_after keeps the prompt prefix and the remainder after a marker, dropping the reasoning trace. Adding extract_after="</think>" to the same budget-forcing plan returns only the answer that follows the closing tag. The thinking shapes the answer but is not shown.

In [6]:

extract_plan = budget_forcing_plan(16, 32)
extract_pipeline = SteeringPipeline(
    controls=[PhasedDecoding(plan=extract_plan, extract_after="</think>")], model=model, tokenizer=tokenizer,
)
extract_pipeline.steer()

out = extract_pipeline.generate(input_ids=bf_inputs["input_ids"], max_new_tokens=256, do_sample=False,
                                pad_token_id=tokenizer.eos_token_id, return_full_sequence=True)
answer_only = tokenizer.decode(out[0], skip_special_tokens=True)
print("answer only (reasoning trace dropped):")
print(wrap(answer_only, 100))

extract_plan = budget_forcing_plan(16, 32) extract_pipeline = SteeringPipeline( controls=[PhasedDecoding(plan=extract_plan, extract_after="</think>")], model=model, tokenizer=tokenizer, ) extract_pipeline.steer() out = extract_pipeline.generate(input_ids=bf_inputs["input_ids"], max_new_tokens=256, do_sample=False, pad_token_id=tokenizer.eos_token_id, return_full_sequence=True) answer_only = tokenizer.decode(out[0], skip_special_tokens=True) print("answer only (reasoning trace dropped):") print(wrap(answer_only, 100))

answer only (reasoning trace dropped):
<think> Let me solve 12 * 7 step by step.Sure! Let's break down the multiplication of 12 and 7 step
by step.  ### Step-by-Step Multiplication:  1. **Multiply the ones place:**    - The ones digit of
12 is 2.    - Multiply 2 by 7:      \[      2 \times 7 = 14      \]    - Write down 4 in the ones
place of our answer (since we are only considering the ones place so far).  2. **Multiply the tens
place:**    - The tens digit of 12 is 1.    - Since there is no tens place in 7, it means \(70\)
(which is \(7 \times 10\)) will be multiplied by 1.    - So, \(1 \times 70 = 70\).    - Add this to
what we have so far from the previous step (the 4):      \[      4 + 70 = 74      \]  So, the final
result of multiplying 12 by 7 is: \[ 12 \times 7 = 84 \]

Response prefill¶

A two-phase plan can force the answer to begin with a fixed string, then generate from there. This is response prefill, where the forced opening commits the model to a framing before it generates. The contrast below runs the same prompt unprefilled and prefilled with a fixed opener, making the effect of the committed opening on the rest of the answer visible.

In [7]:

prefill_prompt = "Should I learn to play the piano as an adult?"
prefill_chat = tokenizer.apply_chat_template(
    [{"role": "user", "content": prefill_prompt}], tokenize=False, add_generation_prompt=True
)
prefill_inputs = tokenizer(prefill_chat, return_tensors="pt").to(device)
prefill_gen = {"max_new_tokens": 50, "do_sample": False, "pad_token_id": tokenizer.eos_token_id, "return_full_sequence": True}

plain_plan = [{"generate": {}}]
prefilled_plan = [{"fixed": "Absolutely, and here is exactly how to start:\n"}, {"generate": {}}]

table = []
for label, plan in [("no prefill", plain_plan), ("prefilled", prefilled_plan)]:
    pipeline = SteeringPipeline(controls=[PhasedDecoding(plan=plan)], model=model, tokenizer=tokenizer)
    pipeline.steer()
    out = pipeline.generate(input_ids=prefill_inputs["input_ids"], **prefill_gen)
    completion = tokenizer.decode(out[0][prefill_inputs["input_ids"].size(1):], skip_special_tokens=True)
    table.append([label, wrap(completion, 74)])

print(f"Prompt: {prefill_prompt}")
print(tabulate(table, headers=["config", "answer"], tablefmt="grid", maxcolwidths=[12, 74]))

prefill_prompt = "Should I learn to play the piano as an adult?" prefill_chat = tokenizer.apply_chat_template( [{"role": "user", "content": prefill_prompt}], tokenize=False, add_generation_prompt=True ) prefill_inputs = tokenizer(prefill_chat, return_tensors="pt").to(device) prefill_gen = {"max_new_tokens": 50, "do_sample": False, "pad_token_id": tokenizer.eos_token_id, "return_full_sequence": True} plain_plan = [{"generate": {}}] prefilled_plan = [{"fixed": "Absolutely, and here is exactly how to start:\n"}, {"generate": {}}] table = [] for label, plan in [("no prefill", plain_plan), ("prefilled", prefilled_plan)]: pipeline = SteeringPipeline(controls=[PhasedDecoding(plan=plan)], model=model, tokenizer=tokenizer) pipeline.steer() out = pipeline.generate(input_ids=prefill_inputs["input_ids"], **prefill_gen) completion = tokenizer.decode(out[0][prefill_inputs["input_ids"].size(1):], skip_special_tokens=True) table.append([label, wrap(completion, 74)]) print(f"Prompt: {prefill_prompt}") print(tabulate(table, headers=["config", "answer"], tablefmt="grid", maxcolwidths=[12, 74]))

Prompt: Should I learn to play the piano as an adult?
+------------+----------------------------------------------------------------------------+
| config     | answer                                                                     |
+============+============================================================================+
| no prefill | Yes, learning to play the piano as an adult can be a rewarding experience! |
|            | Playing an instrument like the piano can improve your cognitive skills     |
|            | such as memory and concentration, enhance your creativity, and provide a   |
|            | sense of accomplishment.  If you're interested in learning                 |
+------------+----------------------------------------------------------------------------+
| prefilled  | Absolutely, and here is exactly how to start: 1. Choose a good teacher:    |
|            | Look for a qualified music instructor who can teach you proper technique   |
|            | and help you develop your skills. 2. Invest in quality equipment: A decent |
|            | piano or keyboard will be essential for practicing regularly. 3. Set       |
|            | realistic goals                                                            |
+------------+----------------------------------------------------------------------------+

Thinking intervention¶

Thinking intervention (Wu et al., 2025, arXiv:2503.24370) rewrites the prompt to splice guidance into the model's reasoning stream. As a plan it is a single replacing fixed phase (the intervention-rewritten prompt) followed by a generate phase, with extract_after="</think>" stripping the reasoning span and returning only the answer. The intervention itself is a (prompt_text, params) -> str callable. Here it prepends a short guidance sentence and a </think> marker.

In [8]:

def intervention(prompt, params):
    return f"Reason carefully and show each step. </think> {prompt}"

ti_prompt = tokenizer("What is 6 times 7?", return_tensors="pt").input_ids.to(device)

pd = PhasedDecoding(
    plan=[{"fixed": intervention, "replace": True, "add_special_tokens": True}, {"generate": {}}],
    extract_after="</think>",
)
pd_pipeline = SteeringPipeline(controls=[pd], model=model, tokenizer=tokenizer)
pd_pipeline.steer()
torch.manual_seed(0)
out = pd_pipeline.generate(input_ids=ti_prompt, max_new_tokens=8, do_sample=False, eos_token_id=None)
print(tokenizer.decode(out[0], skip_special_tokens=True))

def intervention(prompt, params): return f"Reason carefully and show each step. </think> {prompt}" ti_prompt = tokenizer("What is 6 times 7?", return_tensors="pt").input_ids.to(device) pd = PhasedDecoding( plan=[{"fixed": intervention, "replace": True, "add_special_tokens": True}, {"generate": {}}], extract_after="</think>", ) pd_pipeline = SteeringPipeline(controls=[pd], model=model, tokenizer=tokenizer) pd_pipeline.steer() torch.manual_seed(0) out = pd_pipeline.generate(input_ids=ti_prompt, max_new_tokens=8, do_sample=False, eos_token_id=None) print(tokenizer.decode(out[0], skip_special_tokens=True))

What is 6 times 7? To calculate \( 6 \times 

Phases and stops¶

A StoppingRules control composes into every generated phase of the plan. The criteria are anchored to the original prompt at composition time. The stop is therefore global, firing inside a generated phase relative to the whole stream rather than relative to the phase. Below, a two-phase plan runs with a substring stop, and the stop halts generation as soon as the marker appears in the stream. This is the same global behavior described in the semantics section of the stopping-rules notebook.

In [9]:

stops_prompt = "List a few uses for a paperclip, then add a blank line and a closing remark."
stops_chat = tokenizer.apply_chat_template(
    [{"role": "user", "content": stops_prompt}], tokenize=False, add_generation_prompt=True
)
stops_inputs = tokenizer(stops_chat, return_tensors="pt").to(device)
stops_gen = {"max_new_tokens": 120, "do_sample": False, "pad_token_id": tokenizer.eos_token_id, "return_full_sequence": True}

two_phase_plan = [{"fixed": "Here are some uses:\n"}, {"generate": {}}]

table = []
for label, controls in [
    ("plan only", [PhasedDecoding(plan=two_phase_plan)]),
    ("plan + stop at \\n\\n", [PhasedDecoding(plan=two_phase_plan), StoppingRules(stop_texts=["\n\n"])]),
]:
    pipeline = SteeringPipeline(controls=controls, model=model, tokenizer=tokenizer)
    pipeline.steer()
    out = pipeline.generate(input_ids=stops_inputs["input_ids"], **stops_gen)
    completion = tokenizer.decode(out[0][stops_inputs["input_ids"].size(1):], skip_special_tokens=True)
    table.append([label, out[0].size(0) - stops_inputs["input_ids"].size(1), wrap(completion, 66)])

print(f"Prompt: {stops_prompt}")
print(tabulate(table, headers=["config", "new tokens", "generated"], tablefmt="grid", maxcolwidths=[20, 10, 66]))

stops_prompt = "List a few uses for a paperclip, then add a blank line and a closing remark." stops_chat = tokenizer.apply_chat_template( [{"role": "user", "content": stops_prompt}], tokenize=False, add_generation_prompt=True ) stops_inputs = tokenizer(stops_chat, return_tensors="pt").to(device) stops_gen = {"max_new_tokens": 120, "do_sample": False, "pad_token_id": tokenizer.eos_token_id, "return_full_sequence": True} two_phase_plan = [{"fixed": "Here are some uses:\n"}, {"generate": {}}] table = [] for label, controls in [ ("plan only", [PhasedDecoding(plan=two_phase_plan)]), ("plan + stop at \\n\\n", [PhasedDecoding(plan=two_phase_plan), StoppingRules(stop_texts=["\n\n"])]), ]: pipeline = SteeringPipeline(controls=controls, model=model, tokenizer=tokenizer) pipeline.steer() out = pipeline.generate(input_ids=stops_inputs["input_ids"], **stops_gen) completion = tokenizer.decode(out[0][stops_inputs["input_ids"].size(1):], skip_special_tokens=True) table.append([label, out[0].size(0) - stops_inputs["input_ids"].size(1), wrap(completion, 66)]) print(f"Prompt: {stops_prompt}") print(tabulate(table, headers=["config", "new tokens", "generated"], tablefmt="grid", maxcolwidths=[20, 10, 66]))

Prompt: List a few uses for a paperclip, then add a blank line and a closing remark.
+---------------------+--------------+------------------------------------------------------------------+
| config              |   new tokens | generated                                                        |
+=====================+==============+==================================================================+
| plan only           |           39 | Here are some uses: - Holding papers together in a stack -       |
|                     |              | Clipping documents to bind them - Keeping loose change organized |
|                     |              | Closing remark: A simple tool with many practical applications!  |
+---------------------+--------------+------------------------------------------------------------------+
| plan + stop at \n\n |           27 | Here are some uses: - Holding papers together in a stack -       |
|                     |              | Clipping documents to bind them - Keeping loose change organized |
+---------------------+--------------+------------------------------------------------------------------+

Summary¶

Every config here was an assignment of a PhasedDecoding plan over one instruction model. Budget forcing shaped a reasoning trace by bounding a thinking phase, forcing a "Wait" extension and a closing tag, and generating the answer, with a segmentation display reconstructed from the plan's forced strings. Adding extract_after returned the answer alone. Response prefill committed the answer to a forced opening. A thinking-intervention plan rewrote the prompt through a replacing fixed phase and stripped the reasoning span with extract_after. A StoppingRules control composed into a generated phase, firing globally relative to the whole stream.